Skip to content

Commit 11cea8b

Browse files
committed
Add integration test for vector.shuffle
1 parent 7b52bd7 commit 11cea8b

File tree

1 file changed

+53
-0
lines changed

1 file changed

+53
-0
lines changed
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
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: [2, 1, 3]
6+
module attributes {
7+
gpu.container_module,
8+
spirv.target_env = #spirv.target_env<
9+
#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
10+
} {
11+
gpu.module @kernels {
12+
gpu.func @kernel_vector_shuffle(%arg0 : memref<2xi32>, %arg1 : memref<2xi32>, %arg2 : memref<3xi32>)
13+
kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
14+
%c0 = arith.constant 0 : index
15+
%vec0 = vector.load %arg0[%c0] : memref<2xi32>, vector<2xi32>
16+
%vec1 = vector.load %arg1[%c0] : memref<2xi32>, vector<2xi32>
17+
%result = vector.shuffle %vec0, %vec1[2, 1, 3] : vector<2xi32>, vector<2xi32>
18+
vector.store %result, %arg2[%c0] : memref<3xi32>, vector<3xi32>
19+
gpu.return
20+
}
21+
}
22+
23+
func.func @main() {
24+
// Allocate 3 buffers.
25+
%buf0 = memref.alloc() : memref<2xi32>
26+
%buf1 = memref.alloc() : memref<2xi32>
27+
%buf2 = memref.alloc() : memref<3xi32>
28+
29+
%idx0 = arith.constant 0 : index
30+
%idx1 = arith.constant 1 : index
31+
%idx4 = arith.constant 4 : index
32+
33+
// Initialize input buffer
34+
%buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>
35+
%buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>
36+
vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>
37+
vector.store %buf1_vals, %buf1[%idx0] : memref<2xi32>, vector<2xi32>
38+
39+
// Initialize output buffer.
40+
%value0 = arith.constant 0 : i32
41+
%buf3 = memref.cast %buf2 : memref<3xi32> to memref<?xi32>
42+
call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()
43+
44+
gpu.launch_func @kernels::@kernel_vector_shuffle
45+
blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)
46+
args(%buf0 : memref<2xi32>, %buf1 : memref<2xi32>, %buf2 : memref<3xi32>)
47+
%buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
48+
call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()
49+
return
50+
}
51+
func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
52+
func.func private @printMemrefI32(%ptr : memref<*xi32>)
53+
}

0 commit comments

Comments
 (0)