205 lines · plain
1// RUN: mlir-opt --convert-gpu-to-spirv --cse \2// RUN: --split-input-file --verify-diagnostics %s | FileCheck %s3 4module attributes {5 gpu.container_module,6 spirv.target_env = #spirv.target_env<#spirv.vce<v1.6,7 [Shader, CooperativeMatrixKHR, Float16],8 [SPV_KHR_storage_buffer_storage_class, SPV_KHR_cooperative_matrix]>,9 #spirv.resource_limits<>>} {10 11 gpu.module @kernels {12 // CHECK-LABEL: spirv.func @gpu_wmma_load_op13 // CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer>14 gpu.func @gpu_wmma_load_op(%arg0 : memref<32x32xf16, #spirv.storage_class<StorageBuffer>>) kernel15 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {16 %i = arith.constant 16 : index17 %j = arith.constant 16 : index18 // CHECK: %[[STRIDE:.+]] = spirv.Constant 32 : i3219 // CHECK: spirv.KHR.CooperativeMatrixLoad {{%.*}}, %[[STRIDE]], <RowMajor> :20 // CHECK-SAME: !spirv.ptr<f32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>21 %0 = gpu.subgroup_mma_load_matrix %arg0[%i, %j] {leadDimension = 32 : index} :22 memref<32x32xf16, #spirv.storage_class<StorageBuffer>> -> !gpu.mma_matrix<16x16xf16, "COp">23 24 // CHECK: spirv.KHR.CooperativeMatrixLoad {{%.*}}, %[[STRIDE]], <ColumnMajor> :25 // CHECK-SAME: !spirv.ptr<f32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>26 %1 = gpu.subgroup_mma_load_matrix %arg0[%i, %j] {leadDimension = 32 : index, transpose} :27 memref<32x32xf16, #spirv.storage_class<StorageBuffer>> -> !gpu.mma_matrix<16x16xf16, "COp">28 // CHECK: spirv.Return29 gpu.return30 }31 32 // CHECK-LABEL: spirv.func @gpu_wmma_store_op33 // CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer>34 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>35 gpu.func @gpu_wmma_store_op(%arg0: memref<32x32xf16, #spirv.storage_class<StorageBuffer>>,36 %arg1: !gpu.mma_matrix<16x16xf16, "COp">) kernel37 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {38 %i = arith.constant 16 : index39 %j = arith.constant 16 : index40 // CHECK: %[[STRIDE:.+]] = spirv.Constant 32 : i3241 // CHECK: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, %[[STRIDE]], <RowMajor> :42 // CHECK-SAME: !spirv.ptr<f32, StorageBuffer>, !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>43 gpu.subgroup_mma_store_matrix %arg1, %arg0[%i,%j] {leadDimension = 32 : index} :44 !gpu.mma_matrix<16x16xf16, "COp">, memref<32x32xf16, #spirv.storage_class<StorageBuffer>>45 46 // CHECK: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, %[[STRIDE]], <ColumnMajor> :47 // CHECK-SAME: !spirv.ptr<f32, StorageBuffer>, !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>48 gpu.subgroup_mma_store_matrix %arg1, %arg0[%i,%j] {leadDimension = 32 : index, transpose} :49 !gpu.mma_matrix<16x16xf16, "COp">, memref<32x32xf16, #spirv.storage_class<StorageBuffer>>50 // CHECK: spirv.Return51 gpu.return52 }53 54 // CHECK-LABEL: spirv.func @gpu_wmma_mma_op55 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixA>56 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixB>57 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>58 gpu.func @gpu_wmma_mma_op(%A: !gpu.mma_matrix<16x16xf16, "AOp">,59 %B: !gpu.mma_matrix<16x16xf16, "BOp">,60 %C: !gpu.mma_matrix<16x16xf16, "COp">,61 %ptr: memref<16x16xf16, #spirv.storage_class<StorageBuffer>>) kernel62 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {63 // CHECK: %[[MAD:.*]] = spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}} :64 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixA>,65 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixB>66 // CHECK-SAME: -> !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>67 %D = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x16xf16, "AOp">,68 !gpu.mma_matrix<16x16xf16, "BOp">69 -> !gpu.mma_matrix<16x16xf16, "COp">70 71 %i = arith.constant 0 : index72 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %[[MAD]], %{{.+}}, <RowMajor>73 gpu.subgroup_mma_store_matrix %D, %ptr[%i,%i] {leadDimension = 32 : index} :74 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>75 // CHECK: spirv.Return76 gpu.return77 }78 79 // CHECK-LABEL: spirv.func @gpu_wmma_constant_op80 gpu.func @gpu_wmma_constant_op(%ptr: memref<16x16xf16, #spirv.storage_class<StorageBuffer>>) kernel81 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {82 // CHECK: %[[CST1F:.+]] = spirv.Constant 1.000000e+00 : f1683 %cst = arith.constant 1.0 : f1684 // CHECK: %[[MAT:.+]] = spirv.CompositeConstruct %[[CST1F]] :85 // CHECK-SAME: (f16) -> !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>86 %C = gpu.subgroup_mma_constant_matrix %cst : !gpu.mma_matrix<16x16xf16, "COp">87 88 %i = arith.constant 0 : index89 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %[[MAT]], %{{.+}}, <RowMajor>90 gpu.subgroup_mma_store_matrix %C, %ptr[%i,%i] {leadDimension = 32 : index} :91 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>92 // CHECK: spirv.Return93 gpu.return94 }95 96 // CHECK-LABEL: spirv.func @gpu_wmma_extract_thread_local_op97 // CHECK-SAME: %[[ARG0:.+]]: !spirv.coopmatrix<16x16xf32, Subgroup, MatrixA>98 gpu.func @gpu_wmma_extract_thread_local_op(%m: !gpu.mma_matrix<16x16xf32, "AOp">,99 %ptr: memref<16x16xf32, #spirv.storage_class<StorageBuffer>>) kernel100 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {101 // CHECK: spirv.CompositeExtract %[[ARG0]][0 : i32] : !spirv.coopmatrix<16x16xf32, Subgroup, MatrixA>102 %c0 = arith.constant 0 : index103 %val = gpu.subgroup_mma_extract_thread_local %m[%c0] : !gpu.mma_matrix<16x16xf32, "AOp"> -> f32104 memref.store %val, %ptr[%c0, %c0] : memref<16x16xf32, #spirv.storage_class<StorageBuffer>>105 gpu.return106 }107 108 // CHECK-LABEL: spirv.func @gpu_wmma_insert_thread_local_op109 // CHECK-SAME: %[[ARG0:.+]]: f16110 // CHECK-SAME: %[[ARG1:.+]]: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>111 gpu.func @gpu_wmma_insert_thread_local_op(%val: f16,112 %m: !gpu.mma_matrix<16x16xf16, "COp">,113 %ptr: memref<16x16xf16, #spirv.storage_class<StorageBuffer>>) kernel114 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {115 // CHECK: spirv.CompositeInsert %[[ARG0]], %[[ARG1]][0 : i32] : f16 into !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>116 %c0 = arith.constant 0 : index117 %s0 = gpu.subgroup_mma_insert_thread_local %val, %m[%c0] : f16, !gpu.mma_matrix<16x16xf16, "COp"> -> !gpu.mma_matrix<16x16xf16, "COp">118 gpu.subgroup_mma_store_matrix %s0, %ptr[%c0,%c0] {leadDimension = 16 : index} :119 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>120 gpu.return121 }122 123 // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op_default124 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>125 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>126 gpu.func @gpu_wmma_elementwise_op_default(%A: !gpu.mma_matrix<16x16xf16, "COp">,127 %B: !gpu.mma_matrix<16x16xf16, "COp">,128 %ptr: memref<16x16xf32, #spirv.storage_class<StorageBuffer>>) kernel129 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {130 // CHECK: {{%.*}} = spirv.FAdd {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>131 %C = gpu.subgroup_mma_elementwise addf %A, %B :132 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">133 // CHECK: {{%.*}} = spirv.FNegate {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>134 %D = gpu.subgroup_mma_elementwise negatef %C :135 (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">136 // CHECK: {{%.*}} = spirv.FDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>137 %E = gpu.subgroup_mma_elementwise divf %D, %A :138 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">139 // CHECK: {{%.*}} = spirv.FMul {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>140 %F = gpu.subgroup_mma_elementwise mulf %E, %A :141 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">142 // CHECK: {{%.*}} = spirv.FConvert {{%.*}} :143 // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc> to !spirv.coopmatrix<16x16xf32, Subgroup, MatrixAcc>144 %G = gpu.subgroup_mma_elementwise extf %F :145 (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">146 147 %i = arith.constant 0 : index148 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %{{.+}}, %{{.+}}, <RowMajor>149 gpu.subgroup_mma_store_matrix %G, %ptr[%i,%i] {leadDimension = 32 : index} :150 !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32, #spirv.storage_class<StorageBuffer>>151 // CHECK: spirv.Return152 gpu.return153 }154 155 // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op_matrix_times_scalar156 // CHECK-SAME: %[[A:.+]]: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>157 // CHECK-SAME: %[[S:.+]]: f16158 gpu.func @gpu_wmma_elementwise_op_matrix_times_scalar(159 %A: !gpu.mma_matrix<16x16xf16, "COp">, %scalar: f16,160 %ptr: memref<16x16xf16, #spirv.storage_class<StorageBuffer>>) kernel161 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {162 %i = arith.constant 0 : index163 164 %B = gpu.subgroup_mma_constant_matrix %scalar : !gpu.mma_matrix<16x16xf16, "COp">165 // CHECK: %[[C:.+]] = spirv.MatrixTimesScalar %[[A]], %[[S]] : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>, f16166 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %[[C]], %{{.+}}, <RowMajor>167 %C = gpu.subgroup_mma_elementwise mulf %A, %B :168 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">169 gpu.subgroup_mma_store_matrix %C, %ptr[%i,%i] {leadDimension = 32 : index} :170 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>171 172 // CHECK: %[[D:.+]] = spirv.MatrixTimesScalar %[[C]], %[[S]] : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>, f16173 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %[[D]], %{{.+}}, <RowMajor>174 %D = gpu.subgroup_mma_elementwise mulf %B, %C :175 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">176 gpu.subgroup_mma_store_matrix %D, %ptr[%i,%i] {leadDimension = 32 : index} :177 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>178 // CHECK: spirv.Return179 gpu.return180 }181 182 // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op_matrix_plus_scalar183 // CHECK-SAME: %[[A:.+]]: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>184 // CHECK-SAME: %[[S:.+]]: f16185 gpu.func @gpu_wmma_elementwise_op_matrix_plus_scalar(186 %A : !gpu.mma_matrix<16x16xf16, "COp">, %scalar : f16,187 %ptr: memref<16x16xf16, #spirv.storage_class<StorageBuffer>>) kernel188 attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {189 %i = arith.constant 0 : index190 191 // CHECK: %[[SM:.+]] = spirv.CompositeConstruct %[[S]] : (f16) -> !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>192 %B = gpu.subgroup_mma_constant_matrix %scalar : !gpu.mma_matrix<16x16xf16, "COp">193 // CHECK: %[[C:.+]] = spirv.FAdd %[[A]], %[[SM]] : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc>194 %C = gpu.subgroup_mma_elementwise addf %A, %B :195 (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">196 197 // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %[[C]], %{{.+}}, <RowMajor>198 gpu.subgroup_mma_store_matrix %C, %ptr[%i,%i] {leadDimension = 32 : index} :199 !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16, #spirv.storage_class<StorageBuffer>>200 // CHECK: spirv.Return201 gpu.return202 }203 }204}205