[mlir][spirv] Fix Physical32/Physical64 support for OpenCL

[antiagainst]

Summary

We use use64bitIndex in the option to decide the target device
address bitwidth. This makes it consistent with index type
conversion too.

[kuhar]

The patch LGTM, but I have a design question: have you considered moving PhysicalXY from the module level to target env? I understand that this follows SPIR-V proper, but what do we gain much by keeping those separate?

[antiagainst]

In D144827#4153899, @kuhar wrote:

The patch LGTM, but I have a design question: have you considered moving PhysicalXY from the module level to target env? I understand that this follows SPIR-V proper, but what do we gain much by keeping those separate?

Target environment is actually the requirements towards which we are CodeGen, where these Logical/PhysicalNN is the derived module addressing model details to meeting those requirements. Those are part of the final SPIR-V blob; it's more for driver consumption. So I don't think we should require the developer to figure them out; that's the job of the compiler?

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

 TargetEnvAttr lookupTargetEnv(Operation *op);
 
 /// Queries the target environment recursively from enclosing symbol table ops
 /// containing the given `op` or returns the default target environment as
 /// returned by getDefaultTargetEnv() if not provided.
 TargetEnvAttr lookupTargetEnvOrDefault(Operation *op);
 
 /// Returns addressing model selected based on target environment.
-AddressingModel getAddressingModel(TargetEnvAttr targetAttr);
+AddressingModel getAddressingModel(TargetEnvAttr targetAttr,
+                                   bool use64bitAddress);
 
 /// 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
 
 #endif // MLIR_DIALECT_SPIRV_IR_TARGETANDABI_H

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

lowerAsEntryFunction(gpu::GPUFuncOp funcOp, TypeConverter &typeConverter,
 
   return newFuncOp;
 }
 
 /// Populates `argABI` with spirv.interface_var_abi attributes for lowering
 /// gpu.func to spirv.func if no arguments have the attributes set
 /// already. Returns failure if any argument has the ABI attribute set already.
 static LogicalResult
-getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
+getDefaultABIAttrs(const spirv::TargetEnv &targetEnv, 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()))
       return failure();
     // Vulkan's interface variable requirements needs scalars to be wrapped in a
     // struct. The struct held in storage buffer.
     std::optional<spirv::StorageClass> sc;
     if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
       sc = spirv::StorageClass::StorageBuffer;
-    argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
+    argABI.push_back(
+        spirv::getInterfaceVarABIAttr(0, argIndex, sc, funcOp.getContext()));
   }
   return success();
 }
 
 LogicalResult GPUFuncOpConversion::matchAndRewrite(
     gpu::GPUFuncOp funcOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (!gpu::GPUDialect::isKernel(funcOp))
     return failure();
 
+  auto *typeConverter = getTypeConverter<SPIRVTypeConverter>();
   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
-  if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
+  if (failed(
+          getDefaultABIAttrs(typeConverter->getTargetEnv(), funcOp, argABI))) {
     argABI.clear();
     for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
       // If the ABI is already specified, use it.
       auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
           argIndex, spirv::getInterfaceVarABIAttrName());
       if (!abiAttr) {
         funcOp.emitRemark(
             "match failure: missing 'spirv.interface_var_abi' attribute at "

 
 //===----------------------------------------------------------------------===//
 // ModuleOp with gpu.module.
 //===----------------------------------------------------------------------===//
 
 LogicalResult GPUModuleConversion::matchAndRewrite(
     gpu::GPUModuleOp moduleOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
-  spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(moduleOp);
-  spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv);
+  auto *typeConverter = getTypeConverter<SPIRVTypeConverter>();
+  const spirv::TargetEnv &targetEnv = typeConverter->getTargetEnv();
+  spirv::AddressingModel addressingModel = spirv::getAddressingModel(
+      targetEnv, typeConverter->getOptions().use64bitIndex);
   FailureOr<spirv::MemoryModel> memoryModel = spirv::getMemoryModel(targetEnv);
   if (failed(memoryModel))
-    return moduleOp.emitRemark("match failure: could not selected memory model "
-                               "based on 'spirv.target_env'");
+    return moduleOp.emitRemark(
+        "cannot deduce memory model from 'spirv.target_env'");
 
   // Add a keyword to the module name to avoid symbolic conflict.
   std::string spvModuleName = (kSPIRVModule + moduleOp.getName()).str();
   auto spvModule = rewriter.create<spirv::ModuleOp>(
       moduleOp.getLoc(), addressingModel, *memoryModel, std::nullopt,
       StringRef(spvModuleName));
 
   // Move the region from the module op into the SPIR-V module.

mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp

 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
   if (spirv::TargetEnvAttr attr = spirv::lookupTargetEnv(op))
     return attr;
 
   return getDefaultTargetEnv(op->getContext());
 }
 
 spirv::AddressingModel
-spirv::getAddressingModel(spirv::TargetEnvAttr targetAttr) {
+spirv::getAddressingModel(spirv::TargetEnvAttr targetAttr,
+                          bool use64bitAddress) {
   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;
+      return use64bitAddress ? spirv::AddressingModel::Physical64
+                             : spirv::AddressingModel::Physical32;
     // TODO PhysicalStorageBuffer64 is hard-coded here, but some information
     // should come from TargetEnvAttr to select between PhysicalStorageBuffer64
     // and PhysicalStorageBuffer64EXT
     if (cap == Capability::PhysicalStorageBufferAddresses)
       return spirv::AddressingModel::PhysicalStorageBuffer64;
   }
   // 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)
+    if (cap == spirv::Capability::Kernel)
       return spirv::MemoryModel::OpenCL;
     if (cap == spirv::Capability::Shader)
       return spirv::MemoryModel::GLSL450;
   }
   return failure();
 }

mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir

+// RUN: mlir-opt -split-input-file -convert-gpu-to-spirv="use-64bit-index=false" %s -o - | FileCheck %s --check-prefix=INDEX32
+// RUN: mlir-opt -split-input-file -convert-gpu-to-spirv="use-64bit-index=true" %s -o - | FileCheck %s --check-prefix=INDEX64
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Kernel, Int64], []>, #spirv.resource_limits<>>
+} {
+  func.func @builtin() {
+    %c0 = arith.constant 1 : index
+    gpu.launch_func @kernels::@builtin_workgroup_id_x
+        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
+    return
+  }
+
+  // INDEX32-LABEL:  spirv.module @{{.*}} Physical32 OpenCL
+  // INDEX32: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId") : !spirv.ptr<vector<3xi32>, Input>
+  // INDEX64-LABEL:  spirv.module @{{.*}} Physical64 OpenCL
+  // INDEX64: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId") : !spirv.ptr<vector<3xi64>, Input>
+  gpu.module @kernels {
+    gpu.func @builtin_workgroup_id_x() kernel
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+      // INDEX32: [[ADDRESS:%.*]] = spirv.mlir.addressof [[WORKGROUPID]]
+      // INDEX32-NEXT: [[VEC:%.*]] = spirv.Load "Input" [[ADDRESS]]
+      // INDEX32-NEXT: {{%.*}} = spirv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
+      // INDEX64-NOT: spirv.UConvert
+      %0 = gpu.block_id x
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Kernel, Int64], []>, #spirv.resource_limits<>>
+} {
+  // INDEX32-LABEL:  spirv.module @{{.*}} Physical32 OpenCL
+  // INDEX32: spirv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize") : !spirv.ptr<i32, Input>
+  // INDEX64-LABEL:  spirv.module @{{.*}} Physical64 OpenCL
+  // INDEX64: spirv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize") : !spirv.ptr<i32, Input>
+  gpu.module @kernels {
+    gpu.func @builtin_subgroup_size() kernel
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+      // INDEX32: [[ADDRESS:%.*]] = spirv.mlir.addressof [[SUBGROUPSIZE]]
+      // INDEX32-NEXT: {{%.*}} = spirv.Load "Input" [[ADDRESS]]
+      // INDEX64: spirv.UConvert %{{.+}} : i32 to i64
+      %0 = gpu.subgroup_size : index
+      gpu.return
+    }
+  }
+}

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

-// RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -verify-diagnostics -split-input-file %s -o - | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv="use-64bit-index=true" -verify-diagnostics -split-input-file %s -o - | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel, Addresses], []>, #spirv.resource_limits<>>
 } {
   gpu.module @kernels {
     // CHECK-LABEL: spirv.module @{{.*}} Physical64 OpenCL
     //       CHECK:   spirv.func