diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h --- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h +++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h @@ -64,6 +64,10 @@ Optional 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 @@ /// 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 getExecutionModel(TargetEnvAttr targetAttr); + +/// Returns memory model selected based on target environment. +/// Returns failure if it cannot be selected. +FailureOr getMemoryModel(TargetEnvAttr targetAttr); + } // namespace spirv } // namespace mlir diff --git a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp --- a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp +++ b/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 @@ "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 @@ static LogicalResult getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp, SmallVectorImpl &argABI) { + spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(funcOp); + if (!spirv::needsInterfaceVarABIAttrs(targetEnv)) + return success(); + for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { if (funcOp.getArgAttrOfType( argIndex, spirv::getInterfaceVarABIAttrName())) @@ -272,9 +278,15 @@ LogicalResult GPUModuleConversion::matchAndRewrite( gpu::GPUModuleOp moduleOp, ArrayRef operands, ConversionPatternRewriter &rewriter) const { + spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(moduleOp); + spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv); + FailureOr 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( - 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(); diff --git a/mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp b/mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp --- a/mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp +++ b/mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp @@ -653,7 +653,7 @@ ArrayRef argABIInfo) { // Set the attributes for argument and the function. StringRef argABIAttrName = spirv::getInterfaceVarABIAttrName(); - for (auto argIndex : llvm::seq(0, funcOp.getNumArguments())) { + for (auto argIndex : llvm::seq(0, argABIInfo.size())) { funcOp.setArgAttr(argIndex, argABIAttrName, argABIInfo[argIndex]); } funcOp.setAttr(spirv::getEntryPointABIAttrName(), entryPointInfo); diff --git a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp --- a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp +++ b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp @@ -90,6 +90,16 @@ 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 @@ 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::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::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(); +} diff --git a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp --- a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp +++ b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp @@ -119,8 +119,17 @@ if (failed(getInterfaceVariables(funcOp, interfaceVars))) { return failure(); } + + spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnv(funcOp); + FailureOr 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( - funcOp.getLoc(), spirv::ExecutionModel::GLCompute, funcOp, interfaceVars); + funcOp.getLoc(), executionModel.getValue(), funcOp, interfaceVars); + // Specifies the spv.ExecutionModeOp. auto localSizeAttr = entryPointAttr.local_size(); SmallVector localSize(localSizeAttr.getValues()); diff --git a/mlir/test/Conversion/GPUToSPIRV/module-structure-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/module-structure-opencl.mlir new file mode 100644 --- /dev/null +++ b/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, + {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 [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 + } +} diff --git a/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir b/mlir/test/Dialect/SPIRV/Transforms/abi-interface-opencl.mlir new file mode 100644 --- /dev/null +++ b/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, + {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>, CrossWorkgroup> + // CHECK: spv.EntryPoint "Kernel" [[FN]] + // CHECK: spv.ExecutionMode [[FN]] "LocalSize", 32, 1, 1 + spv.func @kernel( + %arg0: f32, + %arg1: !spv.ptr>, CrossWorkgroup>) "None" + attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} { + spv.Return + } + } +}