391 lines · plain
1// RUN: mlir-opt %s -split-input-file | FileCheck %s2 3// This file contains tests for all sparse MMA (mma.sp.sync) operations in the NVVM dialect4// Based on PTX ISA documentation:5// https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-for-sparse-mma6//7// Sparse MMA operations follow 2:4 structured sparsity where 2 out of every 4 elements8// in the A operand are non-zero. The A operand is provided in compressed form,9// and sparseMetadata provides the sparsity indices.10//11// NOTE: These tests use the default (standard) metadata ordering.12// For ordered metadata tests (PTX ISA 8.5+, sm_90+), see nvvm-mma-sp-ordered.mlir.13 14// =============================================================================15// F16 Sparse MMA Operations (m16n8k16)16// =============================================================================17 18// CHECK-LABEL: @nvvm_mma_sp_m16n8k16_f16_f1619func.func @nvvm_mma_sp_m16n8k16_f16_f16(20 %a0 : vector<2xf16>, %a1 : vector<2xf16>,21 %b0 : vector<2xf16>, %b1 : vector<2xf16>,22 %c0 : vector<2xf16>, %c1 : vector<2xf16>,23 %meta : i32, %sel : i32) {24 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>25 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]26 sparseMetadata[%meta] selector[%sel]27 {shape = #nvvm.shape<m = 16, n = 8, k = 16>}28 : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>29 return30}31 32// CHECK-LABEL: @nvvm_mma_sp_m16n8k16_f16_f3233func.func @nvvm_mma_sp_m16n8k16_f16_f32(34 %a0 : vector<2xf16>, %a1 : vector<2xf16>,35 %b0 : vector<2xf16>, %b1 : vector<2xf16>,36 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,37 %meta : i32, %sel : i32) {38 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>39 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]40 sparseMetadata[%meta] selector[%sel]41 {shape = #nvvm.shape<m = 16, n = 8, k = 16>}42 : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>43 return44}45 46// =============================================================================47// F16 Sparse MMA Operations (m16n8k32)48// =============================================================================49 50// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_f16_f1651func.func @nvvm_mma_sp_m16n8k32_f16_f16(52 %a0 : vector<2xf16>, %a1 : vector<2xf16>, %a2 : vector<2xf16>, %a3 : vector<2xf16>,53 %b0 : vector<2xf16>, %b1 : vector<2xf16>, %b2 : vector<2xf16>, %b3 : vector<2xf16>,54 %c0 : vector<2xf16>, %c1 : vector<2xf16>,55 %meta : i32, %sel : i32) {56 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>57 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1]58 sparseMetadata[%meta] selector[%sel]59 {shape = #nvvm.shape<m = 16, n = 8, k = 32>}60 : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>61 return62}63 64// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_f16_f3265func.func @nvvm_mma_sp_m16n8k32_f16_f32(66 %a0 : vector<2xf16>, %a1 : vector<2xf16>, %a2 : vector<2xf16>, %a3 : vector<2xf16>,67 %b0 : vector<2xf16>, %b1 : vector<2xf16>, %b2 : vector<2xf16>, %b3 : vector<2xf16>,68 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,69 %meta : i32, %sel : i32) {70 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>71 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]72 sparseMetadata[%meta] selector[%sel]73 {shape = #nvvm.shape<m = 16, n = 8, k = 32>}74 : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>75 return76}77 78// =============================================================================79// BF16 Sparse MMA Operations80// =============================================================================81 82// CHECK-LABEL: @nvvm_mma_sp_m16n8k16_bf16_f3283func.func @nvvm_mma_sp_m16n8k16_bf16_f32(84 %a0 : i32, %a1 : i32,85 %b0 : i32, %b1 : i32,86 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,87 %meta : i32, %sel : i32) {88 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<bf16>, multiplicandBPtxType = #nvvm.mma_type<bf16>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>89 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]90 sparseMetadata[%meta] selector[%sel]91 {multiplicandAPtxType = #nvvm.mma_type<bf16>,92 multiplicandBPtxType = #nvvm.mma_type<bf16>,93 shape = #nvvm.shape<m = 16, n = 8, k = 16>}94 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>95 return96}97 98// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_bf16_f3299func.func @nvvm_mma_sp_m16n8k32_bf16_f32(100 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,101 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,102 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,103 %meta : i32, %sel : i32) {104 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<bf16>, multiplicandBPtxType = #nvvm.mma_type<bf16>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>105 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]106 sparseMetadata[%meta] selector[%sel]107 {multiplicandAPtxType = #nvvm.mma_type<bf16>,108 multiplicandBPtxType = #nvvm.mma_type<bf16>,109 shape = #nvvm.shape<m = 16, n = 8, k = 32>}110 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>111 return112}113 114// =============================================================================115// TF32 Sparse MMA Operations116// =============================================================================117 118// CHECK-LABEL: @nvvm_mma_sp_m16n8k8_tf32_f32119func.func @nvvm_mma_sp_m16n8k8_tf32_f32(120 %a0 : i32, %a1 : i32,121 %b0 : i32, %b1 : i32,122 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,123 %meta : i32, %sel : i32) {124 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>, shape = #nvvm.shape<m = 16, n = 8, k = 8>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>125 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]126 sparseMetadata[%meta] selector[%sel]127 {multiplicandAPtxType = #nvvm.mma_type<tf32>,128 multiplicandBPtxType = #nvvm.mma_type<tf32>,129 shape = #nvvm.shape<m = 16, n = 8, k = 8>}130 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>131 return132}133 134// CHECK-LABEL: @nvvm_mma_sp_m16n8k16_tf32_f32135func.func @nvvm_mma_sp_m16n8k16_tf32_f32(136 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,137 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,138 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,139 %meta : i32, %sel : i32) {140 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>141 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]142 sparseMetadata[%meta] selector[%sel]143 {multiplicandAPtxType = #nvvm.mma_type<tf32>,144 multiplicandBPtxType = #nvvm.mma_type<tf32>,145 shape = #nvvm.shape<m = 16, n = 8, k = 16>}146 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>147 return148}149 150// =============================================================================151// Integer (s8) Sparse MMA Operations152// =============================================================================153 154// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_s8_s32155func.func @nvvm_mma_sp_m16n8k32_s8_s32(156 %a0 : i32, %a1 : i32,157 %b0 : i32, %b1 : i32,158 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,159 %meta : i32, %sel : i32) {160 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>161 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]162 sparseMetadata[%meta] selector[%sel]163 {multiplicandAPtxType = #nvvm.mma_type<s8>,164 multiplicandBPtxType = #nvvm.mma_type<s8>,165 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,166 shape = #nvvm.shape<m = 16, n = 8, k = 32>}167 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>168 return169}170 171// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_s8_s32_satfinite172func.func @nvvm_mma_sp_m16n8k32_s8_s32_satfinite(173 %a0 : i32, %a1 : i32,174 %b0 : i32, %b1 : i32,175 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,176 %meta : i32, %sel : i32) {177 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>178 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]179 sparseMetadata[%meta] selector[%sel]180 {multiplicandAPtxType = #nvvm.mma_type<s8>,181 multiplicandBPtxType = #nvvm.mma_type<s8>,182 intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>,183 shape = #nvvm.shape<m = 16, n = 8, k = 32>}184 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>185 return186}187 188// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_s8_s32189func.func @nvvm_mma_sp_m16n8k64_s8_s32(190 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,191 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,192 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,193 %meta : i32, %sel : i32) {194 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>195 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]196 sparseMetadata[%meta] selector[%sel]197 {multiplicandAPtxType = #nvvm.mma_type<s8>,198 multiplicandBPtxType = #nvvm.mma_type<s8>,199 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,200 shape = #nvvm.shape<m = 16, n = 8, k = 64>}201 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>202 return203}204 205// =============================================================================206// Integer (u8) Sparse MMA Operations207// =============================================================================208 209// CHECK-LABEL: @nvvm_mma_sp_m16n8k32_u8_s32210func.func @nvvm_mma_sp_m16n8k32_u8_s32(211 %a0 : i32, %a1 : i32,212 %b0 : i32, %b1 : i32,213 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,214 %meta : i32, %sel : i32) {215 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<u8>, multiplicandBPtxType = #nvvm.mma_type<u8>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>216 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]217 sparseMetadata[%meta] selector[%sel]218 {multiplicandAPtxType = #nvvm.mma_type<u8>,219 multiplicandBPtxType = #nvvm.mma_type<u8>,220 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,221 shape = #nvvm.shape<m = 16, n = 8, k = 32>}222 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>223 return224}225 226// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_u8_s32227func.func @nvvm_mma_sp_m16n8k64_u8_s32(228 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,229 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,230 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,231 %meta : i32, %sel : i32) {232 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<u8>, multiplicandBPtxType = #nvvm.mma_type<u8>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>233 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]234 sparseMetadata[%meta] selector[%sel]235 {multiplicandAPtxType = #nvvm.mma_type<u8>,236 multiplicandBPtxType = #nvvm.mma_type<u8>,237 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,238 shape = #nvvm.shape<m = 16, n = 8, k = 64>}239 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>240 return241}242 243// =============================================================================244// Sub-byte Integer (s4) Sparse MMA Operations245// =============================================================================246 247// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_s4_s32248func.func @nvvm_mma_sp_m16n8k64_s4_s32(249 %a0 : i32, %a1 : i32,250 %b0 : i32, %b1 : i32,251 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,252 %meta : i32, %sel : i32) {253 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<s4>, multiplicandBPtxType = #nvvm.mma_type<s4>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>254 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]255 sparseMetadata[%meta] selector[%sel]256 {multiplicandAPtxType = #nvvm.mma_type<s4>,257 multiplicandBPtxType = #nvvm.mma_type<s4>,258 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,259 shape = #nvvm.shape<m = 16, n = 8, k = 64>}260 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>261 return262}263 264// CHECK-LABEL: @nvvm_mma_sp_m16n8k128_s4_s32265func.func @nvvm_mma_sp_m16n8k128_s4_s32(266 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,267 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,268 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,269 %meta : i32, %sel : i32) {270 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<s4>, multiplicandBPtxType = #nvvm.mma_type<s4>, shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>271 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]272 sparseMetadata[%meta] selector[%sel]273 {multiplicandAPtxType = #nvvm.mma_type<s4>,274 multiplicandBPtxType = #nvvm.mma_type<s4>,275 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,276 shape = #nvvm.shape<m = 16, n = 8, k = 128>}277 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>278 return279}280 281// =============================================================================282// Sub-byte Integer (u4) Sparse MMA Operations283// =============================================================================284 285// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_u4_s32286func.func @nvvm_mma_sp_m16n8k64_u4_s32(287 %a0 : i32, %a1 : i32,288 %b0 : i32, %b1 : i32,289 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,290 %meta : i32, %sel : i32) {291 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<u4>, multiplicandBPtxType = #nvvm.mma_type<u4>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>292 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]293 sparseMetadata[%meta] selector[%sel]294 {multiplicandAPtxType = #nvvm.mma_type<u4>,295 multiplicandBPtxType = #nvvm.mma_type<u4>,296 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,297 shape = #nvvm.shape<m = 16, n = 8, k = 64>}298 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>299 return300}301 302// CHECK-LABEL: @nvvm_mma_sp_m16n8k128_u4_s32303func.func @nvvm_mma_sp_m16n8k128_u4_s32(304 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,305 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,306 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,307 %meta : i32, %sel : i32) {308 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, multiplicandAPtxType = #nvvm.mma_type<u4>, multiplicandBPtxType = #nvvm.mma_type<u4>, shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>309 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]310 sparseMetadata[%meta] selector[%sel]311 {multiplicandAPtxType = #nvvm.mma_type<u4>,312 multiplicandBPtxType = #nvvm.mma_type<u4>,313 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,314 shape = #nvvm.shape<m = 16, n = 8, k = 128>}315 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>316 return317}318 319// =============================================================================320// FP8 (e4m3) Sparse MMA Operations321// =============================================================================322 323// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_e4m3_f16324func.func @nvvm_mma_sp_m16n8k64_e4m3_f16(325 %a0 : i32, %a1 : i32,326 %b0 : i32, %b1 : i32,327 %c0 : vector<2xf16>, %c1 : vector<2xf16>,328 %meta : i32, %sel : i32) {329 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>330 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]331 sparseMetadata[%meta] selector[%sel]332 {multiplicandAPtxType = #nvvm.mma_type<e4m3>,333 multiplicandBPtxType = #nvvm.mma_type<e4m3>,334 shape = #nvvm.shape<m = 16, n = 8, k = 64>}335 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>336 return337}338 339// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_e4m3_f32340func.func @nvvm_mma_sp_m16n8k64_e4m3_f32(341 %a0 : i32, %a1 : i32,342 %b0 : i32, %b1 : i32,343 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,344 %meta : i32, %sel : i32) {345 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>346 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]347 sparseMetadata[%meta] selector[%sel]348 {multiplicandAPtxType = #nvvm.mma_type<e4m3>,349 multiplicandBPtxType = #nvvm.mma_type<e4m3>,350 shape = #nvvm.shape<m = 16, n = 8, k = 64>}351 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>352 return353}354 355// =============================================================================356// FP8 (e5m2) Sparse MMA Operations357// =============================================================================358 359// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_e5m2_f16360func.func @nvvm_mma_sp_m16n8k64_e5m2_f16(361 %a0 : i32, %a1 : i32,362 %b0 : i32, %b1 : i32,363 %c0 : vector<2xf16>, %c1 : vector<2xf16>,364 %meta : i32, %sel : i32) {365 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>366 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]367 sparseMetadata[%meta] selector[%sel]368 {multiplicandAPtxType = #nvvm.mma_type<e5m2>,369 multiplicandBPtxType = #nvvm.mma_type<e5m2>,370 shape = #nvvm.shape<m = 16, n = 8, k = 64>}371 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>372 return373}374 375// CHECK-LABEL: @nvvm_mma_sp_m16n8k64_e5m2_f32376func.func @nvvm_mma_sp_m16n8k64_e5m2_f32(377 %a0 : i32, %a1 : i32,378 %b0 : i32, %b1 : i32,379 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,380 %meta : i32, %sel : i32) {381 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>382 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]383 sparseMetadata[%meta] selector[%sel]384 {multiplicandAPtxType = #nvvm.mma_type<e5m2>,385 multiplicandBPtxType = #nvvm.mma_type<e5m2>,386 shape = #nvvm.shape<m = 16, n = 8, k = 64>}387 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>388 return389}390 391