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/test/lib/Dialect/SPIRV/CMakeLists.txt b/mlir/test/lib/Dialect/SPIRV/CMakeLists.txt --- a/mlir/test/lib/Dialect/SPIRV/CMakeLists.txt +++ b/mlir/test/lib/Dialect/SPIRV/CMakeLists.txt @@ -1,6 +1,7 @@ # Exclude tests from libMLIR.so add_mlir_library(MLIRSPIRVTestPasses TestAvailability.cpp + TestEntryPointAbi.cpp EXCLUDE_FROM_LIBMLIR @@ -9,6 +10,7 @@ ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR LINK_LIBS PUBLIC + MLIRGPU MLIRIR MLIRPass MLIRSPIRV diff --git a/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp b/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp @@ -0,0 +1,67 @@ +//===- TestAvailability.cpp - Test pass for setting Entry point ABI info --===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements a pass that sets the spv.entry_point_abi attribute on +// functions that are to be lowered as entry point functions. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/GPU/GPUDialect.h" +#include "mlir/Dialect/SPIRV/SPIRVDialect.h" +#include "mlir/Dialect/SPIRV/TargetAndABI.h" +#include "mlir/Pass/Pass.h" + +using namespace mlir; + +namespace { +/// 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 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)); + } +} + +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/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();