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

   1 // RUN: mlir-opt %s -test-vulkan-runner-pipeline \
   2 // RUN:   | mlir-vulkan-runner - \
   3 // RUN:     --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
   4 // RUN:     --entry-point-result=void | FileCheck %s
   5
   6 // CHECK: [0, 2, 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_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 : index
  16       %idx1 = arith.constant 1 : index
  17       %idx2 = arith.constant 2 : index
  18       %idx3 = arith.constant 3 : index
  19       %idx4 = arith.constant 4 : index
  20
  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.insertelement %val0, %lhs[%idx0 : index] : vector<2xi32>
  30       %lhs1 = vector.insertelement %val1, %lhs0[%idx1 : index] : vector<2xi32>
  31       %rhs0 = vector.insertelement %val2, %rhs[%idx0 : index] : vector<2xi32>
  32       %rhs1 = vector.insertelement %val3, %rhs0[%idx1 : index] : vector<2xi32>
  33
  34       %interleave = vector.interleave %lhs1, %rhs1 : vector<2xi32> -> vector<4xi32>
  35
  36       %res0 = vector.extractelement %interleave[%idx0 : index] : vector<4xi32>
  37       %res1 = vector.extractelement %interleave[%idx1 : index] : vector<4xi32>
  38       %res2 = vector.extractelement %interleave[%idx2 : index] : vector<4xi32>
  39       %res3 = vector.extractelement %interleave[%idx3 : index] : 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.return
  47     }
  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 : index
  57     %idx1 = arith.constant 1 : index
  58     %idx4 = arith.constant 4 : index
  59
  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 : i32
  68     %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_interleave
  72         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     return
  77   }
  78   func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
  79   func.func private @printMemrefI32(%ptr : memref<*xi32>)
  80 }