[mlir][spirv] Update argmax-kernel lit test for convert-to-spirv. NFC.
#106586
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir Author: Angel Zhang (angelz913) ChangesThis PR updates the Full diff: https://github.com/llvm/llvm-project/pull/106586.diff 1 Files Affected:
diff --git a/mlir/test/Conversion/ConvertToSPIRV/argmax-kernel.mlir b/mlir/test/Conversion/ConvertToSPIRV/argmax-kernel.mlir
index 5cd1fead2527b1..8b1c3e54cbfb83 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/argmax-kernel.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/argmax-kernel.mlir
@@ -5,43 +5,50 @@ module attributes {
spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader, Groups, GroupNonUniformArithmetic, GroupNonUniformBallot], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
} {
// CHECK-LABEL: spirv.module @{{.*}} Logical GLSL450
- // CHECK-DAG: spirv.GlobalVariable @[[$LOCALINVOCATIONIDVAR:.*]] built_in("LocalInvocationId") : !spirv.ptr<vector<3xi32>, Input>
+ // CHECK-DAG: spirv.GlobalVariable @[[$LOCALINVOCATIONIDVAR:.*]] built_in("LocalInvocationId") : !spirv.ptr<vector<3xi32>, Input>
+ // CHECK-DAG: spirv.GlobalVariable @[[$SUBGROUPSIZE:.*]] built_in("SubgroupSize") : !spirv.ptr<i32, Input>
// CHECK-LABEL: spirv.func @argmax
- // CHECK-SAME: %[[ARG0:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>
- // CHECK-SAME: %[[ARG1:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>
+ // CHECK-SAME: %[[ARG0:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<128 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
+ // CHECK-SAME: %[[ARG1:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}
+ // CHECK-SAME: %[[ARG2:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 2)>}
gpu.module @kernels {
- gpu.func @argmax(%input : memref<4xf32>, %output : memref<i32>) kernel
+ gpu.func @argmax(%input : memref<128xf32>, %output : memref<1xi32>, %total_count_buf : memref<1xi32>) kernel
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>} {
// CHECK: %[[C0:.*]] = spirv.Constant 0 : i32
// CHECK: %[[C1:.*]] = spirv.Constant 1 : i32
- // CHECK: %[[C32:.*]] = spirv.Constant 32 : i32
+ // CHECK: %[[ADDRESSSUBGROUPSIZE:.*]] = spirv.mlir.addressof @[[$SUBGROUPSIZE]]
+ // CHECK: %[[SUBGROUPSIZE:.*]] = spirv.Load "Input" %[[ADDRESSSUBGROUPSIZE]]
// CHECK: %[[ADDRESSLOCALINVOCATIONID:.*]] = spirv.mlir.addressof @[[$LOCALINVOCATIONIDVAR]]
// CHECK: %[[LOCALINVOCATIONID:.*]] = spirv.Load "Input" %[[ADDRESSLOCALINVOCATIONID]]
// CHECK: %[[LOCALINVOCATIONIDX:.*]] = spirv.CompositeExtract %[[LOCALINVOCATIONID]]{{\[}}0 : i32{{\]}}
- // CHECK: %[[AC0:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[LOCALINVOCATIONIDX]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+ // CHECK: %[[AC:.*]] = spirv.AccessChain %[[ARG2]][%[[C0]], %[[C0]]] : !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>, i32, i32
+ // CHECK: %[[LOAD:.*]] = spirv.Load "StorageBuffer" %[[AC]] : i32
+ // CHECK: %[[AC0:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[LOCALINVOCATIONIDX]]] : !spirv.ptr<!spirv.struct<(!spirv.array<128 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[LOAD0:.*]] = spirv.Load "StorageBuffer" %[[AC0]] : f32
+ // CHECK: %[[UB:.*]] = spirv.UDiv %[[LOAD]], %[[SUBGROUPSIZE]] : i32
// CHECK: %[[FUNC0:.*]] = spirv.Variable : !spirv.ptr<i32, Function>
// CHECK: %[[FUNC1:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
- %cst_0_idx = arith.constant 0 : index
- %cst_1_i32 = arith.constant 1 : i32
- %cst_1_idx = arith.constant 1 : index
- %cst_32 = arith.constant 32 : i32
- %num_batches = arith.divui %cst_1_i32, %cst_32 : i32
- %tx = gpu.thread_id x
- %tx_i32 = index.castu %tx : index to i32
- %ub = index.castu %num_batches : i32 to index
+ %idx_0 = arith.constant 0 : index
+ %idx_1 = arith.constant 1 : index
+ %lane_count_idx = gpu.subgroup_size : index
+ %lane_count_i32 = index.castu %lane_count_idx : index to i32
+ %lane_id_idx = gpu.thread_id x
+ %lane_id_i32 = index.castu %lane_id_idx : index to i32
+ %total_count = memref.load %total_count_buf[%idx_0] : memref<1xi32>
%lane_res_init = arith.constant 0 : i32
- %lane_max_init = memref.load %input[%tx] : memref<4xf32>
+ %lane_max_init = memref.load %input[%lane_id_idx] : memref<128xf32>
+ %num_batches_i32 = arith.divui %total_count, %lane_count_i32 : i32
+ %num_batches_idx = index.castu %num_batches_i32 : i32 to index
// CHECK: spirv.mlir.loop {
// CHECK: spirv.Branch ^[[HEADER:.*]](%[[C1]], %[[C0]], %[[LOAD0]] : i32, i32, f32)
// CHECK: ^[[HEADER]](%[[INDVAR0:.*]]: i32, %[[INDVAR1:.*]]: i32, %[[INDVAR2:.*]]: f32):
- // CHECK: %[[SLT:.*]] = spirv.SLessThan %[[INDVAR0]], %[[C0]] : i32
+ // CHECK: %[[SLT:.*]] = spirv.SLessThan %[[INDVAR0]], %[[UB]] : i32
// CHECK: spirv.BranchConditional %[[SLT]], ^[[BODY:.*]], ^[[MERGE:.*]]
// CHECK: ^[[BODY]]:
- // CHECK: %[[MUL:.*]] = spirv.IMul %[[INDVAR0]], %[[C32]] : i32
+ // CHECK: %[[MUL:.*]] = spirv.IMul %[[SUBGROUPSIZE]], %[[INDVAR0]] : i32
// CHECK: %[[ADD:.*]] = spirv.IAdd %[[MUL]], %[[LOCALINVOCATIONIDX]] : i32
- // CHECK: %[[AC1:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[ADD]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+ // CHECK: %[[AC1:.*]] = spirv.AccessChain %[[ARG0]][%[[C0]], %[[ADD]]] : !spirv.ptr<!spirv.struct<(!spirv.array<128 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[LOAD1:.*]] = spirv.Load "StorageBuffer" %[[AC1]] : f32
// CHECK: %[[OGT:.*]] = spirv.FOrdGreaterThan %[[LOAD1]], %[[INDVAR2]] : f32
// CHECK: %[[SELECT0:.*]] = spirv.Select %[[OGT]], %[[ADD]], %[[INDVAR1]] : i1, i32
@@ -55,13 +62,13 @@ module attributes {
// CHECK: }
// CHECK-DAG: %[[LANE_RES:.*]] = spirv.Load "Function" %[[FUNC0]] : i32
// CHECK-DAG: %[[LANE_MAX:.*]] = spirv.Load "Function" %[[FUNC1]] : f32
- %lane_res, %lane_max = scf.for %iter = %cst_1_idx to %ub step %cst_1_idx
+ %lane_res, %lane_max = scf.for %iter = %idx_1 to %num_batches_idx step %idx_1
iter_args(%lane_res_iter = %lane_res_init, %lane_max_iter = %lane_max_init) -> (i32, f32) {
%iter_i32 = index.castu %iter : index to i32
- %mul = arith.muli %cst_32, %iter_i32 : i32
- %idx_i32 = arith.addi %mul, %tx_i32 : i32
+ %mul = arith.muli %lane_count_i32, %iter_i32 : i32
+ %idx_i32 = arith.addi %mul, %lane_id_i32 : i32
%idx = index.castu %idx_i32 : i32 to index
- %elem = memref.load %input[%idx] : memref<4xf32>
+ %elem = memref.load %input[%idx] : memref<128xf32>
%gt = arith.cmpf ogt, %elem, %lane_max_iter : f32
%lane_res_next = arith.select %gt, %idx_i32, %lane_res_iter : i32
%lane_max_next = arith.select %gt, %elem, %lane_max_iter : f32
@@ -72,12 +79,12 @@ module attributes {
// CHECK: %[[OEQ:.*]] = spirv.FOrdEqual %[[LANE_MAX]], %[[SUBGROUP_MAX]] : f32
// CHECK: %[[BALLOT:.*]] = spirv.GroupNonUniformBallot <Subgroup> %[[OEQ]] : vector<4xi32>
// CHECK: %[[BALLOTLSB:.*]] = spirv.GroupNonUniformBallotFindLSB <Subgroup> %[[BALLOT]] : vector<4xi32>, i32
- // CHECK: %[[EQ:.*]] = spirv.IEqual %[[LOCALINVOCATIONIDX]], %[[C1]] : i32
+ // CHECK: %[[EQ:.*]] = spirv.IEqual %[[BALLOTLSB]], %[[LOCALINVOCATIONIDX]] : i32
%subgroup_max = gpu.subgroup_reduce maximumf %lane_max : (f32) -> (f32)
%eq = arith.cmpf oeq, %lane_max, %subgroup_max : f32
%ballot = spirv.GroupNonUniformBallot <Subgroup> %eq : vector<4xi32>
%lsb = spirv.GroupNonUniformBallotFindLSB <Subgroup> %ballot : vector<4xi32>, i32
- %cond = arith.cmpi eq, %cst_1_i32, %tx_i32 : i32
+ %cond = arith.cmpi eq, %lsb, %lane_id_i32 : i32
// CHECK: spirv.mlir.selection {
// CHECK: spirv.BranchConditional %[[EQ]], ^[[TRUE:.*]], ^[[FALSE:.*]]
@@ -89,7 +96,7 @@ module attributes {
// CHECK: spirv.mlir.merge
// CHECK: }
scf.if %cond {
- memref.store %lane_res, %output[] : memref<i32>
+ memref.store %lane_res, %output[%idx_0] : memref<1xi32>
}
// CHECK: spirv.Return
|
angelz913
changed the title
convert-to-spirvargmax-kernel lit test for convert-to-spirv
kuhar
changed the title
argmax-kernel lit test for convert-to-spirvargmax-kernel lit test for convert-to-spirv. NFC.
|
This seems pretty large as a unit-test for conversion, can you elaborate on this test? |
This is effectively an integration tests for the convert-to-spirv pass. Unit tests are placed in each dialect conversion pass. Here want to make sure that a small complete kernel can be converted like we'd expect, including the interaction between dialects, control flow, type legalization, memory operations, etc. Historically, we saw failures in some combination of these like indices not being properly legalized. We picked argmax because it's a very simple example of code that does something that's not completely contrived: it scans the input buffer, loops over the content, and uses a spirv op as an 'intrinsic' for subgroup operations. We do not plan to add more of similar tests unless we can demonstrate that they fill some actual gaps in the coverage. We had a short-lived e2e tests for argmax, but that was a mistake -- this is better suited for a LIT test IMO. |
|
This is a very weird way to test though: seems completely arbitrary to me and not the kind of unit-tests I expect to see at all. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR updates the
argmax-kernellit test underConvertToSPIRV(for demonstrating the ability of using SPIR-V ops among other higher-level dialects). The original test was landed in #105010 but had a few mistakes, including using constants for variables like subgroup size and total number of elements in the input, etc. The PR fixed them so that this kernel module could potentially be used for integration tests. However, adding more integration tests is not planned at the moment.