mlir/test/Conversion/GPUToSPIRV/load-store.mlir

   1 // RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s
   2
   3 module attributes {
   4   gpu.container_module,
   5   spv.target_env = #spv.target_env<
   6     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   7 } {
   8   func @load_store(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>) {
   9     %c0 = arith.constant 0 : index
  10     %c12 = arith.constant 12 : index
  11     %0 = arith.subi %c12, %c0 : index
  12     %c1 = arith.constant 1 : index
  13     %c0_0 = arith.constant 0 : index
  14     %c4 = arith.constant 4 : index
  15     %1 = arith.subi %c4, %c0_0 : index
  16     %c1_1 = arith.constant 1 : index
  17     %c1_2 = arith.constant 1 : index
  18     gpu.launch_func @kernels::@load_store_kernel
  19         blocks in (%0, %c1_2, %c1_2) threads in (%1, %c1_2, %c1_2)
  20         args(%arg0 : memref<12x4xf32>, %arg1 : memref<12x4xf32>, %arg2 : memref<12x4xf32>,
  21              %c0 : index, %c0_0 : index, %c1 : index, %c1_1 : index)
  22     return
  23   }
  24
  25   // CHECK-LABEL: spv.module @{{.*}} Logical GLSL450
  26   gpu.module @kernels {
  27     // CHECK-DAG: spv.GlobalVariable @[[NUMWORKGROUPSVAR:.*]] built_in("NumWorkgroups") : !spv.ptr<vector<3xi32>, Input>
  28     // CHECK-DAG: spv.GlobalVariable @[[$LOCALINVOCATIONIDVAR:.*]] built_in("LocalInvocationId") : !spv.ptr<vector<3xi32>, Input>
  29     // CHECK-DAG: spv.GlobalVariable @[[$WORKGROUPIDVAR:.*]] built_in("WorkgroupId") : !spv.ptr<vector<3xi32>, Input>
  30     // CHECK-LABEL:    spv.func @load_store_kernel
  31     // CHECK-SAME: %[[ARG0:.*]]: !spv.ptr<!spv.struct<(!spv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spv.interface_var_abi = #spv.interface_var_abi<(0, 0)>}
  32     // CHECK-SAME: %[[ARG1:.*]]: !spv.ptr<!spv.struct<(!spv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spv.interface_var_abi = #spv.interface_var_abi<(0, 1)>}
  33     // CHECK-SAME: %[[ARG2:.*]]: !spv.ptr<!spv.struct<(!spv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spv.interface_var_abi = #spv.interface_var_abi<(0, 2)>}
  34     // CHECK-SAME: %[[ARG3:.*]]: i32 {spv.interface_var_abi = #spv.interface_var_abi<(0, 3), StorageBuffer>}
  35     // CHECK-SAME: %[[ARG4:.*]]: i32 {spv.interface_var_abi = #spv.interface_var_abi<(0, 4), StorageBuffer>}
  36     // CHECK-SAME: %[[ARG5:.*]]: i32 {spv.interface_var_abi = #spv.interface_var_abi<(0, 5), StorageBuffer>}
  37     // CHECK-SAME: %[[ARG6:.*]]: i32 {spv.interface_var_abi = #spv.interface_var_abi<(0, 6), StorageBuffer>}
  38     gpu.func @load_store_kernel(%arg0: memref<12x4xf32>, %arg1: memref<12x4xf32>, %arg2: memref<12x4xf32>, %arg3: index, %arg4: index, %arg5: index, %arg6: index) kernel
  39       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
  40       // CHECK: %[[ADDRESSWORKGROUPID:.*]] = spv.mlir.addressof @[[$WORKGROUPIDVAR]]
  41       // CHECK: %[[WORKGROUPID:.*]] = spv.Load "Input" %[[ADDRESSWORKGROUPID]]
  42       // CHECK: %[[WORKGROUPIDX:.*]] = spv.CompositeExtract %[[WORKGROUPID]]{{\[}}0 : i32{{\]}}
  43       // CHECK: %[[ADDRESSLOCALINVOCATIONID:.*]] = spv.mlir.addressof @[[$LOCALINVOCATIONIDVAR]]
  44       // CHECK: %[[LOCALINVOCATIONID:.*]] = spv.Load "Input" %[[ADDRESSLOCALINVOCATIONID]]
  45       // CHECK: %[[LOCALINVOCATIONIDX:.*]] = spv.CompositeExtract %[[LOCALINVOCATIONID]]{{\[}}0 : i32{{\]}}
  46       %0 = "gpu.block_id"() {dimension = "x"} : () -> index
  47       %1 = "gpu.block_id"() {dimension = "y"} : () -> index
  48       %2 = "gpu.block_id"() {dimension = "z"} : () -> index
  49       %3 = "gpu.thread_id"() {dimension = "x"} : () -> index
  50       %4 = "gpu.thread_id"() {dimension = "y"} : () -> index
  51       %5 = "gpu.thread_id"() {dimension = "z"} : () -> index
  52       %6 = "gpu.grid_dim"() {dimension = "x"} : () -> index
  53       %7 = "gpu.grid_dim"() {dimension = "y"} : () -> index
  54       %8 = "gpu.grid_dim"() {dimension = "z"} : () -> index
  55       %9 = "gpu.block_dim"() {dimension = "x"} : () -> index
  56       %10 = "gpu.block_dim"() {dimension = "y"} : () -> index
  57       %11 = "gpu.block_dim"() {dimension = "z"} : () -> index
  58       // CHECK: %[[INDEX1:.*]] = spv.IAdd %[[ARG3]], %[[WORKGROUPIDX]]
  59       %12 = arith.addi %arg3, %0 : index
  60       // CHECK: %[[INDEX2:.*]] = spv.IAdd %[[ARG4]], %[[LOCALINVOCATIONIDX]]
  61       %13 = arith.addi %arg4, %3 : index
  62       // CHECK: %[[ZERO:.*]] = spv.Constant 0 : i32
  63       // CHECK: %[[OFFSET1_0:.*]] = spv.Constant 0 : i32
  64       // CHECK: %[[STRIDE1_1:.*]] = spv.Constant 4 : i32
  65       // CHECK: %[[UPDATE1_1:.*]] = spv.IMul %[[STRIDE1_1]], %[[INDEX1]] : i32
  66       // CHECK: %[[OFFSET1_1:.*]] = spv.IAdd %[[OFFSET1_0]], %[[UPDATE1_1]] : i32
  67       // CHECK: %[[STRIDE1_2:.*]] = spv.Constant 1 : i32
  68       // CHECK: %[[UPDATE1_2:.*]] = spv.IMul %[[STRIDE1_2]], %[[INDEX2]] : i32
  69       // CHECK: %[[OFFSET1_2:.*]] = spv.IAdd %[[OFFSET1_1]], %[[UPDATE1_2]] : i32
  70       // CHECK: %[[PTR1:.*]] = spv.AccessChain %[[ARG0]]{{\[}}%[[ZERO]], %[[OFFSET1_2]]{{\]}}
  71       // CHECK-NEXT: %[[VAL1:.*]] = spv.Load "StorageBuffer" %[[PTR1]]
  72       %14 = memref.load %arg0[%12, %13] : memref<12x4xf32>
  73       // CHECK: %[[PTR2:.*]] = spv.AccessChain %[[ARG1]]{{\[}}{{%.*}}, {{%.*}}{{\]}}
  74       // CHECK-NEXT: %[[VAL2:.*]] = spv.Load "StorageBuffer" %[[PTR2]]
  75       %15 = memref.load %arg1[%12, %13] : memref<12x4xf32>
  76       // CHECK: %[[VAL3:.*]] = spv.FAdd %[[VAL1]], %[[VAL2]]
  77       %16 = arith.addf %14, %15 : f32
  78       // CHECK: %[[PTR3:.*]] = spv.AccessChain %[[ARG2]]{{\[}}{{%.*}}, {{%.*}}{{\]}}
  79       // CHECK-NEXT: spv.Store "StorageBuffer" %[[PTR3]], %[[VAL3]]
  80       memref.store %16, %arg2[%12, %13] : memref<12x4xf32>
  81       gpu.return
  82     }
  83   }
  84 }