Skip to content

Commit b1b7643

Browse files
authored
[mlir][spirv] Add integration test for vector.deinterleave (#95469)
This commit is dependent on #95313.
1 parent 8e0ba08 commit b1b7643

File tree

1 file changed

+72
-0
lines changed

1 file changed

+72
-0
lines changed
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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: [0, 2]
6+
// CHECK: [1, 3]
7+
module attributes {
8+
gpu.container_module,
9+
spirv.target_env = #spirv.target_env<
10+
#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
11+
} {
12+
gpu.module @kernels {
13+
gpu.func @kernel_vector_deinterleave(%arg0 : memref<4xi32>, %arg1 : memref<2xi32>, %arg2 : memref<2xi32>)
14+
kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
15+
16+
%src = arith.constant dense<[0, 0, 0, 0]> : vector<4xi32>
17+
18+
%val0 = memref.load %arg0[0] : memref<4xi32>
19+
%val1 = memref.load %arg0[1] : memref<4xi32>
20+
%val2 = memref.load %arg0[2] : memref<4xi32>
21+
%val3 = memref.load %arg0[3] : memref<4xi32>
22+
23+
%src0 = vector.insert %val0, %src[0] : i32 into vector<4xi32>
24+
%src1 = vector.insert %val1, %src0[1] : i32 into vector<4xi32>
25+
%src2 = vector.insert %val2, %src1[2] : i32 into vector<4xi32>
26+
%src3 = vector.insert %val3, %src2[3] : i32 into vector<4xi32>
27+
28+
%res0, %res1 = vector.deinterleave %src3 : vector<4xi32> -> vector<2xi32>
29+
30+
%res0_0 = vector.extract %res0[0] : i32 from vector<2xi32>
31+
%res0_1 = vector.extract %res0[1] : i32 from vector<2xi32>
32+
%res1_0 = vector.extract %res1[0] : i32 from vector<2xi32>
33+
%res1_1 = vector.extract %res1[1] : i32 from vector<2xi32>
34+
35+
memref.store %res0_0, %arg1[0]: memref<2xi32>
36+
memref.store %res0_1, %arg1[1]: memref<2xi32>
37+
memref.store %res1_0, %arg2[0]: memref<2xi32>
38+
memref.store %res1_1, %arg2[1]: memref<2xi32>
39+
40+
gpu.return
41+
}
42+
}
43+
44+
func.func @main() {
45+
// Allocate 3 buffers.
46+
%buf0 = memref.alloc() : memref<4xi32>
47+
%buf1 = memref.alloc() : memref<2xi32>
48+
%buf2 = memref.alloc() : memref<2xi32>
49+
50+
// Initialize input buffer.
51+
%buf0_vals = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
52+
vector.store %buf0_vals, %buf0[0] : memref<4xi32>, vector<4xi32>
53+
54+
// Initialize output buffers.
55+
%value0 = arith.constant 0 : i32
56+
%buf3 = memref.cast %buf1 : memref<2xi32> to memref<?xi32>
57+
%buf4 = memref.cast %buf2 : memref<2xi32> to memref<?xi32>
58+
call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()
59+
call @fillResource1DInt(%buf4, %value0) : (memref<?xi32>, i32) -> ()
60+
61+
gpu.launch_func @kernels::@kernel_vector_deinterleave
62+
blocks in (4, 1, 1) threads in (1, 1, 1)
63+
args(%buf0 : memref<4xi32>, %buf1 : memref<2xi32>, %buf2 : memref<2xi32>)
64+
%buf5 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
65+
%buf6 = memref.cast %buf4 : memref<?xi32> to memref<*xi32>
66+
call @printMemrefI32(%buf5) : (memref<*xi32>) -> ()
67+
call @printMemrefI32(%buf6) : (memref<*xi32>) -> ()
68+
return
69+
}
70+
func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
71+
func.func private @printMemrefI32(%ptr : memref<*xi32>)
72+
}

0 commit comments

Comments
 (0)