412 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 ORDERED metadata.4// The ordered metadata variant was introduced in PTX ISA 8.5 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// Ordered metadata provides an alternative metadata ordering for 2:4 structured sparsity10// that can offer better performance on newer architectures.11 12// =============================================================================13// F16 Sparse MMA Operations with Ordered Metadata (m16n8k16)14// =============================================================================15 16// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k16_f16_f1617func.func @nvvm_mma_sp_ordered_m16n8k16_f16_f16(18 %a0 : vector<2xf16>, %a1 : vector<2xf16>,19 %b0 : vector<2xf16>, %b1 : vector<2xf16>,20 %c0 : vector<2xf16>, %c1 : vector<2xf16>,21 %meta : i32, %sel : i32) {22 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>23 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]24 sparseMetadata[%meta] selector[%sel]25 {orderedMetadata,26 shape = #nvvm.shape<m = 16, n = 8, k = 16>}27 : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>28 return29}30 31// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k16_f16_f3232func.func @nvvm_mma_sp_ordered_m16n8k16_f16_f32(33 %a0 : vector<2xf16>, %a1 : vector<2xf16>,34 %b0 : vector<2xf16>, %b1 : vector<2xf16>,35 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,36 %meta : i32, %sel : i32) {37 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>38 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]39 sparseMetadata[%meta] selector[%sel]40 {orderedMetadata,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 with Ordered Metadata (m16n8k32)48// =============================================================================49 50// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_f16_f1651func.func @nvvm_mma_sp_ordered_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[{{.*}}] {orderedMetadata, 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 {orderedMetadata,60 shape = #nvvm.shape<m = 16, n = 8, k = 32>}61 : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>62 return63}64 65// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_f16_f3266func.func @nvvm_mma_sp_ordered_m16n8k32_f16_f32(67 %a0 : vector<2xf16>, %a1 : vector<2xf16>, %a2 : vector<2xf16>, %a3 : vector<2xf16>,68 %b0 : vector<2xf16>, %b1 : vector<2xf16>, %b2 : vector<2xf16>, %b3 : vector<2xf16>,69 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,70 %meta : i32, %sel : i32) {71 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>72 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]73 sparseMetadata[%meta] selector[%sel]74 {orderedMetadata,75 shape = #nvvm.shape<m = 16, n = 8, k = 32>}76 : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>77 return78}79 80// =============================================================================81// BF16 Sparse MMA Operations with Ordered Metadata82// =============================================================================83 84// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k16_bf16_f3285func.func @nvvm_mma_sp_ordered_m16n8k16_bf16_f32(86 %a0 : i32, %a1 : i32,87 %b0 : i32, %b1 : i32,88 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,89 %meta : i32, %sel : i32) {90 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<bf16>, multiplicandBPtxType = #nvvm.mma_type<bf16>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>91 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]92 sparseMetadata[%meta] selector[%sel]93 {orderedMetadata,94 multiplicandAPtxType = #nvvm.mma_type<bf16>,95 multiplicandBPtxType = #nvvm.mma_type<bf16>,96 shape = #nvvm.shape<m = 16, n = 8, k = 16>}97 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>98 return99}100 101// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_bf16_f32102func.func @nvvm_mma_sp_ordered_m16n8k32_bf16_f32(103 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,104 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,105 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,106 %meta : i32, %sel : i32) {107 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<bf16>, multiplicandBPtxType = #nvvm.mma_type<bf16>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>108 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]109 sparseMetadata[%meta] selector[%sel]110 {orderedMetadata,111 multiplicandAPtxType = #nvvm.mma_type<bf16>,112 multiplicandBPtxType = #nvvm.mma_type<bf16>,113 shape = #nvvm.shape<m = 16, n = 8, k = 32>}114 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>115 return116}117 118// =============================================================================119// TF32 Sparse MMA Operations with Ordered Metadata120// =============================================================================121 122// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k8_tf32_f32123func.func @nvvm_mma_sp_ordered_m16n8k8_tf32_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[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 8>} : (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 {orderedMetadata,132 multiplicandAPtxType = #nvvm.mma_type<tf32>,133 multiplicandBPtxType = #nvvm.mma_type<tf32>,134 shape = #nvvm.shape<m = 16, n = 8, k = 8>}135 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>136 return137}138 139// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k16_tf32_f32140func.func @nvvm_mma_sp_ordered_m16n8k16_tf32_f32(141 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,142 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,143 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,144 %meta : i32, %sel : i32) {145 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}, {{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>146 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]147 sparseMetadata[%meta] selector[%sel]148 {orderedMetadata,149 multiplicandAPtxType = #nvvm.mma_type<tf32>,150 multiplicandBPtxType = #nvvm.mma_type<tf32>,151 shape = #nvvm.shape<m = 16, n = 8, k = 16>}152 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>153 return154}155 156// =============================================================================157// Integer (s8) Sparse MMA Operations with Ordered Metadata158// =============================================================================159 160// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_s8_s32161func.func @nvvm_mma_sp_ordered_m16n8k32_s8_s32(162 %a0 : i32, %a1 : i32,163 %b0 : i32, %b1 : i32,164 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,165 %meta : i32, %sel : i32) {166 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>167 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]168 sparseMetadata[%meta] selector[%sel]169 {orderedMetadata,170 multiplicandAPtxType = #nvvm.mma_type<s8>,171 multiplicandBPtxType = #nvvm.mma_type<s8>,172 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,173 shape = #nvvm.shape<m = 16, n = 8, k = 32>}174 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>175 return176}177 178// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_s8_s32_satfinite179func.func @nvvm_mma_sp_ordered_m16n8k32_s8_s32_satfinite(180 %a0 : i32, %a1 : i32,181 %b0 : i32, %b1 : i32,182 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,183 %meta : i32, %sel : i32) {184 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>185 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]186 sparseMetadata[%meta] selector[%sel]187 {orderedMetadata,188 multiplicandAPtxType = #nvvm.mma_type<s8>,189 multiplicandBPtxType = #nvvm.mma_type<s8>,190 intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>,191 shape = #nvvm.shape<m = 16, n = 8, k = 32>}192 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>193 return194}195 196// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_s8_s32197func.func @nvvm_mma_sp_ordered_m16n8k64_s8_s32(198 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,199 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,200 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,201 %meta : i32, %sel : i32) {202 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>203 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]204 sparseMetadata[%meta] selector[%sel]205 {orderedMetadata,206 multiplicandAPtxType = #nvvm.mma_type<s8>,207 multiplicandBPtxType = #nvvm.mma_type<s8>,208 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,209 shape = #nvvm.shape<m = 16, n = 8, k = 64>}210 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>211 return212}213 214// =============================================================================215// Integer (u8) Sparse MMA Operations with Ordered Metadata216// =============================================================================217 218// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k32_u8_s32219func.func @nvvm_mma_sp_ordered_m16n8k32_u8_s32(220 %a0 : i32, %a1 : i32,221 %b0 : i32, %b1 : i32,222 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,223 %meta : i32, %sel : i32) {224 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>225 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]226 sparseMetadata[%meta] selector[%sel]227 {orderedMetadata,228 multiplicandAPtxType = #nvvm.mma_type<u8>,229 multiplicandBPtxType = #nvvm.mma_type<u8>,230 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,231 shape = #nvvm.shape<m = 16, n = 8, k = 32>}232 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>233 return234}235 236// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_u8_s32237func.func @nvvm_mma_sp_ordered_m16n8k64_u8_s32(238 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,239 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,240 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,241 %meta : i32, %sel : i32) {242 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>243 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]244 sparseMetadata[%meta] selector[%sel]245 {orderedMetadata,246 multiplicandAPtxType = #nvvm.mma_type<u8>,247 multiplicandBPtxType = #nvvm.mma_type<u8>,248 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,249 shape = #nvvm.shape<m = 16, n = 8, k = 64>}250 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>251 return252}253 254// =============================================================================255// Sub-byte Integer (s4) Sparse MMA Operations with Ordered Metadata256// =============================================================================257 258// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_s4_s32259func.func @nvvm_mma_sp_ordered_m16n8k64_s4_s32(260 %a0 : i32, %a1 : i32,261 %b0 : i32, %b1 : i32,262 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,263 %meta : i32, %sel : i32) {264 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>265 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]266 sparseMetadata[%meta] selector[%sel]267 {orderedMetadata,268 multiplicandAPtxType = #nvvm.mma_type<s4>,269 multiplicandBPtxType = #nvvm.mma_type<s4>,270 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,271 shape = #nvvm.shape<m = 16, n = 8, k = 64>}272 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>273 return274}275 276// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k128_s4_s32277func.func @nvvm_mma_sp_ordered_m16n8k128_s4_s32(278 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,279 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,280 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,281 %meta : i32, %sel : i32) {282 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>283 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]284 sparseMetadata[%meta] selector[%sel]285 {orderedMetadata,286 multiplicandAPtxType = #nvvm.mma_type<s4>,287 multiplicandBPtxType = #nvvm.mma_type<s4>,288 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,289 shape = #nvvm.shape<m = 16, n = 8, k = 128>}290 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>291 return292}293 294// =============================================================================295// Sub-byte Integer (u4) Sparse MMA Operations with Ordered Metadata296// =============================================================================297 298// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_u4_s32299func.func @nvvm_mma_sp_ordered_m16n8k64_u4_s32(300 %a0 : i32, %a1 : i32,301 %b0 : i32, %b1 : i32,302 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,303 %meta : i32, %sel : i32) {304 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>305 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]306 sparseMetadata[%meta] selector[%sel]307 {orderedMetadata,308 multiplicandAPtxType = #nvvm.mma_type<u4>,309 multiplicandBPtxType = #nvvm.mma_type<u4>,310 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,311 shape = #nvvm.shape<m = 16, n = 8, k = 64>}312 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>313 return314}315 316// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k128_u4_s32317func.func @nvvm_mma_sp_ordered_m16n8k128_u4_s32(318 %a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,319 %b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32,320 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32,321 %meta : i32, %sel : i32) {322 // 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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>323 %0 = nvvm.mma.sp.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1, %b2, %b3] C[%c0, %c1, %c2, %c3]324 sparseMetadata[%meta] selector[%sel]325 {orderedMetadata,326 multiplicandAPtxType = #nvvm.mma_type<u4>,327 multiplicandBPtxType = #nvvm.mma_type<u4>,328 intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>,329 shape = #nvvm.shape<m = 16, n = 8, k = 128>}330 : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>331 return332}333 334// =============================================================================335// FP8 (e4m3) Sparse MMA Operations with Ordered Metadata336// NOTE: FP8 ordered metadata requires PTX ISA 8.7+ and sm_90+337// =============================================================================338 339// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_e4m3_f16340func.func @nvvm_mma_sp_ordered_m16n8k64_e4m3_f16(341 %a0 : i32, %a1 : i32,342 %b0 : i32, %b1 : i32,343 %c0 : vector<2xf16>, %c1 : vector<2xf16>,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>, orderedMetadata, shape = #nvvm.shape<m = 16, n = 8, k = 64>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>346 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]347 sparseMetadata[%meta] selector[%sel]348 {orderedMetadata,349 multiplicandAPtxType = #nvvm.mma_type<e4m3>,350 multiplicandBPtxType = #nvvm.mma_type<e4m3>,351 shape = #nvvm.shape<m = 16, n = 8, k = 64>}352 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>353 return354}355 356// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_e4m3_f32357func.func @nvvm_mma_sp_ordered_m16n8k64_e4m3_f32(358 %a0 : i32, %a1 : i32,359 %b0 : i32, %b1 : i32,360 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,361 %meta : i32, %sel : i32) {362 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {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)>363 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]364 sparseMetadata[%meta] selector[%sel]365 {orderedMetadata,366 multiplicandAPtxType = #nvvm.mma_type<e4m3>,367 multiplicandBPtxType = #nvvm.mma_type<e4m3>,368 shape = #nvvm.shape<m = 16, n = 8, k = 64>}369 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>370 return371}372 373// =============================================================================374// FP8 (e5m2) Sparse MMA Operations with Ordered Metadata375// NOTE: FP8 ordered metadata requires PTX ISA 8.7+ and sm_90+376// =============================================================================377 378// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_e5m2_f16379func.func @nvvm_mma_sp_ordered_m16n8k64_e5m2_f16(380 %a0 : i32, %a1 : i32,381 %b0 : i32, %b1 : i32,382 %c0 : vector<2xf16>, %c1 : vector<2xf16>,383 %meta : i32, %sel : i32) {384 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {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>)>385 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]386 sparseMetadata[%meta] selector[%sel]387 {orderedMetadata,388 multiplicandAPtxType = #nvvm.mma_type<e5m2>,389 multiplicandBPtxType = #nvvm.mma_type<e5m2>,390 shape = #nvvm.shape<m = 16, n = 8, k = 64>}391 : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>392 return393}394 395// CHECK-LABEL: @nvvm_mma_sp_ordered_m16n8k64_e5m2_f32396func.func @nvvm_mma_sp_ordered_m16n8k64_e5m2_f32(397 %a0 : i32, %a1 : i32,398 %b0 : i32, %b1 : i32,399 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,400 %meta : i32, %sel : i32) {401 // CHECK: nvvm.mma.sp.sync A[{{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] sparseMetadata[{{.*}}] selector[{{.*}}] {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)>402 %0 = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]403 sparseMetadata[%meta] selector[%sel]404 {orderedMetadata,405 multiplicandAPtxType = #nvvm.mma_type<e5m2>,406 multiplicandBPtxType = #nvvm.mma_type<e5m2>,407 shape = #nvvm.shape<m = 16, n = 8, k = 64>}408 : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>409 return410}411 412