diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h --- a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h +++ b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h @@ -121,7 +121,8 @@ 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. diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -283,9 +283,8 @@ /// 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 &argABI) { - spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(funcOp); if (!spirv::needsInterfaceVarABIAttrs(targetEnv)) return success(); @@ -298,7 +297,8 @@ std::optional 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(); } @@ -309,8 +309,10 @@ if (!gpu::GPUDialect::isKernel(funcOp)) return failure(); + auto *typeConverter = getTypeConverter(); SmallVector argABI; - if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) { + if (failed( + getDefaultABIAttrs(typeConverter->getTargetEnv(), funcOp, argABI))) { argABI.clear(); for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { // If the ABI is already specified, use it. @@ -349,12 +351,14 @@ 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(); + const spirv::TargetEnv &targetEnv = typeConverter->getTargetEnv(); + spirv::AddressingModel addressingModel = spirv::getAddressingModel( + targetEnv, typeConverter->getOptions().use64bitIndex); FailureOr 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(); diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp --- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp +++ b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp @@ -205,12 +205,12 @@ } 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 @@ -235,7 +235,7 @@ FailureOr 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; diff --git a/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir @@ -0,0 +1,52 @@ +// 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, #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, Input> + // INDEX64-LABEL: spirv.module @{{.*}} Physical64 OpenCL + // INDEX64: spirv.GlobalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId") : !spirv.ptr, Input> + gpu.module @kernels { + gpu.func @builtin_workgroup_id_x() kernel + attributes {spirv.entry_point_abi = #spirv.entry_point_abi} { + // 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, #spirv.resource_limits<>> +} { + // INDEX32-LABEL: spirv.module @{{.*}} Physical32 OpenCL + // INDEX32: spirv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize") : !spirv.ptr + // INDEX64-LABEL: spirv.module @{{.*}} Physical64 OpenCL + // INDEX64: spirv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize") : !spirv.ptr + gpu.module @kernels { + gpu.func @builtin_subgroup_size() kernel + attributes {spirv.entry_point_abi = #spirv.entry_point_abi} { + // 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 + } + } +} diff --git a/mlir/test/Conversion/GPUToSPIRV/builtins.mlir b/mlir/test/Conversion/GPUToSPIRV/builtins-vulkan.mlir rename from mlir/test/Conversion/GPUToSPIRV/builtins.mlir rename to mlir/test/Conversion/GPUToSPIRV/builtins-vulkan.mlir diff --git a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir --- a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir @@ -1,4 +1,4 @@ -// 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,