[mlir][spirv] Add correct handling of Kernel and Addresses capabilities

This change adds initial support needed to generate OpenCL compliant SPIRV.
If Kernel capability is declared then memory model becomes OpenCL.
If Addresses capability is declared then addressing model becomes Physical64.
Additionally for Kernel capability interface variable ABI attributes are not
generated as entry point function is expected to have normal arguments.

Differential Revision: https://reviews.llvm.org/D85196

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

@@ -64,6 +64,10 @@ InterfaceVarABIAttr getInterfaceVarABIAttr(unsigned descriptorSet,
                                            Optional<StorageClass> storageClass,
                                            MLIRContext *context);
 
+/// Returns whether the given SPIR-V target (described by TargetEnvAttr) needs
+/// ABI attributes for interface variables (spv.interface_var_abi).
+bool needsInterfaceVarABIAttrs(TargetEnvAttr targetAttr);
+
 /// Returns the attribute name for specifying entry point information.
 StringRef getEntryPointABIAttrName();
 
@@ -100,6 +104,17 @@ TargetEnvAttr lookupTargetEnv(Operation *op);
 /// returned by getDefaultTargetEnv() if not provided.
 TargetEnvAttr lookupTargetEnvOrDefault(Operation *op);
 
+/// Returns addressing model selected based on target environment.
+AddressingModel getAddressingModel(TargetEnvAttr targetAttr);
+
+/// Returns execution model selected based on target environment.
+/// Returns failure if it cannot be selected.
+FailureOr<ExecutionModel> getExecutionModel(TargetEnvAttr targetAttr);
+
+/// Returns memory model selected based on target environment.
+/// Returns failure if it cannot be selected.
+FailureOr<MemoryModel> getMemoryModel(TargetEnvAttr targetAttr);
+
 } // namespace spirv
 } // namespace mlir
 

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVLowering.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/IR/Module.h"
 
 using namespace mlir;
@@ -170,9 +171,10 @@ lowerAsEntryFunction(gpu::GPUFuncOp funcOp, SPIRVTypeConverter &typeConverter,
                      "with no return values right now");
     return nullptr;
   }
-  if (fnType.getNumInputs() != argABIInfo.size()) {
+  if (!argABIInfo.empty() && fnType.getNumInputs() != argABIInfo.size()) {
     funcOp.emitError(
-        "lowering as entry functions requires ABI info for all arguments");
+        "lowering as entry functions requires ABI info for all arguments "
+        "or none of them");
     return nullptr;
   }
   // Update the signature to valid SPIR-V types and add the ABI
@@ -213,6 +215,10 @@ lowerAsEntryFunction(gpu::GPUFuncOp funcOp, SPIRVTypeConverter &typeConverter,
 static LogicalResult
 getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
                    SmallVectorImpl<spirv::InterfaceVarABIAttr> &argABI) {
+  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(funcOp);
+  if (!spirv::needsInterfaceVarABIAttrs(targetEnv))
+    return success();
+
   for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
     if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
             argIndex, spirv::getInterfaceVarABIAttrName()))
@@ -272,9 +278,15 @@ LogicalResult GPUFuncOpConversion::matchAndRewrite(
 LogicalResult GPUModuleConversion::matchAndRewrite(
     gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
+  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(moduleOp);
+  spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv);
+  FailureOr<spirv::MemoryModel> memoryModel = spirv::getMemoryModel(targetEnv);
+  if (failed(memoryModel))
+    return moduleOp.emitRemark("match failure: could not selected memory model "
+                               "based on 'spv.target_env'");
+
   auto spvModule = rewriter.create<spirv::ModuleOp>(
-      moduleOp.getLoc(), spirv::AddressingModel::Logical,
-      spirv::MemoryModel::GLSL450);
+      moduleOp.getLoc(), addressingModel, memoryModel.getValue());
 
   // Move the region from the module op into the SPIR-V module.
   Region &spvModuleRegion = spvModule.body();

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

@@ -653,7 +653,7 @@ mlir::spirv::setABIAttrs(spirv::FuncOp funcOp,
                          ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo) {
   // Set the attributes for argument and the function.
   StringRef argABIAttrName = spirv::getInterfaceVarABIAttrName();
-  for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
+  for (auto argIndex : llvm::seq<unsigned>(0, argABIInfo.size())) {
     funcOp.setArgAttr(argIndex, argABIAttrName, argABIInfo[argIndex]);
   }
   funcOp.setAttr(spirv::getEntryPointABIAttrName(), entryPointInfo);

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

@@ -90,6 +90,16 @@ spirv::getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding,
                                          context);
 }
 
+bool spirv::needsInterfaceVarABIAttrs(spirv::TargetEnvAttr targetAttr) {
+  for (spirv::Capability cap : targetAttr.getCapabilities()) {
+    if (cap == spirv::Capability::Kernel)
+      return false;
+    if (cap == spirv::Capability::Shader)
+      return true;
+  }
+  return false;
+}
+
 StringRef spirv::getEntryPointABIAttrName() { return "spv.entry_point_abi"; }
 
 spirv::EntryPointABIAttr
@@ -165,3 +175,37 @@ spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
 
   return getDefaultTargetEnv(op->getContext());
 }
+
+spirv::AddressingModel
+spirv::getAddressingModel(spirv::TargetEnvAttr targetAttr) {
+  for (spirv::Capability cap : targetAttr.getCapabilities()) {
+    // TODO: Physical64 is hard-coded here, but some information should come
+    // from TargetEnvAttr to selected between Physical32 and Physical64.
+    if (cap == Capability::Kernel)
+      return spirv::AddressingModel::Physical64;
+  }
+  // Logical addressing doesn't need any capabilities so return it as default.
+  return spirv::AddressingModel::Logical;
+}
+
+FailureOr<spirv::ExecutionModel>
+spirv::getExecutionModel(spirv::TargetEnvAttr targetAttr) {
+  for (spirv::Capability cap : targetAttr.getCapabilities()) {
+    if (cap == spirv::Capability::Kernel)
+      return spirv::ExecutionModel::Kernel;
+    if (cap == spirv::Capability::Shader)
+      return spirv::ExecutionModel::GLCompute;
+  }
+  return failure();
+}
+
+FailureOr<spirv::MemoryModel>
+spirv::getMemoryModel(spirv::TargetEnvAttr targetAttr) {
+  for (spirv::Capability cap : targetAttr.getCapabilities()) {
+    if (cap == spirv::Capability::Addresses)
+      return spirv::MemoryModel::OpenCL;
+    if (cap == spirv::Capability::Shader)
+      return spirv::MemoryModel::GLSL450;
+  }
+  return failure();
+}

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

@@ -119,8 +119,17 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
   if (failed(getInterfaceVariables(funcOp, interfaceVars))) {
     return failure();
   }
+
+  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnv(funcOp);
+  FailureOr<spirv::ExecutionModel> executionModel =
+      spirv::getExecutionModel(targetEnv);
+  if (failed(executionModel))
+    return funcOp.emitRemark("lower entry point failure: could not select "
+                             "execution model based on 'spv.target_env'");
+
   builder.create<spirv::EntryPointOp>(
-      funcOp.getLoc(), spirv::ExecutionModel::GLCompute, funcOp, interfaceVars);
+      funcOp.getLoc(), executionModel.getValue(), funcOp, interfaceVars);
+
   // Specifies the spv.ExecutionModeOp.
   auto localSizeAttr = entryPointAttr.local_size();
   SmallVector<int32_t, 3> localSize(localSizeAttr.getValues<int32_t>());

mlir/test/Conversion/GPUToSPIRV/module-structure-opencl.mlir

@@ -0,0 +1,32 @@
+// RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
+
+module attributes {
+  gpu.container_module,
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Kernel, Addresses], []>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+  gpu.module @kernels {
+    // CHECK-LABEL: spv.module Physical64 OpenCL
+    //       CHECK:   spv.func
+    //  CHECK-SAME:     {{%.*}}: f32
+    //   CHECK-NOT:     spv.interface_var_abi
+    //  CHECK-SAME:     {{%.*}}: !spv.ptr<!spv.struct<!spv.array<12 x f32, stride=4> [0]>, CrossWorkgroup>
+    //   CHECK-NOT:     spv.interface_var_abi
+    //  CHECK-SAME:     spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>}
+    gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, 11>) kernel
+        attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
+      gpu.return
+    }
+  }
+
+  func @main() {
+    %0 = "op"() : () -> (f32)
+    %1 = "op"() : () -> (memref<12xf32, 11>)
+    %cst = constant 1 : index
+    "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernels::@basic_module_structure }
+        : (index, index, index, index, index, index, f32, memref<12xf32, 11>) -> ()
+    return
+  }
+}

mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt -spirv-lower-abi-attrs -verify-diagnostics %s -o - | FileCheck %s
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Kernel, Addresses], []>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+  spv.module Physical64 OpenCL {
+    // CHECK-LABEL: spv.module
+    //       CHECK:   spv.func [[FN:@.*]](
+    //  CHECK-SAME:     {{%.*}}: f32
+    //  CHECK-SAME:     {{%.*}}: !spv.ptr<!spv.struct<!spv.array<12 x f32>>, CrossWorkgroup>
+    //       CHECK:   spv.EntryPoint "Kernel" [[FN]]
+    //       CHECK:   spv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1
+    spv.func @kernel(
+      %arg0: f32,
+      %arg1: !spv.ptr<!spv.struct<!spv.array<12 x f32>>, CrossWorkgroup>) "None"
+    attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} {
+      spv.Return
+    }
+  }
+}