diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.td b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.td index 22fd54221c33..628cf849d85b 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.td @@ -27,7 +27,7 @@ include "mlir/Dialect/SPIRV/IR/SPIRVBase.td" // points in the generated SPIR-V module: // 1) WorkGroup Size. def SPV_EntryPointABIAttr : StructAttr<"EntryPointABIAttr", SPIRV_Dialect, [ - StructFieldAttr<"local_size", I32ElementsAttr> + StructFieldAttr<"local_size", OptionalAttr> ]>; def SPV_ExtensionArrayAttr : TypedArrayAttrBase< diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp index 8c5627c0aa8a..546b0ac38f8d 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -55,7 +55,8 @@ public: /// attribute on the surrounding FuncOp is used to replace the gpu::BlockDimOp. class WorkGroupSizeConversion : public OpConversionPattern { public: - using OpConversionPattern::OpConversionPattern; + WorkGroupSizeConversion(TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context, /*benefit*/ 10) {} LogicalResult matchAndRewrite(gpu::BlockDimOp op, OpAdaptor adaptor, @@ -159,6 +160,9 @@ LogicalResult WorkGroupSizeConversion::matchAndRewrite( gpu::BlockDimOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op); + if (!workGroupSizeAttr) + return failure(); + auto val = workGroupSizeAttr .getValues()[static_cast(op.dimension())]; auto convertedType = @@ -366,6 +370,7 @@ void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter, GPUModuleEndConversion, GPUReturnOpConversion, LaunchConfigConversion, LaunchConfigConversion, + LaunchConfigConversion, LaunchConfigConversion, SingleDimLaunchConfigConversion localSize, MLIRContext *context) { + if (localSize.empty()) + return spirv::EntryPointABIAttr::get(nullptr, context); + assert(localSize.size() == 3); return spirv::EntryPointABIAttr::get( DenseElementsAttr::get( diff --git a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp index 6094ad8bf224..71042491c57a 100644 --- a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp +++ b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp @@ -136,10 +136,13 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp, // Specifies the spv.ExecutionModeOp. auto localSizeAttr = entryPointAttr.local_size(); - SmallVector localSize(localSizeAttr.getValues()); - builder.create( - funcOp.getLoc(), funcOp, spirv::ExecutionMode::LocalSize, localSize); - funcOp->removeAttr(entryPointAttrName); + if (localSizeAttr) { + auto values = localSizeAttr.getValues(); + SmallVector localSize(values); + builder.create( + funcOp.getLoc(), funcOp, spirv::ExecutionMode::LocalSize, localSize); + funcOp->removeAttr(entryPointAttrName); + } return success(); } diff --git a/mlir/test/Conversion/GPUToSPIRV/builtins.mlir b/mlir/test/Conversion/GPUToSPIRV/builtins.mlir index 43cacf23e7a9..edbd9839ce69 100644 --- a/mlir/test/Conversion/GPUToSPIRV/builtins.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/builtins.mlir @@ -223,6 +223,78 @@ module attributes {gpu.container_module} { // ----- +module attributes {gpu.container_module} { + func @builtin() { + %c0 = arith.constant 1 : index + gpu.launch_func @kernels::@builtin_workgroup_size_x + blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) + return + } + + // CHECK-LABEL: spv.module @{{.*}} + // CHECK: spv.GlobalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize") + gpu.module @kernels { + gpu.func @builtin_workgroup_size_x() kernel + attributes {spv.entry_point_abi = {}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPSIZE]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}} + %0 = gpu.block_dim x + gpu.return + } + } +} + +// ----- + +module attributes {gpu.container_module} { + func @builtin() { + %c0 = arith.constant 1 : index + gpu.launch_func @kernels::@builtin_workgroup_size_y + blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) + return + } + + // CHECK-LABEL: spv.module @{{.*}} + // CHECK: spv.GlobalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize") + gpu.module @kernels { + gpu.func @builtin_workgroup_size_y() kernel + attributes {spv.entry_point_abi = {}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPSIZE]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}} + %0 = gpu.block_dim y + gpu.return + } + } +} + +// ----- + +module attributes {gpu.container_module} { + func @builtin() { + %c0 = arith.constant 1 : index + gpu.launch_func @kernels::@builtin_workgroup_size_z + blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) + return + } + + // CHECK-LABEL: spv.module @{{.*}} + // CHECK: spv.GlobalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize") + gpu.module @kernels { + gpu.func @builtin_workgroup_size_z() kernel + attributes {spv.entry_point_abi = {}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPSIZE]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}} + %0 = gpu.block_dim z + gpu.return + } + } +} + +// ----- + module attributes {gpu.container_module} { // CHECK-LABEL: spv.module @{{.*}} Logical GLSL450 // CHECK: spv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize")