diff --git a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp --- a/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp @@ -21,6 +21,7 @@ #include "mlir/Dialect/SPIRV/SPIRVDialect.h" #include "mlir/Dialect/SPIRV/SPIRVLowering.h" #include "mlir/Dialect/SPIRV/SPIRVOps.h" +#include "mlir/Pass/Pass.h" using namespace mlir; @@ -37,6 +38,26 @@ struct GPUToSPIRVPass : public ConvertGPUToSPIRVBase { void runOnOperation() override; }; + +/// Pass to set the spv.entry_point_abi +class TestSpirvEntryPointABIPass + : public PassWrapper> { +public: + TestSpirvEntryPointABIPass() = default; + TestSpirvEntryPointABIPass(ArrayRef sizes) { workgroupSize = sizes; } + TestSpirvEntryPointABIPass(const TestSpirvEntryPointABIPass &) {} + void runOnOperation() override; + +private: + Pass::ListOption workgroupSize{ + *this, "workgroup-size", + llvm::cl::desc( + "Workgroup size to use for all gpu.func kernels in the module, " + "specified with x-dimension first, y-dimension next and z-dimension " + "last. Unspecified dimensions will be set to 1"), + llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated}; +}; } // namespace void GPUToSPIRVPass::runOnOperation() { @@ -68,6 +89,32 @@ return signalPassFailure(); } +void TestSpirvEntryPointABIPass::runOnOperation() { + gpu::GPUModuleOp gpuModule = getOperation(); + MLIRContext *context = &getContext(); + StringRef attrName = spirv::getEntryPointABIAttrName(); + for (gpu::GPUFuncOp gpuFunc : gpuModule.getOps()) { + if (!gpu::GPUDialect::isKernel(gpuFunc) || gpuFunc.getAttr(attrName)) + continue; + SmallVector workgroupSize32 = llvm::to_vector<3>( + llvm::map_range(workgroupSize, [](int64_t v) -> int32_t { + return static_cast(v); + })); + workgroupSize32.resize(3, 1); + gpuFunc.setAttr(attrName, + spirv::getEntryPointABIAttr(workgroupSize32, context)); + } +} + std::unique_ptr> mlir::createConvertGPUToSPIRVPass() { return std::make_unique(); } + +namespace mlir { +void registerTestSpirvEntryPointABIPass() { + PassRegistration registration( + "test-spirv-entry-point-abi", + "Set the spv.entry_point_abi attribute on GPU kernel function within the " + "module, intended for testing only"); +} +} // namespace mlir diff --git a/mlir/test/Conversion/GPUToSPIRV/test_spirv_entry_point.mlir b/mlir/test/Conversion/GPUToSPIRV/test_spirv_entry_point.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUToSPIRV/test_spirv_entry_point.mlir @@ -0,0 +1,14 @@ +// RUN: mlir-opt -test-spirv-entry-point-abi %s | FileCheck %s -check-prefix=DEFAULT +// RUN: mlir-opt -test-spirv-entry-point-abi="workgroup-size=32" %s | FileCheck %s -check-prefix=WG32 + +// DEFAULT: gpu.func @foo() +// DEFAULT-SAME: spv.entry_point_abi = {local_size = dense<1> : vector<3xi32>} + +// WG32: gpu.func @foo() +// WG32-SAME: spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>} + +gpu.module @kernels { + gpu.func @foo() kernel { + gpu.return + } +} \ No newline at end of file diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -65,6 +65,7 @@ void registerTestPreparationPassWithAllowedMemrefResults(); void registerTestReducer(); void registerTestGpuParallelLoopMappingPass(); +void registerTestSpirvEntryPointABIPass(); void registerTestSCFUtilsPass(); void registerTestVectorConversions(); void registerVectorizerTestPass(); @@ -140,6 +141,7 @@ registerTestPreparationPassWithAllowedMemrefResults(); registerTestReducer(); registerTestGpuParallelLoopMappingPass(); + registerTestSpirvEntryPointABIPass(); registerTestSCFUtilsPass(); registerTestVectorConversions(); registerVectorizerTestPass();