Skip to content

Commit 93eda08

Browse files
angelz913angelz913
authored
[mlir][spirv] Support gpu in convert-to-spirv pass (#105010)
This PR adds conversion patterns for GPU to the `convert-to-spirv` pass, introduced in #95942. Now the pass is able to convert each `gpu.module` and its ops within a `builtin.module` into a `spirv.module`. **Future Plans** - Use `gpu.launch_func` to invoke kernel from host functions - Potentially integrate into the `mlir-vulkan-runner` for e2e testing
1 parent e0441d5 commit 93eda08

File tree

5 files changed

+188
-0
lines changed

5 files changed

+188
-0
lines changed

‎mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ add_mlir_conversion_library(MLIRConvertToSPIRVPass
1515
MLIRArithToSPIRV
1616
MLIRArithTransforms
1717
MLIRFuncToSPIRV
18+
MLIRGPUToSPIRV
1819
MLIRIndexToSPIRV
1920
MLIRIR
2021
MLIRMemRefToSPIRV

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include "mlir/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.h"
1010
#include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
1111
#include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
12+
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
1213
#include "mlir/Conversion/IndexToSPIRV/IndexToSPIRV.h"
1314
#include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
1415
#include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
@@ -79,6 +80,7 @@ struct ConvertToSPIRVPass final
7980
arith::populateArithToSPIRVPatterns(typeConverter, patterns);
8081
populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
8182
populateFuncToSPIRVPatterns(typeConverter, patterns);
83+
populateGPUToSPIRVPatterns(typeConverter, patterns);
8284
index::populateIndexToSPIRVPatterns(typeConverter, patterns);
8385
populateMemRefToSPIRVPatterns(typeConverter, patterns);
8486
populateVectorToSPIRVPatterns(typeConverter, patterns);

‎mlir/test/Conversion/ConvertToSPIRV/argmax-kernel.mlir

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
// RUN: mlir-opt -convert-to-spirv -cse %s | FileCheck %s
2+
3+
module attributes {
4+
gpu.container_module,
5+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader, Groups, GroupNonUniformArithmetic, GroupNonUniformBallot], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
6+
} {
7+
// CHECK-LABEL: spirv.module @{{.*}} Logical GLSL450
8+
// CHECK-DAG: spirv.GlobalVariable @[[$LOCALINVOCATIONIDVAR:.*]] built_in("LocalInvocationId") : !spirv.ptr<vector<3xi32>, Input>
9+
// CHECK-LABEL: spirv.func @argmax
10+
// CHECK-SAME: %[[ARG0:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>
11+
// CHECK-SAME: %[[ARG1:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>
12+
gpu.module @kernels {
13+
gpu.func @argmax(%input : memref<4xf32>, %output : memref<i32>) kernel
14+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>} {
15+
// CHECK: %[[C0:.*]] = spirv.Constant 0 : i32
16+
// CHECK: %[[C1:.*]] = spirv.Constant 1 : i32
17+
// CHECK: %[[C32:.*]] = spirv.Constant 32 : i32
18+
// CHECK: %[[ADDRESSLOCALINVOCATIONID:.*]] = spirv.mlir.addressof @[[$LOCALINVOCATIONIDVAR]]
19+
// CHECK: %[[LOCALINVOCATIONID:.*]] = spirv.Load "Input" %[[ADDRESSLOCALINVOCATIONID]]
20+
// CHECK: %[[LOCALINVOCATIONIDX:.*]] = spirv.CompositeExtract %[[LOCALINVOCATIONID]]{{\[}}0 : i32{{\]}}
21+
// CHECK: %[[AC0:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[LOCALINVOCATIONIDX]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
22+
// CHECK: %[[LOAD0:.*]] = spirv.Load "StorageBuffer" %[[AC0]] : f32
23+
// CHECK: %[[FUNC0:.*]] = spirv.Variable : !spirv.ptr<i32, Function>
24+
// CHECK: %[[FUNC1:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
25+
%cst_0_idx = arith.constant 0 : index
26+
%cst_1_i32 = arith.constant 1 : i32
27+
%cst_1_idx = arith.constant 1 : index
28+
%cst_32 = arith.constant 32 : i32
29+
%num_batches = arith.divui %cst_1_i32, %cst_32 : i32
30+
%tx = gpu.thread_id x
31+
%tx_i32 = index.castu %tx : index to i32
32+
%ub = index.castu %num_batches : i32 to index
33+
%lane_res_init = arith.constant 0 : i32
34+
%lane_max_init = memref.load %input[%tx] : memref<4xf32>
35+
36+
// CHECK: spirv.mlir.loop {
37+
// CHECK: spirv.Branch ^[[HEADER:.*]](%[[C1]], %[[C0]], %[[LOAD0]] : i32, i32, f32)
38+
// CHECK: ^[[HEADER]](%[[INDVAR0:.*]]: i32, %[[INDVAR1:.*]]: i32, %[[INDVAR2:.*]]: f32):
39+
// CHECK: %[[SLT:.*]] = spirv.SLessThan %[[INDVAR0]], %[[C0]] : i32
40+
// CHECK: spirv.BranchConditional %[[SLT]], ^[[BODY:.*]], ^[[MERGE:.*]]
41+
// CHECK: ^[[BODY]]:
42+
// CHECK: %[[MUL:.*]] = spirv.IMul %[[INDVAR0]], %[[C32]] : i32
43+
// CHECK: %[[ADD:.*]] = spirv.IAdd %[[MUL]], %[[LOCALINVOCATIONIDX]] : i32
44+
// CHECK: %[[AC1:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[ADD]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
45+
// CHECK: %[[LOAD1:.*]] = spirv.Load "StorageBuffer" %[[AC1]] : f32
46+
// CHECK: %[[OGT:.*]] = spirv.FOrdGreaterThan %[[LOAD1]], %[[INDVAR2]] : f32
47+
// CHECK: %[[SELECT0:.*]] = spirv.Select %[[OGT]], %[[ADD]], %[[INDVAR1]] : i1, i32
48+
// CHECK: %[[SELECT1:.*]] = spirv.Select %[[OGT]], %[[LOAD1]], %[[INDVAR2]] : i1, f32
49+
// CHECK: spirv.Store "Function" %[[FUNC0]], %[[SELECT0]] : i32
50+
// CHECK: spirv.Store "Function" %[[FUNC1]], %[[SELECT1]] : f32
51+
// CHECK: %[[ADD1:.*]] = spirv.IAdd %[[INDVAR0]], %[[C1]] : i32
52+
// CHECK: spirv.Branch ^[[HEADER]](%[[ADD1]], %[[SELECT0]], %[[SELECT1]] : i32, i32, f32)
53+
// CHECK: ^[[MERGE]]:
54+
// CHECK: spirv.mlir.merge
55+
// CHECK: }
56+
// CHECK-DAG: %[[LANE_RES:.*]] = spirv.Load "Function" %[[FUNC0]] : i32
57+
// CHECK-DAG: %[[LANE_MAX:.*]] = spirv.Load "Function" %[[FUNC1]] : f32
58+
%lane_res, %lane_max = scf.for %iter = %cst_1_idx to %ub step %cst_1_idx
59+
iter_args(%lane_res_iter = %lane_res_init, %lane_max_iter = %lane_max_init) -> (i32, f32) {
60+
%iter_i32 = index.castu %iter : index to i32
61+
%mul = arith.muli %cst_32, %iter_i32 : i32
62+
%idx_i32 = arith.addi %mul, %tx_i32 : i32
63+
%idx = index.castu %idx_i32 : i32 to index
64+
%elem = memref.load %input[%idx] : memref<4xf32>
65+
%gt = arith.cmpf ogt, %elem, %lane_max_iter : f32
66+
%lane_res_next = arith.select %gt, %idx_i32, %lane_res_iter : i32
67+
%lane_max_next = arith.select %gt, %elem, %lane_max_iter : f32
68+
scf.yield %lane_res_next, %lane_max_next : i32, f32
69+
}
70+
71+
// CHECK: %[[SUBGROUP_MAX:.*]] = spirv.GroupNonUniformFMax "Subgroup" "Reduce" %[[LANE_MAX]] : f32
72+
// CHECK: %[[OEQ:.*]] = spirv.FOrdEqual %[[LANE_MAX]], %[[SUBGROUP_MAX]] : f32
73+
// CHECK: %[[BALLOT:.*]] = spirv.GroupNonUniformBallot <Subgroup> %[[OEQ]] : vector<4xi32>
74+
// CHECK: %[[BALLOTLSB:.*]] = spirv.GroupNonUniformBallotFindLSB <Subgroup> %[[BALLOT]] : vector<4xi32>, i32
75+
// CHECK: %[[EQ:.*]] = spirv.IEqual %[[LOCALINVOCATIONIDX]], %[[C1]] : i32
76+
%subgroup_max = gpu.subgroup_reduce maximumf %lane_max : (f32) -> (f32)
77+
%eq = arith.cmpf oeq, %lane_max, %subgroup_max : f32
78+
%ballot = spirv.GroupNonUniformBallot <Subgroup> %eq : vector<4xi32>
79+
%lsb = spirv.GroupNonUniformBallotFindLSB <Subgroup> %ballot : vector<4xi32>, i32
80+
%cond = arith.cmpi eq, %cst_1_i32, %tx_i32 : i32
81+
82+
// CHECK: spirv.mlir.selection {
83+
// CHECK: spirv.BranchConditional %[[EQ]], ^[[TRUE:.*]], ^[[FALSE:.*]]
84+
// CHECK: ^[[TRUE]]:
85+
// CHECK: %[[AC2:.*]] = spirv.AccessChain %[[ARG1]][%[[C0]], %[[C0]]] : !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>, i32, i32
86+
// CHECK: spirv.Store "StorageBuffer" %[[AC2]], %[[LANE_RES]] : i32
87+
// CHECK: spirv.Branch ^[[FALSE]]
88+
// CHECK: ^[[FALSE]]:
89+
// CHECK: spirv.mlir.merge
90+
// CHECK: }
91+
scf.if %cond {
92+
memref.store %lane_res, %output[] : memref<i32>
93+
}
94+
95+
// CHECK: spirv.Return
96+
gpu.return
97+
}
98+
}
99+
}

‎mlir/test/Conversion/ConvertToSPIRV/gpu.mlir

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// RUN: mlir-opt -convert-to-spirv -split-input-file %s | FileCheck %s
2+
3+
module attributes {
4+
gpu.container_module,
5+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
6+
} {
7+
8+
gpu.module @kernels {
9+
// CHECK-LABEL: spirv.func @all_reduce
10+
// CHECK-SAME: (%[[ARG0:.*]]: f32)
11+
// CHECK: %{{.*}} = spirv.GroupNonUniformFAdd "Workgroup" "Reduce" %[[ARG0]] : f32
12+
gpu.func @all_reduce(%arg0 : f32) kernel
13+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
14+
%reduced = gpu.all_reduce add %arg0 {} : (f32) -> (f32)
15+
gpu.return
16+
}
17+
}
18+
19+
}
20+
21+
// -----
22+
23+
module attributes {
24+
gpu.container_module,
25+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
26+
} {
27+
28+
gpu.module @kernels {
29+
// CHECK-LABEL: spirv.func @subgroup_reduce
30+
// CHECK-SAME: (%[[ARG0:.*]]: f32)
31+
// CHECK: %{{.*}} = spirv.GroupNonUniformFAdd "Subgroup" "Reduce" %[[ARG0]] : f32
32+
gpu.func @subgroup_reduce(%arg0 : f32) kernel
33+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
34+
%reduced = gpu.subgroup_reduce add %arg0 {} : (f32) -> (f32)
35+
gpu.return
36+
}
37+
}
38+
39+
}
40+
41+
// -----
42+
43+
module attributes {
44+
gpu.container_module,
45+
spirv.target_env = #spirv.target_env<
46+
#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
47+
} {
48+
49+
// CHECK-LABEL: spirv.module @{{.*}} Logical GLSL450
50+
// CHECK-LABEL: spirv.func @load_store
51+
// CHECK-SAME: %[[ARG0:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
52+
// CHECK-SAME: %[[ARG1:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}
53+
// CHECK-SAME: %[[ARG2:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 2)>}
54+
gpu.module @kernels {
55+
gpu.func @load_store(%arg0: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg1: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg2: memref<12x4xf32, #spirv.storage_class<StorageBuffer>>, %arg3: index, %arg4: index, %arg5: index, %arg6: index) kernel
56+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
57+
// CHECK: %[[PTR1:.*]] = spirv.AccessChain %[[ARG0]]
58+
// CHECK-NEXT: spirv.Load "StorageBuffer" %[[PTR1]]
59+
// CHECK: %[[PTR2:.*]] = spirv.AccessChain %[[ARG1]]
60+
// CHECK-NEXT: spirv.Load "StorageBuffer" %[[PTR2]]
61+
// CHECK: spirv.FAdd
62+
// CHECK: %[[PTR3:.*]] = spirv.AccessChain %[[ARG2]]
63+
// CHECK-NEXT: spirv.Store "StorageBuffer" %[[PTR3]]
64+
%0 = gpu.block_id x
65+
%1 = gpu.block_id y
66+
%2 = gpu.block_id z
67+
%3 = gpu.thread_id x
68+
%4 = gpu.thread_id y
69+
%5 = gpu.thread_id z
70+
%6 = gpu.grid_dim x
71+
%7 = gpu.grid_dim y
72+
%8 = gpu.grid_dim z
73+
%9 = gpu.block_dim x
74+
%10 = gpu.block_dim y
75+
%11 = gpu.block_dim z
76+
%12 = arith.addi %arg3, %0 : index
77+
%13 = arith.addi %arg4, %3 : index
78+
%14 = memref.load %arg0[%12, %13] : memref<12x4xf32, #spirv.storage_class<StorageBuffer>>
79+
%15 = memref.load %arg1[%12, %13] : memref<12x4xf32, #spirv.storage_class<StorageBuffer>>
80+
%16 = arith.addf %14, %15 : f32
81+
memref.store %16, %arg2[%12, %13] : memref<12x4xf32, #spirv.storage_class<StorageBuffer>>
82+
gpu.return
83+
}
84+
}
85+
}

‎utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8401,6 +8401,7 @@ cc_library(
84018401
":ArithTransforms",
84028402
":ConversionPassIncGen",
84038403
":FuncToSPIRV",
8404+
":GPUToSPIRV",
84048405
":IR",
84058406
":IndexToSPIRV",
84068407
":MemRefToSPIRV",

0 commit comments

Comments
 (0)