1// RUN: mlir-opt -convert-gpu-to-spirv %s -o - | FileCheck %s 2 3module 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} 81