81 lines · plain
1// RUN: mlir-opt %s -test-vulkan-runner-pipeline \2// RUN: | mlir-runner - \3// RUN: --shared-libs=%mlir_vulkan_runtime,%mlir_runner_utils \4// RUN: --entry-point-result=void | FileCheck %s5 6// CHECK: [0, 2, 1, 3]7module 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_interleave(%arg0 : memref<2xi32>, %arg1 : memref<2xi32>, %arg2 : memref<4xi32>)14 kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {15 %idx0 = arith.constant 0 : index16 %idx1 = arith.constant 1 : index17 %idx2 = arith.constant 2 : index18 %idx3 = arith.constant 3 : index19 %idx4 = arith.constant 4 : index20 21 %lhs = arith.constant dense<[0, 0]> : vector<2xi32>22 %rhs = arith.constant dense<[0, 0]> : vector<2xi32>23 24 %val0 = memref.load %arg0[%idx0] : memref<2xi32>25 %val1 = memref.load %arg0[%idx1] : memref<2xi32>26 %val2 = memref.load %arg1[%idx0] : memref<2xi32>27 %val3 = memref.load %arg1[%idx1] : memref<2xi32>28 29 %lhs0 = vector.insert %val0, %lhs[%idx0] : i32 into vector<2xi32>30 %lhs1 = vector.insert %val1, %lhs0[%idx1] : i32 into vector<2xi32>31 %rhs0 = vector.insert %val2, %rhs[%idx0] : i32 into vector<2xi32>32 %rhs1 = vector.insert %val3, %rhs0[%idx1] : i32 into vector<2xi32>33 34 %interleave = vector.interleave %lhs1, %rhs1 : vector<2xi32> -> vector<4xi32>35 36 %res0 = vector.extract %interleave[%idx0] : i32 from vector<4xi32>37 %res1 = vector.extract %interleave[%idx1] : i32 from vector<4xi32>38 %res2 = vector.extract %interleave[%idx2] : i32 from vector<4xi32>39 %res3 = vector.extract %interleave[%idx3] : i32 from vector<4xi32>40 41 memref.store %res0, %arg2[%idx0]: memref<4xi32>42 memref.store %res1, %arg2[%idx1]: memref<4xi32>43 memref.store %res2, %arg2[%idx2]: memref<4xi32>44 memref.store %res3, %arg2[%idx3]: memref<4xi32>45 46 gpu.return47 }48 }49 50 func.func @main() {51 // Allocate 3 buffers.52 %buf0 = memref.alloc() : memref<2xi32>53 %buf1 = memref.alloc() : memref<2xi32>54 %buf2 = memref.alloc() : memref<4xi32>55 56 %idx0 = arith.constant 0 : index57 %idx1 = arith.constant 1 : index58 %idx4 = arith.constant 4 : index59 60 // Initialize input buffer.61 %buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>62 %buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>63 vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>64 vector.store %buf1_vals, %buf1[%idx0] : memref<2xi32>, vector<2xi32>65 66 // Initialize output buffer.67 %value0 = arith.constant 0 : i3268 %buf3 = memref.cast %buf2 : memref<4xi32> to memref<?xi32>69 call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()70 71 gpu.launch_func @kernels::@kernel_vector_interleave72 blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)73 args(%buf0 : memref<2xi32>, %buf1 : memref<2xi32>, %buf2 : memref<4xi32>)74 %buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>75 call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()76 return77 }78 func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)79 func.func private @printMemrefI32(%ptr : memref<*xi32>)80}81