Skip to content

Commit 94cbf47

Browse files
angelz913angelz913
committed
[mlir][spirv] Add an argmax integration test with mlir-vulkan-runner
1 parent d880f5a commit 94cbf47

File tree

3 files changed

+106
-1
lines changed

3 files changed

+106
-1
lines changed

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
// RUN: mlir-vulkan-runner %s \
2+
// RUN: --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
3+
// RUN: --entry-point-result=void | FileCheck %s
4+
5+
// CHECK: [15]
6+
module attributes {
7+
gpu.container_module,
8+
spirv.target_env = #spirv.target_env<
9+
#spirv.vce<v1.3, [Shader, Groups, GroupNonUniformArithmetic, GroupNonUniformBallot], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
10+
} {
11+
gpu.module @kernels {
12+
gpu.func @kernel_argmax(%input : memref<128xi32>, %output : memref<1xi32>, %total_count_buf : memref<1xi32>) kernel
13+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>} {
14+
%idx0 = arith.constant 0 : index
15+
%idx1 = arith.constant 1 : index
16+
17+
%total_count = memref.load %total_count_buf[%idx0] : memref<1xi32>
18+
%lane_count_idx = gpu.subgroup_size : index
19+
%lane_count_i32 = index.castu %lane_count_idx : index to i32
20+
%lane_id_idx = gpu.thread_id x
21+
%lane_id_i32 = index.castu %lane_id_idx : index to i32
22+
%lane_res_init = arith.constant 0 : i32
23+
%lane_max_init = memref.load %input[%lane_id_idx] : memref<128xi32>
24+
%num_batches_i32 = arith.divui %total_count, %lane_count_i32 : i32
25+
%num_batches_idx = index.castu %num_batches_i32 : i32 to index
26+
27+
%lane_res, %lane_max = scf.for %iter = %idx1 to %num_batches_idx step %idx1
28+
iter_args(%lane_res_iter = %lane_res_init, %lane_max_iter = %lane_max_init) -> (i32, i32) {
29+
%iter_i32 = index.castu %iter : index to i32
30+
%mul = arith.muli %lane_count_i32, %iter_i32 : i32
31+
%idx_i32 = arith.addi %mul, %lane_id_i32 : i32
32+
%idx = index.castu %idx_i32 : i32 to index
33+
%elem = memref.load %input[%idx] : memref<128xi32>
34+
%gt = arith.cmpi sgt, %elem, %lane_max_iter : i32
35+
%lane_res_next = arith.select %gt, %idx_i32, %lane_res_iter : i32
36+
%lane_max_next = arith.select %gt, %elem, %lane_max_iter : i32
37+
scf.yield %lane_res_next, %lane_max_next : i32, i32
38+
}
39+
40+
%subgroup_max = gpu.subgroup_reduce maxsi %lane_max : (i32) -> (i32)
41+
%eq = arith.cmpi eq, %lane_max, %subgroup_max : i32
42+
%ballot = spirv.GroupNonUniformBallot <Subgroup> %eq : vector<4xi32>
43+
%lsb = spirv.GroupNonUniformBallotFindLSB <Subgroup> %ballot : vector<4xi32>, i32
44+
%cond = arith.cmpi eq, %lsb, %lane_id_i32 : i32
45+
46+
scf.if %cond {
47+
memref.store %lane_res, %output[%idx0] : memref<1xi32>
48+
}
49+
50+
gpu.return
51+
}
52+
}
53+
54+
func.func @main() {
55+
// Allocate 3 buffers.
56+
%in_buf = memref.alloc() : memref<128xi32>
57+
%out_buf = memref.alloc() : memref<1xi32>
58+
%total_count_buf = memref.alloc() : memref<1xi32>
59+
60+
// Constants.
61+
%cst0 = arith.constant 0 : i32
62+
%idx0 = arith.constant 0 : index
63+
%idx1 = arith.constant 1 : index
64+
%idx16 = arith.constant 16 : index
65+
%idx32 = arith.constant 32 : index
66+
%idx48 = arith.constant 48 : index
67+
%idx64 = arith.constant 64 : index
68+
%idx80 = arith.constant 80 : index
69+
%idx96 = arith.constant 96 : index
70+
%idx112 = arith.constant 112 : index
71+
72+
// Initialize input buffer.
73+
%in_vec = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xi32>
74+
vector.store %in_vec, %in_buf[%idx0] : memref<128xi32>, vector<16xi32>
75+
vector.store %in_vec, %in_buf[%idx16] : memref<128xi32>, vector<16xi32>
76+
vector.store %in_vec, %in_buf[%idx32] : memref<128xi32>, vector<16xi32>
77+
vector.store %in_vec, %in_buf[%idx48] : memref<128xi32>, vector<16xi32>
78+
vector.store %in_vec, %in_buf[%idx64] : memref<128xi32>, vector<16xi32>
79+
vector.store %in_vec, %in_buf[%idx80] : memref<128xi32>, vector<16xi32>
80+
vector.store %in_vec, %in_buf[%idx96] : memref<128xi32>, vector<16xi32>
81+
vector.store %in_vec, %in_buf[%idx112] : memref<128xi32>, vector<16xi32>
82+
83+
// Initialize output buffer.
84+
%out_buf2 = memref.cast %out_buf : memref<1xi32> to memref<?xi32>
85+
call @fillResource1DInt(%out_buf2, %cst0) : (memref<?xi32>, i32) -> ()
86+
87+
// Total number of scalars.
88+
%total_count = arith.constant 128 : i32
89+
%total_count_buf2 = memref.cast %total_count_buf : memref<1xi32> to memref<?xi32>
90+
call @fillResource1DInt(%total_count_buf2, %total_count) : (memref<?xi32>, i32) -> ()
91+
92+
// Launch kernel function and print output.
93+
gpu.launch_func @kernels::@kernel_argmax
94+
blocks in (%idx1, %idx1, %idx1) threads in (%idx32, %idx1, %idx1)
95+
args(%in_buf : memref<128xi32>, %out_buf : memref<1xi32>, %total_count_buf : memref<1xi32>)
96+
%out_buf3 = memref.cast %out_buf2 : memref<?xi32> to memref<*xi32>
97+
call @printMemrefI32(%out_buf3) : (memref<*xi32>) -> ()
98+
return
99+
}
100+
func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
101+
func.func private @printMemrefI32(%ptr : memref<*xi32>)
102+
}

‎mlir/tools/mlir-vulkan-runner/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ if (MLIR_ENABLE_VULKAN_RUNNER)
5757
MLIRExecutionEngine
5858
MLIRFuncDialect
5959
MLIRGPUDialect
60+
MLIRIndexDialect
6061
MLIRIR
6162
MLIRJitRunner
6263
MLIRLLVMDialect

‎mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include "mlir/Dialect/Func/IR/FuncOps.h"
2525
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
2626
#include "mlir/Dialect/GPU/Transforms/Passes.h"
27+
#include "mlir/Dialect/Index/IR/IndexDialect.h"
2728
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
2829
#include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
2930
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -110,7 +111,8 @@ int main(int argc, char **argv) {
110111
registry.insert<mlir::arith::ArithDialect, mlir::LLVM::LLVMDialect,
111112
mlir::gpu::GPUDialect, mlir::spirv::SPIRVDialect,
112113
mlir::scf::SCFDialect, mlir::func::FuncDialect,
113-
mlir::memref::MemRefDialect, mlir::vector::VectorDialect>();
114+
mlir::memref::MemRefDialect, mlir::vector::VectorDialect,
115+
mlir::index::IndexDialect>();
114116
mlir::registerBuiltinDialectTranslation(registry);
115117
mlir::registerLLVMDialectTranslation(registry);
116118

0 commit comments

Comments
 (0)