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