222 lines · plain
1// RUN: mlir-opt %s -split-input-file | FileCheck %s2 3// This file contains tests for sparse MMA (mma.sp.sync) operations with KIND variants.4// The kind::f8f6f4 variant was introduced in PTX ISA 8.7 for sm_90+ architectures.5//6// Based on PTX ISA documentation:7// https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-for-sparse-mma8//9// KIND::F8F6F4 enables:10// - Additional FP8 types: e3m2, e2m3, e2m111// - F16 accumulator for m16n8k64 FP8 operations12// - Mixed-precision FP8 computations13//14// Requirements:15// - ONLY works with ordered metadata (sp::ordered_metadata)16// - ONLY for shape m16n8k6417// - ONLY for FP8 types (not integers or other floats)18 19// =============================================================================20// FP8 e4m3 Sparse MMA with KIND (m16n8k64)21// =============================================================================22 23// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e4m3_f1624func.func @nvvm_mma_sp_kind_m16n8k64_e4m3_f16(25 %a0 : i32, %a1 : i32,26 %b0 : i32, %b1 : i32,27 %c0 : vector<2xf16>, %c1 : vector<2xf16>,28 %meta : i32, %sel : i32) {29 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>30 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]31 sparseMetadata[%meta] selector[%sel]32 {kind = #nvvm.mma_kind<f8f6f4>,33 orderedMetadata,34 multiplicandAPtxType = #nvvm.mma_type<e4m3>,35 multiplicandBPtxType = #nvvm.mma_type<e4m3>,36 shape = #nvvm.shape<m = 16, n = 8, k = 64>}37 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>38 return39}40 41// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e4m3_f3242func.func @nvvm_mma_sp_kind_m16n8k64_e4m3_f32(43 %a0 : i32, %a1 : i32,44 %b0 : i32, %b1 : i32,45 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,46 %meta : i32, %sel : i32) {47 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>48 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]49 sparseMetadata[%meta] selector[%sel]50 {kind = #nvvm.mma_kind<f8f6f4>,51 orderedMetadata,52 multiplicandAPtxType = #nvvm.mma_type<e4m3>,53 multiplicandBPtxType = #nvvm.mma_type<e4m3>,54 shape = #nvvm.shape<m = 16, n = 8, k = 64>}55 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>56 return57}58 59// =============================================================================60// FP8 e5m2 Sparse MMA with KIND (m16n8k64)61// =============================================================================62 63// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e5m2_f1664func.func @nvvm_mma_sp_kind_m16n8k64_e5m2_f16(65 %a0 : i32, %a1 : i32,66 %b0 : i32, %b1 : i32,67 %c0 : vector<2xf16>, %c1 : vector<2xf16>,68 %meta : i32, %sel : i32) {69 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>70 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]71 sparseMetadata[%meta] selector[%sel]72 {kind = #nvvm.mma_kind<f8f6f4>,73 orderedMetadata,74 multiplicandAPtxType = #nvvm.mma_type<e5m2>,75 multiplicandBPtxType = #nvvm.mma_type<e5m2>,76 shape = #nvvm.shape<m = 16, n = 8, k = 64>}77 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>78 return79}80 81// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e5m2_f3282func.func @nvvm_mma_sp_kind_m16n8k64_e5m2_f32(83 %a0 : i32, %a1 : i32,84 %b0 : i32, %b1 : i32,85 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,86 %meta : i32, %sel : i32) {87 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>88 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]89 sparseMetadata[%meta] selector[%sel]90 {kind = #nvvm.mma_kind<f8f6f4>,91 orderedMetadata,92 multiplicandAPtxType = #nvvm.mma_type<e5m2>,93 multiplicandBPtxType = #nvvm.mma_type<e5m2>,94 shape = #nvvm.shape<m = 16, n = 8, k = 64>}95 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>96 return97}98 99// =============================================================================100// FP8 e3m2 Sparse MMA with KIND (m16n8k64)101// NOTE: e3m2 is ONLY available with kind::f8f6f4102// =============================================================================103 104// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e3m2_f16105func.func @nvvm_mma_sp_kind_m16n8k64_e3m2_f16(106 %a0 : i32, %a1 : i32,107 %b0 : i32, %b1 : i32,108 %c0 : vector<2xf16>, %c1 : vector<2xf16>,109 %meta : i32, %sel : i32) {110 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e3m2>, multiplicandBPtxType = #nvvm.mma_type<e3m2>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>111 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]112 sparseMetadata[%meta] selector[%sel]113 {kind = #nvvm.mma_kind<f8f6f4>,114 orderedMetadata,115 multiplicandAPtxType = #nvvm.mma_type<e3m2>,116 multiplicandBPtxType = #nvvm.mma_type<e3m2>,117 shape = #nvvm.shape<m = 16, n = 8, k = 64>}118 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>119 return120}121 122// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e3m2_f32123func.func @nvvm_mma_sp_kind_m16n8k64_e3m2_f32(124 %a0 : i32, %a1 : i32,125 %b0 : i32, %b1 : i32,126 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,127 %meta : i32, %sel : i32) {128 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e3m2>, multiplicandBPtxType = #nvvm.mma_type<e3m2>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>129 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]130 sparseMetadata[%meta] selector[%sel]131 {kind = #nvvm.mma_kind<f8f6f4>,132 orderedMetadata,133 multiplicandAPtxType = #nvvm.mma_type<e3m2>,134 multiplicandBPtxType = #nvvm.mma_type<e3m2>,135 shape = #nvvm.shape<m = 16, n = 8, k = 64>}136 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>137 return138}139 140// =============================================================================141// FP8 e2m3 Sparse MMA with KIND (m16n8k64)142// NOTE: e2m3 is ONLY available with kind::f8f6f4143// =============================================================================144 145// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e2m3_f16146func.func @nvvm_mma_sp_kind_m16n8k64_e2m3_f16(147 %a0 : i32, %a1 : i32,148 %b0 : i32, %b1 : i32,149 %c0 : vector<2xf16>, %c1 : vector<2xf16>,150 %meta : i32, %sel : i32) {151 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e2m3>, multiplicandBPtxType = #nvvm.mma_type<e2m3>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>152 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]153 sparseMetadata[%meta] selector[%sel]154 {kind = #nvvm.mma_kind<f8f6f4>,155 orderedMetadata,156 multiplicandAPtxType = #nvvm.mma_type<e2m3>,157 multiplicandBPtxType = #nvvm.mma_type<e2m3>,158 shape = #nvvm.shape<m = 16, n = 8, k = 64>}159 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>160 return161}162 163// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e2m3_f32164func.func @nvvm_mma_sp_kind_m16n8k64_e2m3_f32(165 %a0 : i32, %a1 : i32,166 %b0 : i32, %b1 : i32,167 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,168 %meta : i32, %sel : i32) {169 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e2m3>, multiplicandBPtxType = #nvvm.mma_type<e2m3>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>170 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]171 sparseMetadata[%meta] selector[%sel]172 {kind = #nvvm.mma_kind<f8f6f4>,173 orderedMetadata,174 multiplicandAPtxType = #nvvm.mma_type<e2m3>,175 multiplicandBPtxType = #nvvm.mma_type<e2m3>,176 shape = #nvvm.shape<m = 16, n = 8, k = 64>}177 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>178 return179}180 181// =============================================================================182// FP8 e2m1 Sparse MMA with KIND (m16n8k64)183// NOTE: e2m1 is ONLY available with kind::f8f6f4184// =============================================================================185 186// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e2m1_f16187func.func @nvvm_mma_sp_kind_m16n8k64_e2m1_f16(188 %a0 : i32, %a1 : i32,189 %b0 : i32, %b1 : i32,190 %c0 : vector<2xf16>, %c1 : vector<2xf16>,191 %meta : i32, %sel : i32) {192 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e2m1>, multiplicandBPtxType = #nvvm.mma_type<e2m1>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>193 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]194 sparseMetadata[%meta] selector[%sel]195 {kind = #nvvm.mma_kind<f8f6f4>,196 orderedMetadata,197 multiplicandAPtxType = #nvvm.mma_type<e2m1>,198 multiplicandBPtxType = #nvvm.mma_type<e2m1>,199 shape = #nvvm.shape<m = 16, n = 8, k = 64>}200 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>201 return202}203 204// CHECK-LABEL: @nvvm_mma_sp_kind_m16n8k64_e2m1_f32205func.func @nvvm_mma_sp_kind_m16n8k64_e2m1_f32(206 %a0 : i32, %a1 : i32,207 %b0 : i32, %b1 : i32,208 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,209 %meta : i32, %sel : i32) {210 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {kind = #nvvm.mma_kind<f8f6f4>, multiplicandAPtxType = #nvvm.mma_type<e2m1>, multiplicandBPtxType = #nvvm.mma_type<e2m1>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>211 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]212 sparseMetadata[%meta] selector[%sel]213 {kind = #nvvm.mma_kind<f8f6f4>,214 orderedMetadata,215 multiplicandAPtxType = #nvvm.mma_type<e2m1>,216 multiplicandBPtxType = #nvvm.mma_type<e2m1>,217 shape = #nvvm.shape<m = 16, n = 8, k = 64>}218 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>219 return220}221 222