Skip to content

Commit 7724be9

Browse files
andfau-amdandfau-amd
authored
[mlir][spirv] Do SPIR-V serialization in -test-vulkan-runner-pipeline (llvm#121494)
This commit is a further incremental step toward moving the whole mlir-vulkan-runner MLIR pass pipeline into mlir-opt (see llvm#73457). The previous step was b225b3adf7b78387c9fcb97a3ff0e0a1e26eafe2, which moved all device passes prior to SPIR-V serialization into a new mlir-opt test pass, `-test-vulkan-runner-pipeline`. This commit changes how SPIR-V serialization is accomplished for Vulkan runner tests. Until now, this was done by the Vulkan-specific ConvertGpuLaunchFuncToVulkanLaunchFunc pass. With this commit, this responsibility is removed from that pass, and is instead done with the existing generic GpuModuleToBinaryPass. In addition, the SPIR-V serialization step is no longer done inside mlir-vulkan-runner, but rather inside mlir-opt (in the `-test-vulkan-runner-pipeline` pass). Both of these changes represent a greater alignment between mlir-vulkan-runner and the other GPU integration tests. Notably, the IR shapes produced by the mlir-opt pipelines for the Vulkan and SYCL runners are now much more similar, with both using a gpu.binary op for the serialized SPIR-V kernel. In order to enable this, this commit includes these supporting changes: - ConvertToSPIRVPass is enhanced to support producing the IR shape where a spirv.module is nested inside a gpu.module, since this is what GpuModuleToBinaryPass expects. - ConvertGPULaunchFuncToVulkanLaunchFunc is changed to remove its SPIR-V serialization functionality, and instead now extracts the SPIR-V from a gpu.binary operation (as produced by ConvertToSPIRVPass). - `-test-vulkan-runner-pipeline` now attaches SPIR-V target information required by GpuModuleToBinaryPass. - The WebGPU pass option, which had been removed from mlir-vulkan-runner in the previous commit in this series, is restored as an option to `-test-vulkan-runner-pipeline` instead, so that the WebGPU pass continues being inserted into the pipeline just before SPIR-V serialization.
1 parent 473510a commit 7724be9

File tree

10 files changed

+121
-53
lines changed

10 files changed

+121
-53
lines changed

‎mlir/include/mlir/Conversion/Passes.td‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,10 @@ def ConvertToSPIRVPass : Pass<"convert-to-spirv"> {
6161
"Run vector unrolling to convert vector types in function bodies">,
6262
Option<"convertGPUModules", "convert-gpu-modules", "bool",
6363
/*default=*/"false",
64-
"Clone and convert GPU modules">
64+
"Clone and convert GPU modules">,
65+
Option<"nestInGPUModule", "nest-in-gpu-module", "bool",
66+
/*default=*/"false",
67+
"Put converted SPIR-V module inside the gpu.module instead of alongside it.">,
6568
];
6669
}
6770

‎mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,10 @@ struct ConvertToSPIRVPass final
108108
SmallVector<Operation *, 1> gpuModules;
109109
OpBuilder builder(context);
110110
op->walk([&](gpu::GPUModuleOp gpuModule) {
111-
builder.setInsertionPoint(gpuModule);
111+
if (nestInGPUModule)
112+
builder.setInsertionPointToStart(gpuModule.getBody());
113+
else
114+
builder.setInsertionPoint(gpuModule);
112115
gpuModules.push_back(builder.clone(*gpuModule));
113116
});
114117
// Run conversion for each module independently as they can have

‎mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ void GPUToSPIRVPass::runOnOperation() {
7171
// launch op still needs the original GPU kernel module.
7272
// For Vulkan Shader capabilities, we insert the newly converted SPIR-V
7373
// module right after the original GPU module, as that's the expectation of
74-
// the in-tree Vulkan runner.
74+
// the in-tree SPIR-V CPU runner (the Vulkan runner does not use this pass).
7575
// For OpenCL Kernel capabilities, we insert the newly converted SPIR-V
7676
// module inside the original GPU module, as that's the expectaion of the
7777
// normal GPU compilation pipeline.

‎mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp‎

Lines changed: 36 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,8 @@
77
//===----------------------------------------------------------------------===//
88
//
99
// This file implements a pass to convert gpu launch function into a vulkan
10-
// launch function. Creates a SPIR-V binary shader from the `spirv::ModuleOp`
11-
// using `spirv::serialize` function, attaches binary data and entry point name
12-
// as an attributes to vulkan launch call op.
10+
// launch function. Extracts the SPIR-V from a `gpu::BinaryOp` and attaches it
11+
// along with the entry point name as attributes to a Vulkan launch call op.
1312
//
1413
//===----------------------------------------------------------------------===//
1514

@@ -40,21 +39,19 @@ static constexpr const char *kVulkanLaunch = "vulkanLaunch";
4039

4140
namespace {
4241

43-
/// A pass to convert gpu launch op to vulkan launch call op, by creating a
44-
/// SPIR-V binary shader from `spirv::ModuleOp` using `spirv::serialize`
45-
/// function and attaching binary data and entry point name as an attributes to
46-
/// created vulkan launch call op.
42+
/// A pass to convert gpu launch op to vulkan launch call op, by extracting a
43+
/// SPIR-V binary shader from a `gpu::BinaryOp` and attaching binary data and
44+
/// entry point name as an attributes to created vulkan launch call op.
4745
class ConvertGpuLaunchFuncToVulkanLaunchFunc
4846
: public impl::ConvertGpuLaunchFuncToVulkanLaunchFuncBase<
4947
ConvertGpuLaunchFuncToVulkanLaunchFunc> {
5048
public:
5149
void runOnOperation() override;
5250

5351
private:
54-
/// Creates a SPIR-V binary shader from the given `module` using
55-
/// `spirv::serialize` function.
56-
LogicalResult createBinaryShader(ModuleOp module,
57-
std::vector<char> &binaryShader);
52+
/// Extracts a SPIR-V binary shader from the given `module`, if any.
53+
/// Note that this also removes the binary from the IR.
54+
FailureOr<StringAttr> getBinaryShader(ModuleOp module);
5855

5956
/// Converts the given `launchOp` to vulkan launch call.
6057
void convertGpuLaunchFunc(gpu::LaunchFuncOp launchOp);
@@ -135,22 +132,35 @@ LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::declareVulkanLaunchFunc(
135132
return success();
136133
}
137134

138-
LogicalResult ConvertGpuLaunchFuncToVulkanLaunchFunc::createBinaryShader(
139-
ModuleOp module, std::vector<char> &binaryShader) {
135+
FailureOr<StringAttr>
136+
ConvertGpuLaunchFuncToVulkanLaunchFunc::getBinaryShader(ModuleOp module) {
140137
bool done = false;
141-
SmallVector<uint32_t, 0> binary;
142-
for (auto spirvModule : module.getOps<spirv::ModuleOp>()) {
138+
StringAttr binaryAttr;
139+
gpu::BinaryOp binaryToErase;
140+
for (auto gpuBinary : module.getOps<gpu::BinaryOp>()) {
143141
if (done)
144-
return spirvModule.emitError("should only contain one 'spirv.module' op");
142+
return gpuBinary.emitError("should only contain one 'gpu.binary' op");
145143
done = true;
146144

147-
if (failed(spirv::serialize(spirvModule, binary)))
148-
return failure();
145+
ArrayRef<Attribute> objects = gpuBinary.getObjectsAttr().getValue();
146+
if (objects.size() != 1)
147+
return gpuBinary.emitError("should only contain a single object");
148+
149+
auto object = cast<gpu::ObjectAttr>(objects[0]);
150+
151+
if (!isa<spirv::TargetEnvAttr>(object.getTarget()))
152+
return gpuBinary.emitError(
153+
"should contain an object with a SPIR-V target environment");
154+
155+
binaryAttr = object.getObject();
156+
binaryToErase = gpuBinary;
149157
}
150-
binaryShader.resize(binary.size() * sizeof(uint32_t));
151-
std::memcpy(binaryShader.data(), reinterpret_cast<char *>(binary.data()),
152-
binaryShader.size());
153-
return success();
158+
if (!done)
159+
return module.emitError("should contain a 'gpu.binary' op");
160+
161+
// Remove the binary to avoid confusing later conversion passes.
162+
binaryToErase.erase();
163+
return binaryAttr;
154164
}
155165

156166
void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
@@ -159,9 +169,9 @@ void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
159169
OpBuilder builder(launchOp);
160170
Location loc = launchOp.getLoc();
161171

162-
// Serialize `spirv::Module` into binary form.
163-
std::vector<char> binary;
164-
if (failed(createBinaryShader(module, binary)))
172+
FailureOr<StringAttr> binaryAttr = getBinaryShader(module);
173+
// Extract SPIR-V from `gpu.binary` op.
174+
if (failed(binaryAttr))
165175
return signalPassFailure();
166176

167177
// Declare vulkan launch function.
@@ -182,9 +192,7 @@ void ConvertGpuLaunchFuncToVulkanLaunchFunc::convertGpuLaunchFunc(
182192
vulkanLaunchOperands);
183193

184194
// Set SPIR-V binary shader data as an attribute.
185-
vulkanLaunchCallOp->setAttr(
186-
kSPIRVBlobAttrName,
187-
builder.getStringAttr(StringRef(binary.data(), binary.size())));
195+
vulkanLaunchCallOp->setAttr(kSPIRVBlobAttrName, *binaryAttr);
188196

189197
// Set entry point name as an attribute.
190198
vulkanLaunchCallOp->setAttr(kSPIRVEntryPointAttrName,

‎mlir/test/Conversion/ConvertToSPIRV/convert-gpu-modules-nested.mlir‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// RUN: mlir-opt -convert-to-spirv="convert-gpu-modules=true nest-in-gpu-module=true run-signature-conversion=false run-vector-unrolling=false" %s | FileCheck %s
2+
3+
module attributes {
4+
gpu.container_module,
5+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], []>, #spirv.resource_limits<>>
6+
} {
7+
// CHECK-LABEL: func.func @main
8+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
9+
// CHECK: gpu.launch_func @[[$KERNELS_1:.*]]::@[[$BUILTIN_WG_ID_X:.*]] blocks in (%[[C1]], %[[C1]], %[[C1]]) threads in (%[[C1]], %[[C1]], %[[C1]])
10+
func.func @main() {
11+
%c1 = arith.constant 1 : index
12+
gpu.launch_func @kernels_1::@builtin_workgroup_id_x
13+
blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)
14+
return
15+
}
16+
17+
// CHECK: gpu.module @[[$KERNELS_1]]
18+
// CHECK: spirv.module @{{.*}} Logical GLSL450
19+
// CHECK: spirv.func @[[$BUILTIN_WG_ID_X]]
20+
// CHECK: spirv.mlir.addressof
21+
// CHECK: spirv.Load "Input"
22+
// CHECK: spirv.CompositeExtract
23+
gpu.module @kernels_1 {
24+
gpu.func @builtin_workgroup_id_x() kernel
25+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
26+
%0 = gpu.block_id x
27+
gpu.return
28+
}
29+
}
30+
}

‎mlir/test/Conversion/GPUToVulkan/lower-gpu-launch-vulkan-launch.mlir‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,24 @@
1-
// RUN: mlir-opt %s -convert-gpu-launch-to-vulkan-launch | FileCheck %s
1+
// RUN: mlir-opt %s -pass-pipeline='builtin.module(spirv-attach-target{ver=v1.0 caps=Shader exts=SPV_KHR_storage_buffer_storage_class},gpu-module-to-binary,convert-gpu-launch-to-vulkan-launch)' | FileCheck %s
22

33
// CHECK: %[[resource:.*]] = memref.alloc() : memref<12xf32>
44
// CHECK: %[[index:.*]] = arith.constant 1 : index
55
// CHECK: call @vulkanLaunch(%[[index]], %[[index]], %[[index]], %[[resource]]) {spirv_blob = "{{.*}}", spirv_element_types = [f32], spirv_entry_point = "kernel"}
66

77
module attributes {gpu.container_module} {
8-
spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]> {
9-
spirv.GlobalVariable @kernel_arg_0 bind(0, 0) : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
10-
spirv.func @kernel() "None" attributes {workgroup_attributions = 0 : i64} {
11-
%0 = spirv.mlir.addressof @kernel_arg_0 : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
12-
%2 = spirv.Constant 0 : i32
13-
%3 = spirv.mlir.addressof @kernel_arg_0 : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
14-
%4 = spirv.AccessChain %0[%2, %2] : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<f32, StorageBuffer>
15-
%5 = spirv.Load "StorageBuffer" %4 : f32
16-
spirv.Return
17-
}
18-
spirv.EntryPoint "GLCompute" @kernel
19-
spirv.ExecutionMode @kernel "LocalSize", 1, 1, 1
20-
}
218
gpu.module @kernels {
9+
spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]> {
10+
spirv.GlobalVariable @kernel_arg_0 bind(0, 0) : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
11+
spirv.func @kernel() "None" attributes {workgroup_attributions = 0 : i64} {
12+
%0 = spirv.mlir.addressof @kernel_arg_0 : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
13+
%2 = spirv.Constant 0 : i32
14+
%3 = spirv.mlir.addressof @kernel_arg_0 : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>
15+
%4 = spirv.AccessChain %0[%2, %2] : !spirv.ptr<!spirv.struct<(!spirv.array<12 x f32, stride=4> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<f32, StorageBuffer>
16+
%5 = spirv.Load "StorageBuffer" %4 : f32
17+
spirv.Return
18+
}
19+
spirv.EntryPoint "GLCompute" @kernel
20+
spirv.ExecutionMode @kernel "LocalSize", 1, 1, 1
21+
}
2222
gpu.func @kernel(%arg0: memref<12xf32>) kernel {
2323
gpu.return
2424
}

‎mlir/test/lib/Pass/TestVulkanRunnerPipeline.cpp‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,33 +12,57 @@
1212

1313
#include "mlir/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.h"
1414
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
15+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1516
#include "mlir/Dialect/GPU/Transforms/Passes.h"
1617
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
1718
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
1819
#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
1920
#include "mlir/Pass/PassManager.h"
21+
#include "mlir/Pass/PassOptions.h"
2022

2123
using namespace mlir;
2224

2325
namespace {
2426

25-
void buildTestVulkanRunnerPipeline(OpPassManager &passManager) {
27+
struct VulkanRunnerPipelineOptions
28+
: PassPipelineOptions<VulkanRunnerPipelineOptions> {
29+
Option<bool> spirvWebGPUPrepare{
30+
*this, "spirv-webgpu-prepare",
31+
llvm::cl::desc("Run MLIR transforms used when targetting WebGPU")};
32+
};
33+
34+
void buildTestVulkanRunnerPipeline(OpPassManager &passManager,
35+
const VulkanRunnerPipelineOptions &options) {
2636
passManager.addPass(createGpuKernelOutliningPass());
2737
passManager.addPass(memref::createFoldMemRefAliasOpsPass());
2838

39+
GpuSPIRVAttachTargetOptions attachTargetOptions{};
40+
attachTargetOptions.spirvVersion = "v1.0";
41+
attachTargetOptions.spirvCapabilities.push_back("Shader");
42+
attachTargetOptions.spirvExtensions.push_back(
43+
"SPV_KHR_storage_buffer_storage_class");
44+
passManager.addPass(createGpuSPIRVAttachTarget(attachTargetOptions));
45+
2946
ConvertToSPIRVPassOptions convertToSPIRVOptions{};
3047
convertToSPIRVOptions.convertGPUModules = true;
48+
convertToSPIRVOptions.nestInGPUModule = true;
3149
passManager.addPass(createConvertToSPIRVPass(convertToSPIRVOptions));
32-
OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
33-
modulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());
34-
modulePM.addPass(spirv::createSPIRVUpdateVCEPass());
50+
51+
OpPassManager &spirvModulePM =
52+
passManager.nest<gpu::GPUModuleOp>().nest<spirv::ModuleOp>();
53+
spirvModulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());
54+
spirvModulePM.addPass(spirv::createSPIRVUpdateVCEPass());
55+
if (options.spirvWebGPUPrepare)
56+
spirvModulePM.addPass(spirv::createSPIRVWebGPUPreparePass());
57+
58+
passManager.addPass(createGpuModuleToBinaryPass());
3559
}
3660

3761
} // namespace
3862

3963
namespace mlir::test {
4064
void registerTestVulkanRunnerPipeline() {
41-
PassPipelineRegistration<>(
65+
PassPipelineRegistration<VulkanRunnerPipelineOptions>(
4266
"test-vulkan-runner-pipeline",
4367
"Runs a series of passes for lowering GPU-dialect MLIR to "
4468
"SPIR-V-dialect MLIR intended for mlir-vulkan-runner.",

‎mlir/test/mlir-vulkan-runner/addui_extended.mlir‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
77
// RUN: --entry-point-result=void | FileCheck %s
88

9-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline -spirv-webgpu-prepare \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=spirv-webgpu-prepare \
1010
// RUN: | mlir-vulkan-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s

‎mlir/test/mlir-vulkan-runner/smul_extended.mlir‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
77
// RUN: --entry-point-result=void | FileCheck %s
88

9-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline -spirv-webgpu-prepare \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=spirv-webgpu-prepare \
1010
// RUN: | mlir-vulkan-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s

‎mlir/test/mlir-vulkan-runner/umul_extended.mlir‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
77
// RUN: --entry-point-result=void | FileCheck %s
88

9-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline -spirv-webgpu-prepare \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=spirv-webgpu-prepare \
1010
// RUN: | mlir-vulkan-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s

0 commit comments

Comments
 (0)