Skip to content

Commit 99a562b

Browse files
andfau-amdandfau-amd
authored
[mlir][spirv] Add mgpu* wrappers for Vulkan runtime, migrate some tests (#123114)
This commit adds new wrappers around the MLIR Vulkan runtime which implement the mgpu* APIs (as generated by GPUToLLVMConversionPass), adds an optional LLVM lowering to the Vulkan runner mlir-opt pipeline based on GPUToLLVMConversionPass, and migrates several of the mlir-vulkan-runner tests to use mlir-cpu-runner instead, together with the new pipeline and wrappers. This is a further incremental step towards eliminating mlir-vulkan-runner and its associated pipeline, passes and wrappers (#73457). This commit does not migrate all of the tests to the new system, because changes to the mgpuLaunchKernel ABI will be necessary to support the tests that use multi-dimensional memref arguments.
1 parent 8c75ecb commit 99a562b

File tree

14 files changed

+208
-32
lines changed

14 files changed

+208
-32
lines changed

‎mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ struct FunctionCallBuilder {
6464
/// populate converter for gpu types.
6565
void populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
6666
RewritePatternSet &patterns,
67-
bool kernelBarePtrCallConv = false);
67+
bool kernelBarePtrCallConv = false,
68+
bool typeCheckKernelArgs = false);
6869

6970
/// A function that maps a MemorySpace enum to a target-specific integer value.
7071
using MemorySpaceMapping = std::function<unsigned(gpu::AddressSpace)>;

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -517,6 +517,12 @@ def GpuToLLVMConversionPass : Pass<"gpu-to-llvm", "ModuleOp"> {
517517
/*default=*/"false",
518518
"Use bare pointers to pass memref arguments to kernels. "
519519
"The kernel must use the same setting for this option."
520+
>,
521+
Option<"typeCheckKernelArgs", "type-check-kernel-args", "bool",
522+
/*default=*/"false",
523+
"Require all kernel arguments to be memrefs of rank 1 and with a "
524+
"32-bit element size. This is a temporary option that will be "
525+
"removed; TODO(https://github.com/llvm/llvm-project/issues/73457)."
520526
>
521527
];
522528

‎mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -427,16 +427,19 @@ class LegalizeLaunchFuncOpPattern
427427
: public ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp> {
428428
public:
429429
LegalizeLaunchFuncOpPattern(const LLVMTypeConverter &typeConverter,
430-
bool kernelBarePtrCallConv)
430+
bool kernelBarePtrCallConv,
431+
bool typeCheckKernelArgs)
431432
: ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(typeConverter),
432-
kernelBarePtrCallConv(kernelBarePtrCallConv) {}
433+
kernelBarePtrCallConv(kernelBarePtrCallConv),
434+
typeCheckKernelArgs(typeCheckKernelArgs) {}
433435

434436
private:
435437
LogicalResult
436438
matchAndRewrite(gpu::LaunchFuncOp launchOp, OpAdaptor adaptor,
437439
ConversionPatternRewriter &rewriter) const override;
438440

439441
bool kernelBarePtrCallConv;
442+
bool typeCheckKernelArgs;
440443
};
441444

442445
/// A rewrite pattern to convert gpu.memcpy operations into a GPU runtime
@@ -563,8 +566,8 @@ void GpuToLLVMConversionPass::runOnOperation() {
563566
populateFinalizeMemRefToLLVMConversionPatterns(converter, patterns);
564567
populateAsyncStructuralTypeConversionsAndLegality(converter, patterns,
565568
target);
566-
populateGpuToLLVMConversionPatterns(converter, patterns,
567-
kernelBarePtrCallConv);
569+
populateGpuToLLVMConversionPatterns(
570+
converter, patterns, kernelBarePtrCallConv, typeCheckKernelArgs);
568571

569572
if (failed(
570573
applyPartialConversion(getOperation(), target, std::move(patterns))))
@@ -966,6 +969,28 @@ LogicalResult LegalizeLaunchFuncOpPattern::matchAndRewrite(
966969
// stream must be created to pass to subsequent operations.
967970
else if (launchOp.getAsyncToken())
968971
stream = streamCreateCallBuilder.create(loc, rewriter, {}).getResult();
972+
973+
if (typeCheckKernelArgs) {
974+
// The current non-bare-pointer ABI is a bad fit for `mgpuLaunchKernel`,
975+
// which takes an untyped list of arguments. The type check here prevents
976+
// accidentally violating the assumption made in vulkan-runtime-wrappers.cpp
977+
// and creating a unchecked runtime ABI mismatch.
978+
// TODO(https://github.com/llvm/llvm-project/issues/73457): Change the ABI
979+
// here to remove the need for this type check.
980+
for (Value arg : launchOp.getKernelOperands()) {
981+
if (auto memrefTy = dyn_cast<MemRefType>(arg.getType())) {
982+
if (memrefTy.getRank() != 1 ||
983+
memrefTy.getElementTypeBitWidth() != 32) {
984+
return rewriter.notifyMatchFailure(
985+
launchOp, "Operand to launch op is not a rank-1 memref with "
986+
"32-bit element type.");
987+
}
988+
} else {
989+
return rewriter.notifyMatchFailure(
990+
launchOp, "Operand to launch op is not a memref.");
991+
}
992+
}
993+
}
969994
// Lower the kernel operands to match kernel parameters.
970995
// Note: If `useBarePtrCallConv` is set in the type converter's options,
971996
// the value of `kernelBarePtrCallConv` will be ignored.
@@ -1737,7 +1762,8 @@ LogicalResult ConvertCreateBsrOpToGpuRuntimeCallPattern::matchAndRewrite(
17371762

17381763
void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
17391764
RewritePatternSet &patterns,
1740-
bool kernelBarePtrCallConv) {
1765+
bool kernelBarePtrCallConv,
1766+
bool typeCheckKernelArgs) {
17411767
addOpaquePointerConversion<gpu::AsyncTokenType>(converter);
17421768
addOpaquePointerConversion<gpu::SparseDnTensorHandleType>(converter);
17431769
addOpaquePointerConversion<gpu::SparseSpMatHandleType>(converter);
@@ -1774,7 +1800,8 @@ void mlir::populateGpuToLLVMConversionPatterns(LLVMTypeConverter &converter,
17741800
ConvertSpGEMMCopyOpToGpuRuntimeCallPattern,
17751801
ConvertSpMatGetSizeOpToGpuRuntimeCallPattern,
17761802
ConvertSetCsrPointersOpToGpuRuntimeCallPattern>(converter);
1777-
patterns.add<LegalizeLaunchFuncOpPattern>(converter, kernelBarePtrCallConv);
1803+
patterns.add<LegalizeLaunchFuncOpPattern>(converter, kernelBarePtrCallConv,
1804+
typeCheckKernelArgs);
17781805
}
17791806

17801807
//===----------------------------------------------------------------------===//

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,13 @@
1111
//===----------------------------------------------------------------------===//
1212

1313
#include "mlir/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.h"
14+
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
1415
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
16+
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
17+
#include "mlir/Dialect/Func/IR/FuncOps.h"
1518
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1619
#include "mlir/Dialect/GPU/Transforms/Passes.h"
20+
#include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
1721
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
1822
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
1923
#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
@@ -29,6 +33,9 @@ struct VulkanRunnerPipelineOptions
2933
Option<bool> spirvWebGPUPrepare{
3034
*this, "spirv-webgpu-prepare",
3135
llvm::cl::desc("Run MLIR transforms used when targetting WebGPU")};
36+
Option<bool> toLlvm{*this, "to-llvm",
37+
llvm::cl::desc("Run MLIR transforms to lower host code "
38+
"to LLVM, intended for mlir-cpu-runner")};
3239
};
3340

3441
void buildTestVulkanRunnerPipeline(OpPassManager &passManager,
@@ -56,6 +63,19 @@ void buildTestVulkanRunnerPipeline(OpPassManager &passManager,
5663
spirvModulePM.addPass(spirv::createSPIRVWebGPUPreparePass());
5764

5865
passManager.addPass(createGpuModuleToBinaryPass());
66+
67+
if (options.toLlvm) {
68+
passManager.addPass(createFinalizeMemRefToLLVMConversionPass());
69+
passManager.nest<func::FuncOp>().addPass(
70+
LLVM::createRequestCWrappersPass());
71+
// vulkan-runtime-wrappers.cpp uses the non-bare-pointer calling convention,
72+
// and the type check is needed to prevent accidental ABI mismatches.
73+
GpuToLLVMConversionPassOptions opt;
74+
opt.hostBarePtrCallConv = false;
75+
opt.kernelBarePtrCallConv = false;
76+
opt.typeCheckKernelArgs = true;
77+
passManager.addPass(createGpuToLLVMConversionPass(opt));
78+
}
5979
}
6080

6181
} // namespace
@@ -65,7 +85,7 @@ void registerTestVulkanRunnerPipeline() {
6585
PassPipelineRegistration<VulkanRunnerPipelineOptions>(
6686
"test-vulkan-runner-pipeline",
6787
"Runs a series of passes for lowering GPU-dialect MLIR to "
68-
"SPIR-V-dialect MLIR intended for mlir-vulkan-runner.",
88+
"SPIR-V-dialect MLIR intended for mlir-vulkan-runner or mlir-cpu-runner.",
6989
buildTestVulkanRunnerPipeline);
7090
}
7191
} // namespace mlir::test

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
33

44
// CHECK: [3.3, 3.3, 3.3, 3.3, 3.3, 3.3, 3.3, 3.3]
55
module attributes {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
33

44
// CHECK: [3.3, 3.3, 3.3, 3.3, 0, 0, 0, 0]
55
module attributes {

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
// Make sure that addition with carry produces expected results
22
// with and without expansion to primitive add/cmp ops for WebGPU.
33

4-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
5-
// RUN: | mlir-vulkan-runner - \
4+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
5+
// RUN: | mlir-cpu-runner - \
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 \
10-
// RUN: | mlir-vulkan-runner - \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline="spirv-webgpu-prepare to-llvm" \
10+
// RUN: | mlir-cpu-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s
1313

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
// Make sure that signed extended multiplication produces expected results
22
// with and without expansion to primitive mul/add ops for WebGPU.
33

4-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
5-
// RUN: | mlir-vulkan-runner - \
4+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
5+
// RUN: | mlir-cpu-runner - \
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 \
10-
// RUN: | mlir-vulkan-runner - \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline="spirv-webgpu-prepare to-llvm" \
10+
// RUN: | mlir-cpu-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s
1313

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s
33

44
// CHECK: Compute shader execution time
55
// CHECK: Command buffer submit time

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
// Make sure that unsigned extended multiplication produces expected results
22
// with and without expansion to primitive mul/add ops for WebGPU.
33

4-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
5-
// RUN: | mlir-vulkan-runner - \
4+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
5+
// RUN: | mlir-cpu-runner - \
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 \
10-
// RUN: | mlir-vulkan-runner - \
9+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline="spirv-webgpu-prepare to-llvm" \
10+
// RUN: | mlir-cpu-runner - \
1111
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
1212
// RUN: --entry-point-result=void | FileCheck %s
1313

‎mlir/test/mlir-vulkan-runner/vector-deinterleave.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - \
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - \
33
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
44
// RUN: --entry-point-result=void | FileCheck %s
55

‎mlir/test/mlir-vulkan-runner/vector-interleave.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - \
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - \
33
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
44
// RUN: --entry-point-result=void | FileCheck %s
55

‎mlir/test/mlir-vulkan-runner/vector-shuffle.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
// RUN: mlir-opt %s -test-vulkan-runner-pipeline \
2-
// RUN: | mlir-vulkan-runner - \
1+
// RUN: mlir-opt %s -test-vulkan-runner-pipeline=to-llvm \
2+
// RUN: | mlir-cpu-runner - \
33
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
44
// RUN: --entry-point-result=void | FileCheck %s
55

0 commit comments

Comments
 (0)