mlir/test/mlir-vulkan-runner/vector-shuffle.mlir

   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, 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<4xi32>)
  13       kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
  14       %idx0 = arith.constant 0 : index
  15       %idx1 = arith.constant 1 : index
  16       %idx2 = arith.constant 2 : index
  17       %idx3 = arith.constant 3 : index
  18       %idx4 = arith.constant 4 : index
  19
  20       %lhs = arith.constant dense<[0, 0]> : vector<2xi32>
  21       %rhs = arith.constant dense<[0, 0]> : vector<2xi32>
  22
  23       %val0 = memref.load %arg0[%idx0] : memref<2xi32>
  24       %val1 = memref.load %arg0[%idx1] : memref<2xi32>
  25       %val2 = memref.load %arg1[%idx0] : memref<2xi32>
  26       %val3 = memref.load %arg1[%idx1] : memref<2xi32>
  27
  28       %lhs0 = vector.insertelement %val0, %lhs[%idx0 : index] : vector<2xi32>
  29       %lhs1 = vector.insertelement %val1, %lhs0[%idx1 : index] : vector<2xi32>
  30       %rhs0 = vector.insertelement %val2, %rhs[%idx0 : index] : vector<2xi32>
  31       %rhs1 = vector.insertelement %val3, %rhs0[%idx1 : index] : vector<2xi32>
  32
  33       %shuffle = vector.shuffle %lhs1, %rhs1[2, 1, 3, 3] : vector<2xi32>, vector<2xi32>
  34
  35       %res0 = vector.extractelement %shuffle[%idx0 : index] : vector<4xi32>
  36       %res1 = vector.extractelement %shuffle[%idx1 : index] : vector<4xi32>
  37       %res2 = vector.extractelement %shuffle[%idx2 : index] : vector<4xi32>
  38       %res3 = vector.extractelement %shuffle[%idx3 : index] : vector<4xi32>
  39
  40       memref.store %res0, %arg2[%idx0]: memref<4xi32>
  41       memref.store %res1, %arg2[%idx1]: memref<4xi32>
  42       memref.store %res2, %arg2[%idx2]: memref<4xi32>
  43       memref.store %res3, %arg2[%idx3]: memref<4xi32>
  44
  45       gpu.return
  46     }
  47   }
  48
  49   func.func @main() {
  50     // Allocate 3 buffers.
  51     %buf0 = memref.alloc() : memref<2xi32>
  52     %buf1 = memref.alloc() : memref<2xi32>
  53     %buf2 = memref.alloc() : memref<4xi32>
  54
  55     %idx0 = arith.constant 0 : index
  56     %idx1 = arith.constant 1 : index
  57     %idx4 = arith.constant 4 : index
  58
  59     // Initialize input buffer.
  60     %buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>
  61     %buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>
  62     vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>
  63     vector.store %buf1_vals, %buf1[%idx0] : memref<2xi32>, vector<2xi32>
  64
  65     // Initialize output buffer.
  66     %value0 = arith.constant 0 : i32
  67     %buf3 = memref.cast %buf2 : memref<4xi32> to memref<?xi32>
  68     call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()
  69
  70     gpu.launch_func @kernels::@kernel_vector_shuffle
  71         blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)
  72         args(%buf0 : memref<2xi32>, %buf1 : memref<2xi32>, %buf2 : memref<4xi32>)
  73     %buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
  74     call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()
  75     return
  76   }
  77   func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
  78   func.func private @printMemrefI32(%ptr : memref<*xi32>)
  79 }