753 lines · plain
1// RUN: mlir-opt --convert-nvvm-to-llvm --convert-arith-to-llvm --split-input-file %s | FileCheck %s2 3// Same below, but using the `ConvertToLLVMPatternInterface` entry point4// and the generic `convert-to-llvm` pass.5// RUN: mlir-opt --convert-to-llvm --split-input-file %s | FileCheck %s6// RUN: mlir-opt --convert-to-llvm="allow-pattern-rollback=0" --split-input-file %s | FileCheck %s7 8// CHECK-LABEL: @init_mbarrier9llvm.func @init_mbarrier(%barrier_gen : !llvm.ptr, %barrier : !llvm.ptr<3>, %count : i32, %pred : i1) {10 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.shared.b64 [$0], $1;", "r,r,b" 11 nvvm.mbarrier.init %barrier, %count, predicate = %pred : !llvm.ptr<3>, i32, i1 12 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" 13 nvvm.mbarrier.init %barrier_gen, %count, predicate = %pred : !llvm.ptr, i32, i114 llvm.return15}16 17// CHECK-LABEL: @init_mbarrier_arrive_expect_tx18llvm.func @init_mbarrier_arrive_expect_tx(%barrier : !llvm.ptr<3>, %txcount : i32, %pred : i1) {19 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.arrive.expect_tx.shared.b64 _, [$0], $1;", "r,r,b"20 nvvm.mbarrier.arrive.expect_tx %barrier, %txcount, predicate = %pred : !llvm.ptr<3>, i32, i1 21 llvm.return22}23 24// CHECK-LABEL: @init_mbarrier_arrive_expect_tx_generic25llvm.func @init_mbarrier_arrive_expect_tx_generic(%barrier : !llvm.ptr, %txcount : i32, %pred : i1) {26 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.arrive.expect_tx.b64 _, [$0], $1;", "l,r,b"27 nvvm.mbarrier.arrive.expect_tx %barrier, %txcount, predicate = %pred : !llvm.ptr, i32, i1 28 llvm.return29}30 31// CHECK-LABEL: @init_mbarrier_try_wait_shared32llvm.func @init_mbarrier_try_wait_shared(%barrier : !llvm.ptr<3>, %ticks : i32, %phase : i32) {33 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att 34 // CHECK-SAME: "{35 // CHECK-SAME: .reg .pred P1;36 // CHECK-SAME: LAB_WAIT: 37 // CHECK-SAME: mbarrier.try_wait.parity.shared.b64 P1, [$0], $1, $2;38 // CHECK-SAME: @P1 bra.uni DONE;39 // CHECK-SAME: bra.uni LAB_WAIT;40 // CHECK-SAME: DONE:41 // CHECK-SAME: }",42 // CHECK-SAME: "r,r,r"43 nvvm.mbarrier.try_wait.parity %barrier, %phase, %ticks : !llvm.ptr<3>, i32, i3244 llvm.return45}46 47// CHECK-LABEL: @init_mbarrier_try_wait48llvm.func @init_mbarrier_try_wait(%barrier : !llvm.ptr, %ticks : i32, %phase : i32){49 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att50 // CHECK-SAME: "{51 // CHECK-SAME: .reg .pred P1;52 // CHECK-SAME: LAB_WAIT: 53 // CHECK-SAME: mbarrier.try_wait.parity.b64 P1, [$0], $1, $2;54 // CHECK-SAME: @P1 bra.uni DONE;55 // CHECK-SAME: bra.uni LAB_WAIT;56 // CHECK-SAME: DONE:57 // CHECK-SAME: }",58 // CHECK-SAME: "l,r,r"59 nvvm.mbarrier.try_wait.parity %barrier, %phase, %ticks : !llvm.ptr, i32, i3260 llvm.return61}62 63// CHECK-LABEL: @async_cp64func.func @async_cp(%dst: !llvm.ptr<3>, %src: !llvm.ptr<1>) {65 // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = ca : !llvm.ptr<3>, !llvm.ptr<1>66 nvvm.cp.async.shared.global %dst, %src, 16, cache = ca : !llvm.ptr<3>, !llvm.ptr<1>67 // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = cg : !llvm.ptr<3>, !llvm.ptr<1>68 nvvm.cp.async.shared.global %dst, %src, 16, cache = cg : !llvm.ptr<3>, !llvm.ptr<1>69 return70}71 72// CHECK-LABEL: @async_cp_zfill73func.func @async_cp_zfill(%dst: !llvm.ptr<3>, %src: !llvm.ptr<1>, %cpSize: i32) {74 // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = cg, %{{.*}} : !llvm.ptr<3>, !llvm.ptr<1>, i3275 nvvm.cp.async.shared.global %dst, %src, 16, cache = cg, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i3276 // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 4, cache = ca, %{{.*}} : !llvm.ptr<3>, !llvm.ptr<1>, i3277 nvvm.cp.async.shared.global %dst, %src, 4, cache = ca, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i3278 return79}80 81// CHECK-LABEL: @cp_async_mbarrier_arrive82func.func @cp_async_mbarrier_arrive(%bar_shared: !llvm.ptr<3>, %bar_gen: !llvm.ptr) {83 // CHECK: nvvm.cp.async.mbarrier.arrive %{{.*}}84 nvvm.cp.async.mbarrier.arrive %bar_gen : !llvm.ptr85 // CHECK: nvvm.cp.async.mbarrier.arrive %{{.*}} {noinc = true}86 nvvm.cp.async.mbarrier.arrive %bar_gen {noinc = true} : !llvm.ptr87 // CHECK: nvvm.cp.async.mbarrier.arrive %{{.*}}88 nvvm.cp.async.mbarrier.arrive %bar_shared : !llvm.ptr<3>89 // CHECK: nvvm.cp.async.mbarrier.arrive %{{.*}} {noinc = true}90 nvvm.cp.async.mbarrier.arrive %bar_shared {noinc = true} : !llvm.ptr<3>91 llvm.return92}93 94// CHECK-LABEL: @tma_load_3d_all95func.func @tma_load_3d_all(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %off0: i16, %off1: i16, %ctamask : i16, %cacheHint : i64, %p : i1) {96 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$9 cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.im2col.multicast::cluster.L2::cache_hint [$0], [$1, {$2,$3,$4} ], [$5],{$6}, $7, $8;", "l,l,r,r,r,r,h,h,l,b"97 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2] im2col[%off0] multicast_mask = %ctamask l2_cache_hint = %cacheHint predicate = %p {mode = #nvvm.tma_load_mode<im2col>} : !llvm.ptr<7>, !llvm.ptr98 return99}100 101// CHECK-LABEL: @tma_load_4d_all102func.func @tma_load_4d_all(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %off0: i16, %off1: i16, %ctamask : i16, %cacheHint : i64, %p : i1) {103 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$11 cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.im2col.multicast::cluster.L2::cache_hint [$0], [$1, {$2,$3,$4,$5} ], [$6],{$7,$8}, $9, $10;", "l,l,r,r,r,r,r,h,h,h,l,b"104 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2,%crd3] im2col[%off0,%off1] multicast_mask = %ctamask l2_cache_hint = %cacheHint predicate = %p {mode = #nvvm.tma_load_mode<im2col>} : !llvm.ptr<7>, !llvm.ptr105 return106}107 108// CHECK-LABEL: @tma_load_5d_all109func.func @tma_load_5d_all(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32, %off0: i16, %off1: i16, %off2: i16, %ctamask : i16, %cacheHint : i64, %p : i1) {110 // CHECK: lvm.inline_asm has_side_effects asm_dialect = att "@$13 cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.im2col.multicast::cluster.L2::cache_hint [$0], [$1, {$2,$3,$4,$5,$6} ], [$7],{$8,$9,$10}, $11, $12;", "l,l,r,r,r,r,r,r,h,h,h,h,l,b"111 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2,%crd3,%crd4] im2col[%off0,%off1,%off2] multicast_mask = %ctamask l2_cache_hint = %cacheHint predicate = %p {mode = #nvvm.tma_load_mode<im2col>} : !llvm.ptr<7>, !llvm.ptr112 return113}114 115// CHECK-LABEL: @tma_load_1d116func.func @tma_load_1d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %p : i1) {117 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$4 cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$2} ], [$3];", "l,l,r,r,b"118 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0] predicate=%p : !llvm.ptr<7>, !llvm.ptr119 return120}121 122// CHECK-LABEL: @tma_load_2d123func.func @tma_load_2d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %p : i1) {124 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$5 cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$2,$3} ], [$4];", "l,l,r,r,r,b"125 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1] predicate=%p : !llvm.ptr<7>, !llvm.ptr126 return127}128 129// CHECK-LABEL: @tma_load_3d130func.func @tma_load_3d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %p : i1) {131 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$6 cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$2,$3,$4} ], [$5];", "l,l,r,r,r,r,b"132 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2] predicate=%p : !llvm.ptr<7>, !llvm.ptr133 return134}135 136// CHECK-LABEL: @tma_load_4d137func.func @tma_load_4d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %p : i1) {138 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$7 cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$2,$3,$4,$5} ], [$6];", "l,l,r,r,r,r,r,b"139 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2,%crd3] predicate=%p : !llvm.ptr<7>, !llvm.ptr140 return141}142 143// CHECK-LABEL: @tma_load_5d144func.func @tma_load_5d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32, %p : i1) {145 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$8 cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$2,$3,$4,$5,$6} ], [$7];", "l,l,r,r,r,r,r,r,b"146 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box[%crd0,%crd1,%crd2,%crd3,%crd4] predicate=%p : !llvm.ptr<7>, !llvm.ptr147 return148}149 150// CHECK-LABEL: @tma_load_multicast1d151func.func @tma_load_multicast1d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %multicastMask : i16, %crd0: i32, %p : i1) {152 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$5 cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [$0], [$1, {$2} ], [$3], $4;", "l,l,r,r,h,b"153 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box [%crd0] multicast_mask = %multicastMask predicate=%p : !llvm.ptr<7>, !llvm.ptr154 return155}156 157// CHECK-LABEL: @tma_load_multicast2d158func.func @tma_load_multicast2d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %multicastMask : i16, %crd0: i32, %crd1: i32, %p : i1) {159 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$6 cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [$0], [$1, {$2,$3} ], [$4], $5;", "l,l,r,r,r,h,b"160 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box [%crd0,%crd1] multicast_mask = %multicastMask predicate=%p : !llvm.ptr<7>, !llvm.ptr161 return162}163 164// CHECK-LABEL: @tma_load_multicast3d165func.func @tma_load_multicast3d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %multicastMask : i16, %crd0: i32, %crd1: i32, %crd2: i32, %p : i1) {166 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$7 cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [$0], [$1, {$2,$3,$4} ], [$5], $6;", "l,l,r,r,r,r,h,b"167 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box [%crd0,%crd1,%crd2] multicast_mask = %multicastMask predicate=%p : !llvm.ptr<7>, !llvm.ptr168 return169}170 171// CHECK-LABEL: @tma_load_multicast4d172func.func @tma_load_multicast4d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %multicastMask : i16, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %p : i1) {173 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$8 cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [$0], [$1, {$2,$3,$4,$5} ], [$6], $7;", "l,l,r,r,r,r,r,h,b"174 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box [%crd0,%crd1,%crd2,%crd3] multicast_mask = %multicastMask predicate=%p : !llvm.ptr<7>, !llvm.ptr175 return176}177 178// CHECK-LABEL: @tma_load_multicast5d179func.func @tma_load_multicast5d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<7>, %barrier: !llvm.ptr<3>, %multicastMask : i16, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32, %p : i1) {180 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$9 cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster [$0], [$1, {$2,$3,$4,$5,$6} ], [$7], $8;", "l,l,r,r,r,r,r,r,h,b"181 nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor, %barrier, box [%crd0,%crd1,%crd2,%crd3,%crd4] multicast_mask = %multicastMask predicate=%p : !llvm.ptr<7>, !llvm.ptr182 return183}184 185// CHECK-LABEL: @tma_store_1d186func.func @tma_store_1d(%tmaDescriptor: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %p : i1) {187 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$3 cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [$0, {$2} ], [$1];", "l,r,r,b"188 nvvm.cp.async.bulk.tensor.global.shared.cta %tmaDescriptor, %src, box[%crd0], predicate=%p : !llvm.ptr, !llvm.ptr<3>189 return190}191 192// CHECK-LABEL: @tma_store_2d193func.func @tma_store_2d(%tmaDescriptor: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %p : i1) {194 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$4 cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [$0, {$2, $3} ], [$1];", "l,r,r,r,b"195 nvvm.cp.async.bulk.tensor.global.shared.cta %tmaDescriptor, %src, box[%crd0,%crd1], predicate=%p : !llvm.ptr, !llvm.ptr<3>196 return197}198 199// CHECK-LABEL: @tma_store_3d200func.func @tma_store_3d(%tmaDescriptor: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %p : i1) {201 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$5 cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [$0, {$2, $3, $4} ], [$1];", "l,r,r,r,r,b"202 nvvm.cp.async.bulk.tensor.global.shared.cta %tmaDescriptor, %src, box[%crd0,%crd1,%crd2], predicate=%p : !llvm.ptr, !llvm.ptr<3>203 return204}205 206// CHECK-LABEL: @tma_store_4d207func.func @tma_store_4d(%tmaDescriptor: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %p : i1) {208 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$6 cp.async.bulk.tensor.4d.global.shared::cta.bulk_group [$0, {$2, $3, $4, $5} ], [$1];", "l,r,r,r,r,r,b"209 nvvm.cp.async.bulk.tensor.global.shared.cta %tmaDescriptor, %src, box[%crd0,%crd1,%crd2,%crd3], predicate=%p : !llvm.ptr, !llvm.ptr<3>210 return211}212 213// CHECK-LABEL: @tma_store_5d214func.func @tma_store_5d(%tmaDescriptor: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32, %p : i1) {215 // CHECK-NEXT: llvm.inline_asm has_side_effects asm_dialect = att "@$7 cp.async.bulk.tensor.5d.global.shared::cta.bulk_group [$0, {$2, $3, $4, $5, $6} ], [$1];", "l,r,r,r,r,r,r,b"216 nvvm.cp.async.bulk.tensor.global.shared.cta %tmaDescriptor, %src, box[%crd0,%crd1,%crd2,%crd3,%crd4], predicate=%p : !llvm.ptr, !llvm.ptr<3>217 return218}219 220// CHECK-LABEL: @wgmma_execute221func.func @wgmma_execute() { 222 nvvm.wgmma.fence.aligned223 nvvm.wgmma.commit.group.sync.aligned224 nvvm.wgmma.wait.group.sync.aligned 0225 // CHECK: nvvm.wgmma.fence.aligned226 // CHECK: nvvm.wgmma.commit.group.sync.aligned227 // CHECK: nvvm.wgmma.wait.group.sync.aligned 0228 229 230 nvvm.wgmma.fence.aligned231 nvvm.wgmma.commit.group.sync.aligned232 nvvm.wgmma.wait.group.sync.aligned 5233 // CHECK: nvvm.wgmma.fence.aligned234 // CHECK: nvvm.wgmma.commit.group.sync.aligned235 // CHECK: nvvm.wgmma.wait.group.sync.aligned 5236 return237}238 239 240// -----241 242!mat64f32 = !llvm.struct<(243 f32, f32, f32, f32, f32, f32, f32, f32, 244 f32, f32, f32, f32, f32, f32, f32, f32)>245 246// CHECK-LABEL: @wgmma_f32_f16_f16(247// CHECK-SAME: %[[ARG0:.+]]: i64, %[[ARG1:.+]]: i64248func.func @wgmma_f32_f16_f16(%descA : i64, %descB : i64) -> !mat64f32{ 249 // CHECK: %[[RES:.*]] = llvm.mlir.undef : !llvm.struct250 // CHECK: %[[A1:.*]] = llvm.mlir.constant(0 : i32) : i32251 // CHECK: %[[A2:.*]] = llvm.mlir.constant(-1 : i32) : i32252 // CHECK: %[[A3:.*]] = llvm.mlir.constant(-1 : i32) : i32253 // CHECK: %[[A4:.*]] = llvm.mlir.constant(1 : i32) : i32254 // CHECK: %[[A5:.*]] = llvm.mlir.constant(0 : i32) : i32255 // CHECK: %[[V0:.*]] = llvm.extractvalue %[[RES]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 256 // CHECK: %[[V4:.*]] = llvm.extractvalue %[[RES]][4] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 257 // CHECK: %[[V11:.*]] = llvm.extractvalue %[[RES]][11] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 258 // CHECK: %[[V13:.*]] = llvm.extractvalue %[[RES]][13] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 259 // CHECK: %[[RES1:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 260 // CHECK-SAME: "{261 // CHECK-SAME: reg .pred p;262 // CHECK-SAME: setp.ne.b32 p, $34, 0;263 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 264 // CHECK-SAME: {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15}, $32, $33, p, $35, $36, $37, $38;\0A}\0A", 265 // CHECK-SAME: "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,l,l,n,n,n,n,n" 266 // CHECK-SAME: %[[V0]], %{{.*}}, %{{.*}}, %{{.*}}, %[[V4]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[V11]], %{{.*}}, %[[V13]], %{{.*}}, %{{.*}}, %[[ARG0]], %[[ARG1]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]] 267 // CHECK-SAME: : (f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, i64, i64, i32, i32, i32, i32, i32) 268 // CHECK-SAME: -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>269 // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i64) : i64270 // CHECK: %[[DESCa:.+]] = llvm.add %[[ARG0]], %[[C2]] : i64271 // CHECK: %[[DESCb:.+]] = llvm.add %[[ARG1]], %[[C2]] : i64272 // CHECK: %[[V0_2:.*]] = llvm.extractvalue %[[RES1]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 273 // CHECK: %[[V4_2:.*]] = llvm.extractvalue %[[RES1]][4] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 274 // CHECK: %[[V11_2:.*]] = llvm.extractvalue %[[RES1]][11] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 275 // CHECK: %[[V13_2:.*]] = llvm.extractvalue %[[RES1]][13] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> 276 // CHECK: %[[RES_2:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 277 // CHECK-SAME: "{278 // CHECK-SAME: .reg .pred p;279 // CHECK-SAME: setp.ne.b32 p, $34, 0;280 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 281 // CHECK-SAME: {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15}, $32, $33, p, $35, $36, $37, $38;\0A}\0A", 282 // CHECK-SAME: "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,l,l,n,n,n,n,n" 283 // CHECK-SAME: %[[V0_2]], %{{.*}}, %{{.*}}, %{{.*}}, %[[V4_2]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[V11_2]], %{{.*}}, %[[V13_2]], %{{.*}}, %{{.*}}, %[[DESCa]], %[[DESCb]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} 284 %result = llvm.mlir.undef : !mat64f32285 %result1 = nvvm.wgmma.mma_async 286 %descA, %descB, %result,287 #nvvm.shape<m = 64, n = 32, k = 16>, 288 D [<f32>, #nvvm.wgmma_scale_out<zero>],289 A [<f16>, #nvvm.wgmma_scale_in<neg>, <col>], 290 B [<f16>, #nvvm.wgmma_scale_in<neg>, <col>]291 :!mat64f32 -> !mat64f32292 %c2 = arith.constant 2 : i64293 %descAnext = arith.addi %descA, %c2 : i64294 %descBnext = arith.addi %descB, %c2 : i64295 %result2 = nvvm.wgmma.mma_async 296 %descAnext, %descBnext, %result1,297 #nvvm.shape<m = 64, n = 32, k = 16>, 298 D [<f32>, #nvvm.wgmma_scale_out<zero>],299 A [<f16>, #nvvm.wgmma_scale_in<neg>, <col>], 300 B [<f16>, #nvvm.wgmma_scale_in<neg>, <col>]301 : !mat64f32 -> !mat64f32302 return %result2 : !mat64f32303}304 305// -----306 307!mat16i32 = !llvm.struct<(i32, i32, i32, i32)>308 309// CHECK-LABEL: @wgmma_s32_s8_s8_satfinite(310// CHECK-SAME: %[[ARG0:.+]]: i64, %[[ARG1:.+]]: i64311func.func @wgmma_s32_s8_s8_satfinite(%descA : i64, %descB : i64) -> !mat16i32{ 312 %result = llvm.mlir.undef : !mat16i32313// CHECK: %[[RES:.*]] = llvm.mlir.undef : !llvm.struct314// CHECK: %[[A1:.*]] = llvm.mlir.constant(1 : i32) : i32315// CHECK: %[[V0:.*]] = llvm.extractvalue %[[RES]][0]316// CHECK: %[[V1:.*]] = llvm.extractvalue %[[RES]][1]317// CHECK: %[[V2:.*]] = llvm.extractvalue %[[RES]][2]318// CHECK: %[[V3:.*]] = llvm.extractvalue %[[RES]][3]319// CHECK: %[[RES_2:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 320// CHECK-SAME: "{321// CHECK-SAME: .reg .pred p;322// CHECK-SAME: setp.ne.b32 p, $10, 0;323// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8.satfinite 324// CHECK-SAME: {$0, $1, $2, $3}, $8, $9, p;\0A}\0A", "=r,=r,=r,=r,0,1,2,3,l,l,n" 325// CHECK-SAME: %[[V0]], %[[V1]], %[[V2]], %[[V3]], %[[ARG0]], %[[ARG1]], %[[A1]] : 326// CHECK-SAME: (i32, i32, i32, i32, i64, i64, i32) -> !llvm.struct<(i32, i32, i32, i32)>327// CHECK: %[[V0_2:.*]] = llvm.extractvalue %[[RES_2]][0]328// CHECK: %[[V1_2:.*]] = llvm.extractvalue %[[RES_2]][1]329// CHECK: %[[V2_2:.*]] = llvm.extractvalue %[[RES_2]][2]330// CHECK: %[[V3_2:.*]] = llvm.extractvalue %[[RES_2]][3]331// CHECK: %[[RES_3:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 332// CHECK-SAME: "{333// CHECK-SAME: .reg .pred p;334// CHECK-SAME: setp.ne.b32 p, $10, 0;335// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8.satfinite 336// CHECK-SAME: {$0, $1, $2, $3}, $8, $9, p;\0A}\0A", 337// CHECK-SAME: "=r,=r,=r,=r,0,1,2,3,l,l,n" 338// CHECK-SAME: %[[V0_2]], %[[V1_2]], %[[V2_2]], %[[V3_2]], %[[ARG0]], %[[ARG1]], %{{.*}}339// CHECK: %[[V0_3:.*]] = llvm.extractvalue %[[RES_3]][0]340// CHECK: %[[V1_3:.*]] = llvm.extractvalue %[[RES_3]][1]341// CHECK: %[[V2_3:.*]] = llvm.extractvalue %[[RES_3]][2]342// CHECK: %[[V3_3:.*]] = llvm.extractvalue %[[RES_3]][3]343// CHECK: %[[RES1:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 344// CHECK-SAME:"{345// CHECK-SAME:.reg .pred p;346// CHECK-SAME: setp.ne.b32 p, $10, 0;347// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8.satfinite348// CHECK-SAME: {$0, $1, $2, $3}, $8, $9, p;\0A}\0A", "=r,=r,=r,=r,0,1,2,3,l,l,n" 349// CHECK-SAME: %[[V0_3]], %[[V1_3]], %[[V2_3]], %[[V3_3]], %[[ARG0]], %[[ARG1]], %{{.*}} 350 %result1 = nvvm.wgmma.mma_async %descA, %descB, %result, 351 #nvvm.shape<m = 64, n = 8, k = 32>, 352 D [<s32>, #nvvm.wgmma_scale_out<one>, <satfinite>],353 A [<s8>, #nvvm.wgmma_scale_in<one>, <row>], 354 B [<s8>, #nvvm.wgmma_scale_in<one>, <col>]355 : !mat16i32 -> !mat16i32356 %result2 = nvvm.wgmma.mma_async %descA, %descB, %result1, 357 #nvvm.shape<m = 64, n = 8, k = 32>, 358 D [<s32>, #nvvm.wgmma_scale_out<one>, <satfinite>],359 A [<s8>, #nvvm.wgmma_scale_in<one>, <row>], 360 B [<s8>, #nvvm.wgmma_scale_in<one>, <col>]361 : !mat16i32 -> !mat16i32362 %result3 = nvvm.wgmma.mma_async %descA, %descB, %result2, 363 #nvvm.shape<m = 64, n = 8, k = 32>, 364 D [<s32>, #nvvm.wgmma_scale_out<one>, <satfinite>],365 A [<s8>, #nvvm.wgmma_scale_in<one>, <row>], 366 B [<s8>, #nvvm.wgmma_scale_in<one>, <col>]367 : !mat16i32 -> !mat16i32368 return %result3 : !mat16i32369}370 371// CHECK-LABEL: @wgmma_s32_u8_u8(372 // CHECK-SAME: %[[ARG0:.+]]: i64, %[[ARG1:.+]]: i64373func.func @wgmma_s32_u8_u8(%descA : i64, %descB : i64) -> !mat16i32 { 374// CHECK: %[[RES:.*]] = llvm.mlir.undef : !llvm.struct375// CHECK: %[[A1:.*]] = llvm.mlir.constant(1 : i32) : i32376// CHECK: %[[V0:.*]] = llvm.extractvalue %[[RES]][0]377// CHECK: %[[V1:.*]] = llvm.extractvalue %[[RES]][1]378// CHECK: %[[V2:.*]] = llvm.extractvalue %[[RES]][2]379// CHECK: %[[V3:.*]] = llvm.extractvalue %[[RES]][3]380// CHECK: %[[RES_2:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 381// CHECK-SAME: "{382// CHECK-SAME: .reg .pred p;383// CHECK-SAME: setp.ne.b32 p, $10, 0;384// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8 {$0, $1, $2, $3}, $8, $9, p;385// CHECK-SAME: }\0A",386// CHECK-SAME: "=r,=r,=r,=r,0,1,2,3,l,l,n" %[[V0]], %[[V1]], %[[V2]], %[[V3]], %[[ARG0]], %[[ARG1]], %[[A1]] : 387// CHECK-SAME:(i32, i32, i32, i32, i64, i64, i32) -> !llvm.struct<(i32, i32, i32, i32)>388// CHECK: %[[V0_2:.*]] = llvm.extractvalue %[[RES_2]][0]389// CHECK: %[[V1_2:.*]] = llvm.extractvalue %[[RES_2]][1]390// CHECK: %[[V2_2:.*]] = llvm.extractvalue %[[RES_2]][2]391// CHECK: %[[V3_2:.*]] = llvm.extractvalue %[[RES_2]][3]392// CHECK: %[[RES_3:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 393// CHECK-SAME:"{394// CHECK-SAME: .reg .pred p;395// CHECK-SAME: setp.ne.b32 p, $10, 0;396// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8 {$0, $1, $2, $3}, $8, $9, p;397// CHECK-SAME: }\0A",398// CHECK-SAME: "=r,=r,=r,=r,0,1,2,3,l,l,n" %[[V0_2]], %[[V1_2]], %[[V2_2]], %[[V3_2]], %[[ARG0]], %[[ARG1]], %{{.*}}399// CHECK: %[[V0_3:.*]] = llvm.extractvalue %[[RES_3]][0]400// CHECK: %[[V1_3:.*]] = llvm.extractvalue %[[RES_3]][1]401// CHECK: %[[V2_3:.*]] = llvm.extractvalue %[[RES_3]][2]402// CHECK: %[[V3_3:.*]] = llvm.extractvalue %[[RES_3]][3]403// CHECK: %[[RES1:.*]] = llvm.inline_asm has_side_effects asm_dialect = att 404// CHECK-SAME:"{405// CHECK-SAME: .reg .pred p;406// CHECK-SAME: setp.ne.b32 p, $10, 0;407// CHECK-SAME: wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8 {$0, $1, $2, $3}, $8, $9, p;408// CHECK-SAME:}\0A", 409// CHECK-SAME:"=r,=r,=r,=r,0,1,2,3,l,l,n" %[[V0_3]], %[[V1_3]], %[[V2_3]], %[[V3_3]], %[[ARG0]], %[[ARG1]], %{{.*}} 410 %result = llvm.mlir.undef : !mat16i32411 %result1 = nvvm.wgmma.mma_async %descA, %descB, %result,412 #nvvm.shape<m = 64, n = 8, k = 32>, 413 D [<s32>, #nvvm.wgmma_scale_out<one>],414 A [<u8>, #nvvm.wgmma_scale_in<one>, <row>], 415 B [<u8>, #nvvm.wgmma_scale_in<one>, <col>]416 : !mat16i32 -> !mat16i32417 %result2 = nvvm.wgmma.mma_async %descA, %descB, %result1,418 #nvvm.shape<m = 64, n = 8, k = 32>, 419 D [<s32>, #nvvm.wgmma_scale_out<one>],420 A [<u8>, #nvvm.wgmma_scale_in<one>, <row>], 421 B [<u8>, #nvvm.wgmma_scale_in<one>, <col>]422 : !mat16i32 -> !mat16i32423 %result3 = nvvm.wgmma.mma_async %descA, %descB, %result2,424 #nvvm.shape<m = 64, n = 8, k = 32>, 425 D [<s32>, #nvvm.wgmma_scale_out<one>],426 A [<u8>, #nvvm.wgmma_scale_in<one>, <row>], 427 B [<u8>, #nvvm.wgmma_scale_in<one>, <col>]428 : !mat16i32 -> !mat16i32429 return %result3 : !mat16i32430}431 432// -----433 434!mat32f32 = !llvm.struct<(435 f32, f32, f32, f32, f32, f32, f32, f32, 436 f32, f32, f32, f32, f32, f32, f32, f32, 437 f32, f32, f32, f32, f32, f32, f32, f32, 438 f32, f32, f32, f32, f32, f32, f32, f32)>439 440// CHECK-LABEL: @wgmma_f32_tf32_tf32441func.func @wgmma_f32_tf32_tf32(%descA : i64, %descB : i64) -> !mat32f32 { 442 // CHECK: %[[RES:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 443 // CHECK-SAME:"{444 // CHECK-SAME: .reg .pred p;445 // CHECK-SAME: setp.ne.b32 p, $66, 0;446 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k8.f32.tf32.tf32 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"447 // CHECK: %[[RES_2:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 448 // CHECK-SAME: "{449 // CHECK-SAME: .reg .pred p;450 // CHECK-SAME: setp.ne.b32 p, $66, 0;451 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k8.f32.tf32.tf32 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"452 %result = llvm.mlir.undef : !mat32f32453 %result1 = nvvm.wgmma.mma_async %descA, %descB, %result,454 #nvvm.shape<m = 64, n = 64, k = 8>, 455 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],456 A [#nvvm.wgmma_type<tf32>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 457 B [#nvvm.wgmma_type<tf32>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]458 : !mat32f32 -> !mat32f32459 %result2 = nvvm.wgmma.mma_async %descA, %descB, %result1,460 #nvvm.shape<m = 64, n = 64, k = 8>, 461 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],462 A [#nvvm.wgmma_type<tf32>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 463 B [#nvvm.wgmma_type<tf32>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]464 : !mat32f32 -> !mat32f32465 return %result2 : !mat32f32466}467 468 469// -----470 471!mat32f32 = !llvm.struct<(472 f32, f32, f32, f32, f32, f32, f32, f32, 473 f32, f32, f32, f32, f32, f32, f32, f32, 474 f32, f32, f32, f32, f32, f32, f32, f32, 475 f32, f32, f32, f32, f32, f32, f32, f32)>476 477// CHECK-LABEL: @wgmma_f32_e4m3_e4m3478func.func @wgmma_f32_e4m3_e4m3(%descA : i64, %descB : i64) -> !mat32f32 { 479 // CHECK: %[[RES:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 480 // CHECK-SAME: "{\0A.reg .pred p;\0Asetp.ne.b32 p, $66, 0;481 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"482 // CHECK: %[[RES_2:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 483 // CHECK-SAME: "{\0A.reg .pred p;\0Asetp.ne.b32 p, $66, 0;484 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"485 %result = llvm.mlir.undef : !mat32f32486 %result1 = nvvm.wgmma.mma_async %descA, %descB, %result,487 #nvvm.shape<m = 64, n = 64, k = 32>, 488 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],489 A [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 490 B [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]491 : !mat32f32 -> !mat32f32492 %result2 = nvvm.wgmma.mma_async %descA, %descB, %result1,493 #nvvm.shape<m = 64, n = 64, k = 32>, 494 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],495 A [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 496 B [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]497 : !mat32f32 -> !mat32f32498 return %result2 : !mat32f32499}500 501// -----502 503!mat32f32 = !llvm.struct<(504 f32, f32, f32, f32, f32, f32, f32, f32, 505 f32, f32, f32, f32, f32, f32, f32, f32, 506 f32, f32, f32, f32, f32, f32, f32, f32, 507 f32, f32, f32, f32, f32, f32, f32, f32)>508 509// CHECK-LABEL: @wgmma_f32_e5m2_e4m3510func.func @wgmma_f32_e5m2_e4m3(%descA : i64, %descB : i64) -> !mat32f32 { 511 // CHECK: %[[RES:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 512 // CHECK-SAME: "{\0A.reg .pred p;\0Asetp.ne.b32 p, $66, 0;513 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e4m3 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"514 // CHECK: %[[RES_2:.+]] = llvm.inline_asm has_side_effects asm_dialect = att 515 // CHECK-SAME: "{\0A.reg .pred p;\0Asetp.ne.b32 p, $66, 0;516 // CHECK-SAME: wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e4m3 {$0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31}, $64, $65, p, $67, $68;\0A}\0A", "=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,=f,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,l,l,n,n,n"517 %result = llvm.mlir.undef : !mat32f32518 %result1 = nvvm.wgmma.mma_async %descA, %descB, %result,519 #nvvm.shape<m = 64, n = 64, k = 32>, 520 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],521 A [#nvvm.wgmma_type<e5m2>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 522 B [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]523 : !mat32f32 -> !mat32f32524 %result2 = nvvm.wgmma.mma_async %descA, %descB, %result1,525 #nvvm.shape<m = 64, n = 64, k = 32>, 526 D [#nvvm.wgmma_type<f32>, #nvvm.wgmma_scale_out<one>],527 A [#nvvm.wgmma_type<e5m2>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<row>], 528 B [#nvvm.wgmma_type<e4m3>, #nvvm.wgmma_scale_in<one>, #nvvm.mma_layout<col>]529 : !mat32f32 -> !mat32f32530 return %result2 : !mat32f32531}532 533// -----534 535func.func @elect_one_leader_sync() { 536 // CHECK: %[[RES:.*]] = nvvm.elect.sync -> i1537 %cnd = nvvm.elect.sync -> i1 538 return 539}540 541// -----542 543// CHECK-LABEL: @test_nvvm_prefetch544llvm.func @test_nvvm_prefetch(%desc : !llvm.ptr, %pred : i1) {545 //CHECK: nvvm.prefetch tensormap, %{{.*}}546 nvvm.prefetch tensormap, %desc : !llvm.ptr547 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$1 prefetch.tensormap [$0];", "l,b"548 nvvm.prefetch tensormap, %desc, predicate = %pred : !llvm.ptr, i1549 llvm.return550}551 552// -----553 554func.func @set_max_register() {555 // CHECK: nvvm.setmaxregister increase 232556 nvvm.setmaxregister increase 232557 558 // CHECK: nvvm.setmaxregister decrease 40559 nvvm.setmaxregister decrease 40560 func.return561}562 563// -----564 565func.func @cp_async_bulk_commit() {566 // CHECK: nvvm.cp.async.bulk.commit.group567 nvvm.cp.async.bulk.commit.group568 func.return569}570 571// -----572 573func.func @cp_async_bulk_wait_group() {574 // CHECK: nvvm.cp.async.bulk.wait_group 1575 // CHECK: nvvm.cp.async.bulk.wait_group 0576 // CHECK: nvvm.cp.async.bulk.wait_group 5 {read}577 // CHECK: nvvm.cp.async.bulk.wait_group 0 {read}578 nvvm.cp.async.bulk.wait_group 1579 nvvm.cp.async.bulk.wait_group 0580 nvvm.cp.async.bulk.wait_group 5 {read}581 nvvm.cp.async.bulk.wait_group 0 {read}582 func.return583}584 585// -----586 587func.func @fence_mbarrier_init() {588 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.mbarrier_init.release.cluster;"589 nvvm.fence.mbarrier.init590 func.return 591}592// -----593 594func.func @fence_proxy() {595 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.proxy.alias;", "" : () -> ()596 nvvm.fence.proxy { kind = #nvvm.proxy_kind<alias>}597 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.proxy.async;", "" : () -> ()598 nvvm.fence.proxy { kind = #nvvm.proxy_kind<async>}599 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.proxy.async.global;", "" : () -> ()600 nvvm.fence.proxy { kind = #nvvm.proxy_kind<async.global>}601 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.proxy.async.shared::cta;", "" : () -> ()602 nvvm.fence.proxy { kind = #nvvm.proxy_kind<async.shared>, space = #nvvm.shared_space<cta>}603 //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "fence.proxy.async.shared::cluster;", "" : () -> ()604 nvvm.fence.proxy { kind = #nvvm.proxy_kind<async.shared>, space = #nvvm.shared_space<cluster>}605 func.return606}607 608// -----609 610// CHECK-LABEL: @llvm_nvvm_barrier_arrive611// CHECK-SAME: (%[[barId:.*]]: i32, %[[numberOfThreads:.*]]: i32)612llvm.func @llvm_nvvm_barrier_arrive(%barID : i32, %numberOfThreads : i32) {613 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "bar.arrive 0, $0;", "r" %[[numberOfThreads]] : (i32) -> ()614 nvvm.barrier.arrive number_of_threads = %numberOfThreads615 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "bar.arrive $0, $1;", "r,r" %[[barId]], %[[numberOfThreads]] : (i32, i32) -> ()616 nvvm.barrier.arrive id = %barID number_of_threads = %numberOfThreads617 llvm.return618}619 620 621// -----622 623llvm.func @init_mbarrier(624 %barrier_gen : !llvm.ptr, 625 %barrier : !llvm.ptr<3>, 626 %count : i32, 627 %pred : i1) {628 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.init.b64 [$0], $1;", "l,r" 629 nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32)630 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b"631 nvvm.inline_ptx "mbarrier.init.b64 [{$r0}], {$r1};" ro (%barrier_gen, %count : !llvm.ptr, i32), predicate = %pred632 llvm.return633}634// -----635 636llvm.func @ex2(%input : f32, %pred : i1) {637 // CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "ex2.approx.ftz.f32 $0, $1;", "=f,f" %{{.*}} : (f32) -> f32638 %0 = nvvm.inline_ptx "ex2.approx.ftz.f32 {$w0}, {$r0};" ro (%input : f32) -> f32639 640 // CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "@$1 ex2.approx.ftz.f32 $0, $1;", "=f,f,b" %{{.*}}, %{{.*}} : (f32, i1) -> f32641 %1 = nvvm.inline_ptx "ex2.approx.ftz.f32 {$w0}, {$r0};" ro (%input : f32), predicate = %pred -> f32642 llvm.return643}644 645// CHECK-LABEL: @multi_return(646// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: i32, %[[arg1:[a-zA-Z0-9_]+]]: i32)647llvm.func @multi_return(%a : i32, %b : i32) -> i32 {648 // CHECK: %[[S1:.+]] = llvm.inline_asm has_side_effects asm_dialect = att "{.reg .pred p; setp.ge.s32 p, $2, $3; selp.s32 $0, $2,$3, p; selp.s32 $1, $2,$3, p;}", "=r,=r,r,r" %[[arg0]], %[[arg1]] : (i32, i32) -> !llvm.struct<(i32, i32)>649 // CHECK: %[[S2:.+]] = llvm.extractvalue %[[S1]][0] : !llvm.struct<(i32, i32)> 650 // CHECK: %[[S3:.+]] = llvm.extractvalue %[[S1]][1] : !llvm.struct<(i32, i32)> 651 // CHECK: %[[S4:.+]] = llvm.add %[[S2]], %[[S3]] : i32652 // CHECK: llvm.return %[[S4]] : i32653 %r1, %r2 = nvvm.inline_ptx "{.reg .pred p; setp.ge.s32 p, {$r0}, {$r1}; selp.s32 {$w0}, {$r0},{$r1}, p; selp.s32 {$w1}, {$r0},{$r1}, p;}"654 ro (%a, %b : i32,i32) -> i32,i32655 %r3 = llvm.add %r1, %r2 : i32656 llvm.return %r3 : i32657}658 659// CHECK-LABEL: @inline_ptx_multi_rw(660// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: i32, %[[arg1:[a-zA-Z0-9_]+]]: i32, %[[arg2:[a-zA-Z0-9_]+]]: f32, %[[arg3:[a-zA-Z0-9_]+]]: f32)661llvm.func @inline_ptx_multi_rw(%a : i32, %b : i32, %rw_c : f32, %rw_d : f32) -> f32 {662// CHECK: %[[S0:.+]] = llvm.inline_asm has_side_effects asm_dialect = att "{.reg .pred p; setp.ge.s32 p, $2, $3; selp.s32 $0, $2,$3, p; selp.s32 $1, $2,$3, p;}", 663// CHECK-SAME: "=f,=f,r,r,0,1" 664// CHECK-SAME: %[[arg2]], %[[arg3]], %[[arg0]], %[[arg1]] 665// CHECK-SAME: : (f32, f32, i32, i32) -> !llvm.struct<(f32, f32)>666// CHECK: %[[S1:.+]] = llvm.extractvalue %[[S0]][0] : !llvm.struct<(f32, f32)> 667// CHECK: %[[S2:.+]] = llvm.extractvalue %[[S0]][1] : !llvm.struct<(f32, f32)> 668// CHECK: %[[S3:.+]] = llvm.fadd %[[S1]], %[[S2]] : f32669// CHECK: llvm.return %[[S3]] : f32670 nvvm.inline_ptx "{.reg .pred p; setp.ge.s32 p, {$r0}, {$r1}; selp.s32 {$rw0}, {$r0},{$r1}, p; selp.s32 {$rw1}, {$r0},{$r1}, p;}"671 ro (%a, %b : i32,i32) 672 rw (%rw_c, %rw_d: f32,f32)673 %r4 = llvm.fadd %rw_c, %rw_d : f32674 llvm.return %r4 : f32675}676 677// CHECK-LABEL: @inline_ptx_multi_rw_r(678// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: i32, %[[arg1:[a-zA-Z0-9_]+]]: i32, %[[arg2:[a-zA-Z0-9_]+]]: f32, %[[arg3:[a-zA-Z0-9_]+]]: f32)679llvm.func @inline_ptx_multi_rw_r(%a : i32, %b : i32, %rw_c : f32, %rw_d : f32) -> f32 {680// CHECK: %[[S0:.+]] = llvm.inline_asm has_side_effects asm_dialect = att "{.reg .pred p; setp.ge.s32 p, $4, $5; selp.s32 $0, $4,$5, p; selp.s32 $1, $4,$5, p; selp.s32 $2, $4,$5, p; selp.s32 $3, $4,$5, p;}", 681// CHECK-SAME: "=f,=f,=r,=r,r,r,0,1" 682// CHECK-SAME: %[[arg2]], %[[arg3]], %[[arg0]], %[[arg1]] : 683// CHECK-SAME: (f32, f32, i32, i32) -> !llvm.struct<(f32, f32, i32, i32)>684// CHECK: %[[S1:.+]] = llvm.extractvalue %[[S0]][0] : !llvm.struct<(f32, f32, i32, i32)> 685// CHECK: %[[S2:.+]] = llvm.extractvalue %[[S0]][1] : !llvm.struct<(f32, f32, i32, i32)> 686// CHECK: %[[S3:.+]] = llvm.extractvalue %[[S0]][2] : !llvm.struct<(f32, f32, i32, i32)> 687// CHECK: %[[S4:.+]] = llvm.extractvalue %[[S0]][3] : !llvm.struct<(f32, f32, i32, i32)> 688// CHECK: %[[S5:.+]] = llvm.add %[[S3]], %[[S4]] : i32689// CHECK: %[[S6:.+]] = llvm.sitofp %[[S5]] : i32 to f32690// CHECK: %[[S7:.+]] = llvm.fadd %[[S1]], %[[S2]] : f32691// CHECK: %[[S8:.+]] = llvm.fadd %[[S6]], %[[S2]] : f32692// CHECK: llvm.return %[[S8]] : f32693 694 %wo0, %wo1 = nvvm.inline_ptx "{.reg .pred p; setp.ge.s32 p, {$r0}, {$r1}; selp.s32 {$rw0}, {$r0},{$r1}, p; selp.s32 {$rw1}, {$r0},{$r1}, p; selp.s32 {$w0}, {$r0},{$r1}, p; selp.s32 {$w1}, {$r0},{$r1}, p;}"695 ro (%a, %b : i32,i32) 696 rw (%rw_c, %rw_d: f32,f32) -> i32,i32697 %r3 = llvm.add %wo0, %wo1 : i32698 %r3f = llvm.sitofp %r3 : i32 to f32699 %r4 = llvm.fadd %rw_c, %rw_d : f32700 %r5 = llvm.fadd %r3f, %rw_d : f32701 llvm.return %r5 : f32702}703 704 705// -----706 707// CHECK-LABEL: @nvvm_pmevent708llvm.func @nvvm_pmevent() {709 // CHECK: %[[S0:.+]] = llvm.mlir.constant(10 : i32) : i32710 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "pmevent $0;", "n" %[[S0]] : (i32) -> ()711 712 nvvm.pmevent id = 10713 // CHECK: %[[S1:.+]] = llvm.mlir.constant(4 : i32) : i32714 // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "pmevent $0;", "n" %[[S1]] : (i32) -> ()715 nvvm.pmevent id = 4716 llvm.return717}718 719// -----720 721llvm.func @inline_ptx_pack_4i8(%src : vector<4xi8>, %mask : i32, %zero: i32) {722// CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "dp4a.s32.s32 $0, $1, $2, $3;", "=r,r,r,r" %{{.*}}, %{{.*}}, %{{.*}} : (vector<4xi8>, i32, i32) -> i32723 %wo0 = nvvm.inline_ptx "dp4a.s32.s32 {$w0}, {$r0}, {$r1}, {$r2};" 724 ro(%src, %mask, %zero : vector<4xi8>, i32, i32) 725 -> i32726 llvm.return 727}728 729llvm.func @inline_ptx_pack_2bf16(%a : f32, %b : f32) {730 // CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "cvt.rn.satfinite.bf16x2.f32 $0, $1, $2;", "=f,f,f" %{{.*}}, %{{.*}} : (f32, f32) -> vector<2xbf16>731 %wo0 = nvvm.inline_ptx "cvt.rn.satfinite.bf16x2.f32 {$w0}, {$r0}, {$r1};" 732 ro(%a, %b : f32, f32) 733 -> vector<2xbf16>734 llvm.return 735}736 737llvm.func @inline_ptx_cvt_rn_e4m3x2_f16x2(%a : i16) {738// CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "cvt.rz.satfinite.ue8m0x2.bf16x2 $0, $1", "=f,h" %{{.*}} : (i16) -> vector<2xbf16>739 %wo0 = nvvm.inline_ptx "cvt.rz.satfinite.ue8m0x2.bf16x2 {$w0}, {$r0}" 740 ro(%a : i16) 741 -> vector<2xbf16>742 llvm.return 743}744 745llvm.func @cvt_i8_bf16(%a : i8) {746 // CHECK: %{{.*}} = llvm.inline_asm has_side_effects asm_dialect = att "{\0A\09.reg .b16 r;\0A\09.reg .b8 s;\0A\09mov.b16 {s,_}, $0;\0A\09cvt.rn.bf16.s8 r, s;\0A\09mov.b16 $1, r;\0A\09", "=h,h" %{{.*}} : (i16) -> i16747 %za = llvm.zext %a : i8 to i16748 %wo0 = nvvm.inline_ptx "{\n\t.reg .b16 r;\n\t.reg .b8 s;\n\tmov.b16 {s,_}, {$w0};\n\tcvt.rn.bf16.s8 r, s;\n\tmov.b16 {$r0}, r;\n\t" 749 ro(%za : i16) 750 -> i16751 llvm.return 752}753