|
| 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