brintos / llvm-project-archived public Read only

1// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s2// RUN: mlir-opt %s -transform-interpreter | FileCheck %s3 4// CHECK-LABEL: @m16n8k16_fp165func.func @m16n8k16_fp16(%arg0: vector<4x2xf16>, %arg1: vector<2x2xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {6  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<2xf16>>7  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<2xf16>>8  // CHECK: llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vector<2xf16>>9  // CHECK: llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vector<2xf16>>10  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>11  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>12  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>13  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>14  // CHECK-NOT: llvm.extractvalue15  // CHECK: [[d:%.+]] = nvvm.mma.sync16  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>17  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>18  // CHECK-DAG: llvm.extractvalue [[d]][0] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>19  // CHECK-DAG: llvm.extractvalue [[d]][1] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>20  // CHECK: llvm.mlir.poison : !llvm.array<2 x vector<2xf16>>21  // CHECK-DAG: llvm.insertvalue {{%.+}}, {{%.+}}[0] : !llvm.array<2 x vector<2xf16>>22  // CHECK-DAG: llvm.insertvalue {{%.+}}, {{%.+}}[1] : !llvm.array<2 x vector<2xf16>>23  return %d : vector<2x2xf16>24}25 26// Same as above but with fp32 acumulation type.27 28// CHECK-LABEL: @m16n8k16_fp16_fp3229func.func @m16n8k16_fp16_fp32(%arg0: vector<4x2xf16>, %arg1: vector<2x2xf16>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {30  // We just need to check the mma instruction and the manipulatin of the result.31  // CHECK: [[d:%.+]] = nvvm.mma.sync32  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>33  // CHECK-SAME: (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>34  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf32>) -> vector<2x2xf32>35  // CHECK: [[undef:%.+]] = llvm.mlir.poison : vector<2xf32>36  // CHECK-DAG: llvm.extractvalue [[d]][0] : !llvm.struct<(f32, f32, f32, f32)>37  // CHECK-DAG: llvm.extractvalue [[d]][1] : !llvm.struct<(f32, f32, f32, f32)>38  // CHECK: [[d00:%.+]] = llvm.insertelement {{%.+}}, [[undef]][{{.*}}] : vector<2xf32>39  // CHECK: [[d01:%.+]] = llvm.insertelement {{%.+}}, [[d00]][{{.*}}] : vector<2xf32>40 41  // CHECK: [[undef:%.+]] = llvm.mlir.poison : vector<2xf32>42  // CHECK-DAG: llvm.extractvalue [[d]][2] : !llvm.struct<(f32, f32, f32, f32)>43  // CHECK-DAG: llvm.extractvalue [[d]][3] : !llvm.struct<(f32, f32, f32, f32)>44  // CHECK: [[d10:%.+]] = llvm.insertelement {{%.+}}, [[undef]][{{.*}}] : vector<2xf32>45  // CHECK: [[d11:%.+]] = llvm.insertelement {{%.+}}, [[d10]][{{.*}}] : vector<2xf32>46 47  // CHECK-DAG: llvm.insertvalue [[d01]], {{%.+}}[0] : !llvm.array<2 x vector<2xf32>>48  // CHECK-DAG: llvm.insertvalue [[d11]], {{%.+}}[1] : !llvm.array<2 x vector<2xf32>>49  return %d : vector<2x2xf32>50}51 52// CHECK-LABEL: @m16n8k8_fp1653func.func @m16n8k8_fp16(%arg0: vector<2x2xf16>, %arg1: vector<1x2xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {54  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>55  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>56  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<1 x vector<2xf16>>57  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>58  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>59  // CHECK-NOT: llvm.extractvalue60  // CHECK: [[d:%.+]] = nvvm.mma.sync61  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 8>62  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 8]} : (vector<2x2xf16>, vector<1x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>63  // CHECK-DAG: llvm.extractvalue [[d]][0] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>64  // CHECK-DAG: llvm.extractvalue [[d]][1] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>65  // CHECK: llvm.mlir.poison : !llvm.array<2 x vector<2xf16>>66  // CHECK-DAG: llvm.insertvalue {{%.+}}, {{%.+}}[0] : !llvm.array<2 x vector<2xf16>>67  // CHECK-DAG: llvm.insertvalue {{%.+}}, {{%.+}}[1] : !llvm.array<2 x vector<2xf16>>68  // CHECK: return69  return %d : vector<2x2xf16>70}71 72// CHECK-LABEL: @m16n8k32_int873func.func @m16n8k32_int8(%arg0: vector<4x4xi8>, %arg1: vector<2x4xi8>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {74  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<4xi8>>75  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3276  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<4xi8>>77  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3278  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<4xi8>>79  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3280  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<4xi8>>81  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3282  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<4xi8>>83  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3284  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<4xi8>>85  // CHECK: llvm.bitcast [[el]] : vector<4xi8> to i3286  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>87  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>88  // CHECK: [[d:%.+]] = nvvm.mma.sync89  // CHECK-SAME: intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>90  // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<s8>91  // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<s8>92  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>93  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<4x4xi8>, vector<2x4xi8>, vector<2x2xi32>) -> vector<2x2xi32>94  return %d : vector<2x2xi32>95}96 97// CHECK-LABEL: @m16n8k32_i498func.func @m16n8k32_i4(%arg0: vector<2x8xi4>, %arg1: vector<1x8xi4>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {99  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<8xi4>>100  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32101  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<8xi4>>102  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32103  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<1 x vector<8xi4>>104  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32105  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>106  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>107  // CHECK: [[d:%.+]] = nvvm.mma.sync108  // CHECK-SAME: intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>109  // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<s4>110  // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<s4>111  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>112  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<2x8xi4>, vector<1x8xi4>, vector<2x2xi32>) -> vector<2x2xi32>113  return %d : vector<2x2xi32>114}115 116// CHECK-LABEL: @m16n8k64_i4117func.func @m16n8k64_i4(%arg0: vector<4x8xi4>, %arg1: vector<2x8xi4>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {118  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<8xi4>>119  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32120  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<8xi4>>121  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32122  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<8xi4>>123  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32124  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<8xi4>>125  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32126  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<8xi4>>127  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32128  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<8xi4>>129  // CHECK: llvm.bitcast [[el]] : vector<8xi4> to i32130  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>131  // CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>132  // CHECK: [[d:%.+]] = nvvm.mma.sync133  // CHECK-SAME: intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>134  // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<s4>135  // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<s4>136  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 64>137  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 64]} : (vector<4x8xi4>, vector<2x8xi4>, vector<2x2xi32>) -> vector<2x2xi32>138  return %d : vector<2x2xi32>139}140 141// CHECK-LABEL: @m8n8k4_f64142func.func @m8n8k4_f64(%arg0: vector<1x1xf64>, %arg1: vector<1x1xf64>, %arg2: vector<1x2xf64>) -> vector<1x2xf64> {143  // CHECK: llvm.extractvalue144  // CHECK: llvm.extractvalue145  // CHECK: llvm.extractvalue146  // CHECK: [[d:%.+]] = nvvm.mma.sync A[{{%.+}}] B[{{%.+}}] C[{{%.+}}, {{%.+}}]147  // CHECK-SAME: shape = #nvvm.shape<m = 8, n = 8, k = 4>148  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [8, 8, 4]} : (vector<1x1xf64>, vector<1x1xf64>, vector<1x2xf64>) -> vector<1x2xf64>149  // CHECK: llvm.mlir.poison : vector<2xf64>150  // CHECK-DAG: llvm.extractvalue [[d]][0] : !llvm.struct<(f64, f64)>151  // CHECK-DAG: llvm.extractvalue [[d]][1] : !llvm.struct<(f64, f64)>152  // CHECK-COUNT-2: llvm.insertelement {{.*}} : vector<2xf64>153  // CHECK-DAG: llvm.insertvalue {{%.+}}, {{%.+}}[0] : !llvm.array<1 x vector<2xf64>>154  // CHECK: return155  return %d : vector<1x2xf64>156}157 158 159// CHECK-LABEL: @ldmatrix_x4160func.func @ldmatrix_x4(%arg0: memref<128x128xf16, 3>) ->  vector<4x2xf16> {161  %c0  = arith.constant 0 : index162  // CHECK: nvvm.ldmatrix {{%.+}} {eltType = #nvvm.ld_st_matrix_elt_type<b16>, layout = #nvvm.mma_layout<row>, num = 4 : i32, shape = #nvvm.ld_st_matrix_shape<m = 8, n = 8>} : {{.*}} -> !llvm.struct<(i32, i32, i32, i32)>163  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf16, 3> -> vector<4x2xf16>164  // CHECK: llvm.extractvalue165  // CHECK: llvm.bitcast166  // CHECK: llvm.insertvalue167  // CHECK: llvm.extractvalue168  // CHECK: llvm.bitcast169  // CHECK: llvm.insertvalue170  // CHECK: llvm.extractvalue171  // CHECK: llvm.bitcast172  // CHECK: llvm.insertvalue173  // CHECK: llvm.extractvalue174  // CHECK: llvm.bitcast175  // CHECK: llvm.insertvalue176  return %a : vector<4x2xf16>177}178 179// CHECK-LABEL: @ldmatrix_x1180func.func @ldmatrix_x1(%arg0: memref<128x128xf16, 3>) ->  vector<1x2xf16> {181  %c0  = arith.constant 0 : index182  // CHECK: nvvm.ldmatrix {{%.+}} {eltType = #nvvm.ld_st_matrix_elt_type<b16>, layout = #nvvm.mma_layout<row>, num = 1 : i32, shape = #nvvm.ld_st_matrix_shape<m = 8, n = 8>} : {{.*}} -> i32183  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 1 : i32} : memref<128x128xf16, 3> -> vector<1x2xf16>184  // CHECK: llvm.bitcast185  // CHECK: llvm.insertvalue186  return %a : vector<1x2xf16>187}188 189// CHECK-LABEL: @m16n8k4_tf32190func.func @m16n8k4_tf32(%arg0: vector<2x1xf32>, %arg1: vector<1x1xf32>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {191  // The A, B operand should be bitcast to i32192  // CHECK: llvm.extractvalue193  // CHECK: llvm.bitcast {{.*}} : vector<1xf32> to i32194  // CHECK: llvm.extractvalue195  // CHECK: llvm.bitcast {{.*}} : vector<1xf32> to i32196  // CHECK: llvm.extractvalue197  // CHECK: llvm.bitcast {{.*}} : vector<1xf32> to i32198 199  // CHECK: [[d:%.+]] = nvvm.mma.sync A[{{%.+}}, {{%.+}}] B[{{%.+}}] C[{{%.+}}, {{%.+}}, {{%.+}}, {{%.+}}]200  // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<tf32>201  // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<tf32>202  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 4>203  // CHECK-SAME: -> !llvm.struct<(f32, f32, f32, f32)>204  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 4], tf32Enabled} : (vector<2x1xf32>, vector<1x1xf32>, vector<2x2xf32>) -> vector<2x2xf32>205  // CHECK: [[undef:%.+]] = llvm.mlir.poison : vector<2xf32>206  // CHECK-DAG: llvm.extractvalue [[d]][0] : !llvm.struct<(f32, f32, f32, f32)>207  // CHECK-DAG: llvm.extractvalue [[d]][1] : !llvm.struct<(f32, f32, f32, f32)>208  // CHECK: [[d00:%.+]] = llvm.insertelement {{%.+}}, [[undef]][{{.*}}] : vector<2xf32>209  // CHECK: [[d01:%.+]] = llvm.insertelement {{%.+}}, [[d00]][{{.*}}] : vector<2xf32>210 211  // CHECK: [[undef:%.+]] = llvm.mlir.poison : vector<2xf32>212  // CHECK-DAG: llvm.extractvalue [[d]][2] : !llvm.struct<(f32, f32, f32, f32)>213  // CHECK-DAG: llvm.extractvalue [[d]][3] : !llvm.struct<(f32, f32, f32, f32)>214  // CHECK: [[d10:%.+]] = llvm.insertelement {{%.+}}, [[undef]][{{.*}}] : vector<2xf32>215  // CHECK: [[d11:%.+]] = llvm.insertelement {{%.+}}, [[d10]][{{.*}}] : vector<2xf32>216 217  // CHECK-DAG: llvm.insertvalue [[d01]], {{%.+}}[0] : !llvm.array<2 x vector<2xf32>>218  // CHECK-DAG: llvm.insertvalue [[d11]], {{%.+}}[1] : !llvm.array<2 x vector<2xf32>>219  return %d : vector<2x2xf32>220}221 222// CHECK-LABEL: @async_cp(223// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index)224func.func @async_cp(225  %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index) {226  // CHECK: %[[IDX1:.*]] = builtin.unrealized_conversion_cast %[[IDX]] : index to i64227  // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<3 x i64>, array<3 x i64>)>228  // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(2048 : index) : i64229  // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX1]], %[[S0]] : i64230  // CHECK-DAG: %[[S1:.*]] = llvm.mlir.constant(128 : index) : i64231  // CHECK-DAG: %[[FI0:.*]] = llvm.mul %[[IDX1]], %[[S1]] : i64232  // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI]], %[[FI0]] : i64233  // CHECK-DAG: %[[FI2:.*]] = llvm.add %[[FI1]], %[[IDX1]] : i64234  // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI2]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>235  // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>236  // CHECK-DAG: %[[S3:.*]] = llvm.mlir.constant(128 : index) : i64237  // CHECK-DAG: %[[FI3:.*]] = llvm.mul %[[IDX1]], %[[S3]]  : i64238  // CHECK-DAG: %[[FI4:.*]] = llvm.add %[[FI3]], %[[IDX1]]  : i64239  // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI4]]] : (!llvm.ptr, i64) -> !llvm.ptr240  // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[ADDRESSSRC]] : !llvm.ptr to !llvm.ptr<1>241  // CHECK-DAG: nvvm.cp.async.shared.global %[[ADDRESSDST]], %[[CAST2]], 16, cache = ca242  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4 : memref<128x128xf32> to memref<3x16x128xf32, 3>243  // CHECK: nvvm.cp.async.commit.group244  %1 = nvgpu.device_async_create_group %0245  // CHECK: nvvm.cp.async.wait.group 1246  nvgpu.device_async_wait %1 { numGroups = 1 : i32 }247 248  // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = cg249  %2 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4 {bypassL1}: memref<128x128xf32> to memref<3x16x128xf32, 3>250  return251}252 253// CHECK-LABEL: @async_cp_i4(254// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index)255func.func @async_cp_i4(256  %src: memref<128x64xi4>, %dst: memref<128x128xi4, 3>, %i : index) -> !nvgpu.device.async.token {257  // CHECK: %[[IDX1:.*]] = builtin.unrealized_conversion_cast %[[IDX]] : index to i64258  // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>259  // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(128 : index) : i64260  // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX1]], %[[S0]] : i64261  // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI]], %[[IDX1]] : i64262  // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI1]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>263  // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>264  // CHECK-DAG: %[[S2:.*]] = llvm.mlir.constant(64 : index) : i64265  // CHECK-DAG: %[[FI2:.*]] = llvm.mul %[[IDX1]], %[[S2]]  : i64266  // CHECK-DAG: %[[FI3:.*]] = llvm.add %[[FI2]], %[[IDX1]]  : i64267  // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI3]]] : (!llvm.ptr, i64) -> !llvm.ptr268  // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[ADDRESSSRC]] : !llvm.ptr to !llvm.ptr<1>269  // CHECK-DAG: nvvm.cp.async.shared.global %[[ADDRESSDST]], %[[CAST2]], 16, cache = ca270  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i], 32 : memref<128x64xi4> to memref<128x128xi4, 3>271  return %0 : !nvgpu.device.async.token272}273 274// CHECK-LABEL: @async_cp_zfill_f32_align4(275// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index, %[[SRCELEMENTS:[a-zA-Z0-9_]+]]: index276func.func @async_cp_zfill_f32_align4(277  %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index, %srcElements : index) {278  // CHECK-DAG: %[[IDX1:.*]] = builtin.unrealized_conversion_cast %[[IDX]] : index to i64279  // CHECK-DAG: %[[SRC1:.*]] = builtin.unrealized_conversion_cast %[[SRCELEMENTS]] : index to i64280  // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<3 x i64>, array<3 x i64>)>281  // CHECK-DAG: %[[S2048:.*]] = llvm.mlir.constant(2048 : index) : i64282  // CHECK-DAG: %[[LI1:.*]] = llvm.mul %[[IDX1]], %[[S2048]] : i64283  // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(128 : index) : i64284  // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX1]], %[[S0]] : i64285  // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI1]], %[[LI]] : i64286  // CHECK-DAG: %[[FI2:.*]] = llvm.add %[[FI1]], %[[IDX1]] : i64287  // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI2]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32288  // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>289  // CHECK-DAG: %[[S2:.*]] = llvm.mlir.constant(128 : index) : i64290  // CHECK-DAG: %[[FI2:.*]] = llvm.mul %[[IDX1]], %[[S2]]  : i64291  // CHECK-DAG: %[[FI3:.*]] = llvm.add %[[FI2]], %[[IDX1]]  : i64292  // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI3]]] : (!llvm.ptr, i64) -> !llvm.ptr293  // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[ADDRESSSRC]] : !llvm.ptr to !llvm.ptr<1>294  // CHECK-DAG: %[[c1:.*]] = llvm.mlir.constant(3 : i32) : i32295  // CHECK-DAG: %[[c2:.*]] = llvm.mlir.constant(32 : i32) : i32296  // CHECK-DAG: %[[c3:.*]] = llvm.trunc %[[SRC1]] : i64 to i32297  // CHECK-DAG: %[[c4:.*]] = llvm.mul %[[c2]], %[[c3]] : i32298  // CHECK-DAG: %[[c5:.*]] = llvm.lshr %[[c4]], %[[c1]] : i32299  // CHECK-DAG: nvvm.cp.async.shared.global %[[ADDRESSDST]], %[[CAST2]], 16, cache = cg, %[[c5]]300  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4, %srcElements {bypassL1}: memref<128x128xf32> to memref<3x16x128xf32, 3>301  // CHECK: nvvm.cp.async.commit.group302  %1 = nvgpu.device_async_create_group %0303  // CHECK: nvvm.cp.async.wait.group 1304  nvgpu.device_async_wait %1 { numGroups = 1 : i32 }305 306  return307}308 309// CHECK-LABEL: @async_cp_zfill_f32_align1(310// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index, %[[SRCELEMENTS:[a-zA-Z0-9_]+]]: index)311func.func @async_cp_zfill_f32_align1(312  %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index, %srcElements : index) {313  // CHECK-DAG: %[[IDX1:.*]] = builtin.unrealized_conversion_cast %[[IDX]] : index to i64314  // CHECK-DAG: %[[SRC1:.*]] = builtin.unrealized_conversion_cast %[[SRCELEMENTS]] : index to i64315  // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<3 x i64>, array<3 x i64>)>316  // CHECK-DAG: %[[S2048:.*]] = llvm.mlir.constant(2048 : index) : i64317  // CHECK-DAG: %[[LI1:.*]] = llvm.mul %[[IDX1]], %[[S2048]] : i64318  // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(128 : index) : i64319  // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX1]], %[[S0]] : i64320  // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI1]], %[[LI]] : i64321  // CHECK-DAG: %[[FI2:.*]] = llvm.add %[[FI1]], %[[IDX1]] : i64322  // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI2]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32323  // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>324  // CHECK-DAG: %[[S2:.*]] = llvm.mlir.constant(128 : index) : i64325  // CHECK-DAG: %[[FI2:.*]] = llvm.mul %[[IDX1]], %[[S2]]  : i64326  // CHECK-DAG: %[[FI3:.*]] = llvm.add %[[FI2]], %[[IDX1]]  : i64327  // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI3]]] : (!llvm.ptr, i64) -> !llvm.ptr328  // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[ADDRESSSRC]] : !llvm.ptr to !llvm.ptr<1>329  // CHECK-DAG: %[[c1:.*]] = llvm.mlir.constant(3 : i32) : i32330  // CHECK-DAG: %[[c2:.*]] = llvm.mlir.constant(32 : i32) : i32331  // CHECK-DAG: %[[c3:.*]] = llvm.trunc %[[SRC1]] : i64 to i32332  // CHECK-DAG: %[[c4:.*]] = llvm.mul %[[c2]], %[[c3]] : i32333  // CHECK-DAG: %[[c5:.*]] = llvm.lshr %[[c4]], %[[c1]] : i32334  // CHECK-DAG: nvvm.cp.async.shared.global %[[ADDRESSDST]], %[[CAST2]], 4, cache = ca, %[[c5]]335  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 1, %srcElements : memref<128x128xf32> to memref<3x16x128xf32, 3>336  // CHECK: nvvm.cp.async.commit.group337  %1 = nvgpu.device_async_create_group %0338  // CHECK: nvvm.cp.async.wait.group 1339  nvgpu.device_async_wait %1 { numGroups = 1 : i32 }340 341  return342}343 344 345// CHECK-LABEL: func @mma_sp_sync_f16_16832(346func.func @mma_sp_sync_f16_16832(%arg0: vector<4x2xf16>,347                                 %arg1: vector<4x2xf16>,348                                 %arg2: vector<2x2xf16>,349                                 %arg3: vector<2xi16>) -> vector<2x2xf16> {350  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<2xf16>>351  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<2xf16>>352  // CHECK: llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vector<2xf16>>353  // CHECK: llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vector<2xf16>>354 355  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<2xf16>>356  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<2xf16>>357  // CHECK: llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vector<2xf16>>358  // CHECK: llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vector<2xf16>>359 360  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>361  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>362 363  // CHECK-NOT: llvm.extractvalue364 365  // CHECK: %[[sparseMetadata:.+]] = llvm.bitcast %{{.+}} : vector<2xi16> to i32366 367  // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att368  // CHECK-SAME: "mma.sp.sync.aligned.m16n8k32.row.col.f16.f16.f16.f16 {$0,$1},{$2,$3,$4,$5},{$6,$7,$8,$9},{$10,$11},$12,0x0;"369  // CHECK-SAME: "=r,=r,r,r,r,r,r,r,r,r,r,r,r"370  // CHECK-SAME: %[[sparseMetadata]] :371  // CHECK-SAME: -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>372 373  %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) {mmaShape = [16, 8, 32]} :374    (vector<4x2xf16>, vector<4x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>375 376  // CHECK-DAG: llvm.extractvalue %[[d]][0] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>377  // CHECK-DAG: llvm.extractvalue %[[d]][1] : !llvm.struct<(vector<2xf16>, vector<2xf16>)>378  //     CHECK: llvm.mlir.poison : !llvm.array<2 x vector<2xf16>>379  //     CHECK: llvm.insertvalue %{{.+}}, %{{.+}}[0] : !llvm.array<2 x vector<2xf16>>380  //     CHECK: llvm.insertvalue %{{.+}}, %{{.+}}[1] : !llvm.array<2 x vector<2xf16>>381  return %d : vector<2x2xf16>382}383 384// CHECK-LABEL: func @mma_sp_sync_f16_16816(385func.func @mma_sp_sync_f16_16816(%arg0: vector<2x2xf16>,386                                 %arg1: vector<2x2xf16>,387                                 %arg2: vector<2x2xf16>,388                                 %arg3: vector<2xi16>) -> vector<2x2xf16> {389 390  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>391  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>392 393  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>394  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>395 396  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>397  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<2 x vector<2xf16>>398 399  // CHECK-NOT: llvm.extractvalue400 401  // CHECK: %[[sparseMetadata:.+]] = llvm.bitcast %{{.+}} : vector<2xi16> to i32402 403  // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att404  // CHECK-SAME: "mma.sp.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {$0,$1},{$2,$3},{$4,$5},{$6,$7},$8,0x0;"405  // CHECK-SAME: "=r,=r,r,r,r,r,r,r,r"406  // CHECK-SAME: %[[sparseMetadata]] :407  // CHECK-SAME: -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>408 409  %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) {mmaShape = [16, 8, 16]} :410    (vector<2x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>411  return %d : vector<2x2xf16>412}413 414// CHECK-LABEL: func @mma_sp_sync_f16_16816_01(415func.func @mma_sp_sync_f16_16816_01(%arg0: vector<2x2xf16>,416                                    %arg1: vector<2x2xf16>,417                                    %arg2: vector<2x2xf16>,418                                    %arg3: vector<2xi16>) -> vector<2x2xf16> {419  //420  // As above, but with sparsity selection 0x01.421  //422  // CHECK: %[[sparseMetadata:.+]] = llvm.bitcast %{{.+}} : vector<2xi16> to i32423  // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att424  // CHECK-SAME: "mma.sp.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {$0,$1},{$2,$3},{$4,$5},{$6,$7},$8,0x1;"425  // CHECK-SAME: "=r,=r,r,r,r,r,r,r,r"426  // CHECK-SAME: %[[sparseMetadata]] :427  // CHECK-SAME: -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>428 429  %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3)430       {mmaShape = [16, 8, 16], sparsitySelector = 1 : i32} :431       (vector<2x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>432  return %d : vector<2x2xf16>433}434 435// CHECK-LABEL: func @mma_sp_sync_i8_16864(436func.func @mma_sp_sync_i8_16864(%arg0: vector<4x4xi8>,437                                %arg1: vector<4x4xi8>,438                                %arg2: vector<2x2xi32>,439                                %arg3: vector<2xi16>) -> vector<2x2xi32> {440 441  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<4xi8>>442  // CHECK: llvm.bitcast %{{.+}} : vector<4xi8> to i32443  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<4xi8>>444  // CHECK: llvm.bitcast %{{.+}} : vector<4xi8> to i32445  // CHECK: llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vector<4xi8>>446  // CHECK: llvm.bitcast %{{.+}} : vector<4xi8> to i32447  // CHECK: llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vector<4xi8>>448 449 450  // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<4xi8>>451  // CHECK: llvm.bitcast %{{.+}} : vector<4xi8> to i32452  // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<4xi8>>453  // CHECK: llvm.bitcast %{{.+}} : vector<4xi8> to i32454 455  // CHECK: llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>456  // CHECK: llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<2xi32>>457 458  // CHECK-NOT: llvm.extractvalue459 460  // CHECK: %[[sparseMetadata:.+]] = llvm.bitcast %{{.+}} : vector<2xi16> to i32461 462  // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att463  // CHECK-SAME: "mma.sp.sync.aligned.m16n8k64.row.col.satfinite.s32.s8.s8.s32 {$0,$1,$2,$3},{$4,$5,$6,$7},{$8,$9,$10,$11},{$12,$13,$14,$15},$16,0x0;"464  // CHECK-SAME: "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r"465  // CHECK-SAME: %[[sparseMetadata]] :466  // CHECK-SAME: -> !llvm.struct<(i32, i32, i32, i32)467 468  %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) {mmaShape = [16, 8, 64]} :469    (vector<4x4xi8>, vector<4x4xi8>, vector<2x2xi32>) -> vector<2x2xi32>470  return %d : vector<2x2xi32>471}472 473!barrierType = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>474!tokenType = !nvgpu.mbarrier.token475 476// CHECK-LABEL: func @mbarrier477func.func @mbarrier() {478  %num_threads = arith.constant 128 : index479  // CHECK: %[[c0:.+]] = arith.constant 0 : index480  // CHECK: %[[mid:.+]] = builtin.unrealized_conversion_cast %[[c0]] : index to i64481  %c0 = arith.constant 0 : index482 483  // CHECK: %[[barMemref:.+]] = memref.get_global @__mbarrier{{.*}} : memref<1xi64, 3>484  %barrier = nvgpu.mbarrier.create -> !barrierType485 486  // CHECK: %[[barStr:.+]] =  builtin.unrealized_conversion_cast %[[barMemref]] : memref<1xi64, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>487  // CHECK: %[[base:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>488  // CHECK: %[[barPtr:.+]] = llvm.getelementptr %[[base]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64489  // CHECK: nvvm.mbarrier.init %[[barPtr]]490    nvgpu.mbarrier.init %barrier[%c0], %num_threads : !barrierType491 492  // CHECK: %[[base2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>493  // CHECK: %[[barPtr2:.+]] = llvm.getelementptr %[[base2]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64494  // CHECK: %[[token:.+]] = nvvm.mbarrier.arrive %[[barPtr2]]495  %token = nvgpu.mbarrier.arrive %barrier[%c0] : !barrierType -> !tokenType496 497  // CHECK: %[[base3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>498  // CHECK: %[[barPtr3:.+]] = llvm.getelementptr %[[base3]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64499  // CHECK: nvvm.mbarrier.test.wait %[[barPtr3]], %[[token]]500  %isDone = nvgpu.mbarrier.test.wait %barrier[%c0], %token : !barrierType, !tokenType501 502  func.return503}504 505// CHECK-LABEL: func @mbarrier_nocomplete506func.func @mbarrier_nocomplete() {507  %num_threads = arith.constant 128 : index508  %count = arith.constant 12 : index509  // CHECK: %[[c0:.+]] = arith.constant 0 : index510  // CHECK: %[[mid:.+]] = builtin.unrealized_conversion_cast %[[c0]] : index to i64511  %c0 = arith.constant 0 : index512 513  // CHECK: %[[barMemref:.+]] = memref.get_global @__mbarrier{{.*}} : memref<1xi64, 3>514  %barrier = nvgpu.mbarrier.create -> !barrierType515 516  // CHECK: %[[barStr:.+]] =  builtin.unrealized_conversion_cast %[[barMemref]] : memref<1xi64, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>517  // CHECK: %[[base:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>518  // CHECK: %[[barPtr:.+]] = llvm.getelementptr %[[base]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64519  // CHECK: nvvm.mbarrier.init %[[barPtr]]520  nvgpu.mbarrier.init %barrier[%c0], %num_threads : !barrierType521 522  // CHECK: %[[base2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>523  // CHECK: %[[barPtr2:.+]] = llvm.getelementptr %[[base2]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64524  // CHECK: %[[token:.+]] = nvvm.mbarrier.arrive.nocomplete %[[barPtr2]]525  %token = nvgpu.mbarrier.arrive.nocomplete %barrier[%c0], %count : !barrierType -> !tokenType526 527  // CHECK: %[[base3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>528  // CHECK: %[[barPtr3:.+]] = llvm.getelementptr %[[base3]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64529  // CHECK: nvvm.mbarrier.test.wait %[[barPtr3]], %[[token]]530  %isDone = nvgpu.mbarrier.test.wait %barrier[%c0], %token : !barrierType, !tokenType531 532  func.return533}534 535// CHECK-LABEL: func @mbarrier_get536//  CHECK-SAME:     %[[ARG0:.*]]: !nvgpu.mbarrier.group{{.*}}537func.func @mbarrier_get(%barriers : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5>) {538  // CHECK: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[ARG0]] : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>539  // CHECK: %[[c2:.+]] = arith.constant 2 : index540  // CHECK: %[[S1:.+]] = builtin.unrealized_conversion_cast %[[c2]] : index to i64541  // CHECK: %[[S2:.+]] = llvm.extractvalue %[[S0]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 542  // CHECK: %[[S3:.+]] = llvm.getelementptr %[[S2]][%[[S1]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64543  // CHECK: %[[S4:.+]] = llvm.ptrtoint %[[S3]] : !llvm.ptr<3> to i32544  %c2 = arith.constant 2 : index545  nvgpu.mbarrier.get %barriers[%c2] : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5> -> i32546 547  // CHECK: %[[c4:.+]] = arith.constant 4 : index548  // CHECK: %[[S5:.+]] = builtin.unrealized_conversion_cast %[[c4]] : index to i64549  // CHECK: %[[S6:.+]] = llvm.extractvalue %[[S0]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 550  // CHECK: %[[S7:.+]] = llvm.getelementptr %[[S6]][%[[S5]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64551  // CHECK: %[[S8:.+]] = llvm.ptrtoint %[[S7]] : !llvm.ptr<3> to i64552  %c4 = arith.constant 4 : index553  nvgpu.mbarrier.get %barriers[%c4] : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5> -> i64554  func.return555}556 557// CHECK-LABEL: func @mbarrier_wait(558//  CHECK-SAME:     %[[ARG0:.*]]: !nvgpu.mbarrier.group{{.*}}, %[[ARG1:.*]]: !nvgpu.mbarrier.token)559func.func @mbarrier_wait(%barriers : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5>, %token : !tokenType) {560// CHECK-DAG: %[[CARG0:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]561// CHECK-DAG: %[[CARG1:.*]] = builtin.unrealized_conversion_cast %[[ARG1]]562  %c0 = arith.constant 0 : index563  %c1 = arith.constant 1 : index564  %n = arith.constant 100 : index565 566  %numBarriers = arith.constant 5 : index567 568  scf.for %i = %c0 to %n step %c1 {569// CHECK: %[[c5:.+]] = arith.constant 5 : index570// CHECK: scf.for %[[i:.*]] =571// CHECK: %[[S2:.+]] = arith.remui %[[i]], %[[c5]] : index572// CHECK: %[[S3:.+]] = builtin.unrealized_conversion_cast %[[S2]] : index to i64573// CHECK: %[[S4:.+]] = llvm.extractvalue %[[CARG0]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>574// CHECK: %[[S5:.+]] = llvm.getelementptr %[[S4]][%[[S3]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64575// CHECK: nvvm.mbarrier.test.wait {{.*}}, %[[CARG1]]576    %mbarId = arith.remui %i, %numBarriers : index577    %isDone = nvgpu.mbarrier.test.wait %barriers[%mbarId], %token : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 5>, !tokenType578  }579  return580}581 582// CHECK-LABEL: func @mbarrier_txcount583func.func @mbarrier_txcount() {584    %num_threads = arith.constant 128 : index585    // CHECK: %[[c0:.+]] = arith.constant 0 : index586    // CHECK: %[[mid:.+]] = builtin.unrealized_conversion_cast %[[c0]] : index to i64587    %c0 = arith.constant 0 : index588 589    // CHECK: %[[barMemref:.+]] = memref.get_global @__mbarrier{{.*}} : memref<1xi64, 3>590    %barrier = nvgpu.mbarrier.create -> !barrierType591 592    // CHECK: %[[barStr:.+]] =  builtin.unrealized_conversion_cast %[[barMemref]] : memref<1xi64, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>593    // CHECK: %[[base:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>594    // CHECK: %[[barPtr:.+]] = llvm.getelementptr %[[base]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64595    // CHECK: nvvm.mbarrier.init %[[barPtr]]596    nvgpu.mbarrier.init %barrier[%c0], %num_threads : !barrierType597 598    %tidxreg = nvvm.read.ptx.sreg.tid.x : i32599    %tidx = arith.index_cast %tidxreg : i32 to index600    %cnd = arith.cmpi eq, %tidx, %c0 : index601 602    scf.if %cnd {603      %txcount = arith.constant 256 : index604      // CHECK: %[[base2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>605      // CHECK: %[[barPtr2:.+]] = llvm.getelementptr %[[base2]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64606      // CHECK: nvvm.mbarrier.arrive.expect_tx %[[barPtr2]]607      nvgpu.mbarrier.arrive.expect_tx %barrier[%c0], %txcount : !barrierType608      scf.yield609    } else {610      %txcount = arith.constant 0 : index611      // CHECK: %[[base2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>612      // CHECK: %[[barPtr2:.+]] = llvm.getelementptr %[[base2]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64613      // CHECK: nvvm.mbarrier.arrive.expect_tx %[[barPtr2]]614      nvgpu.mbarrier.arrive.expect_tx %barrier[%c0], %txcount : !barrierType615      scf.yield616    }617 618 619    %phase_c0 = arith.constant 0 : i1620    %ticks = arith.constant 10000000 : index621    // CHECK: %[[base3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>622    // CHECK: %[[barPtr3:.+]] = llvm.getelementptr %[[base3]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64623    // CHECK: nvvm.mbarrier.try_wait.parity %[[barPtr3]]624    nvgpu.mbarrier.try_wait.parity %barrier[%c0], %phase_c0, %ticks : !barrierType625 626    func.return627}628 629// CHECK-LABEL: func @mbarrier_txcount_pred630func.func @mbarrier_txcount_pred() {631    %mine = arith.constant 1 : index632    // CHECK: %[[c0:.+]] = arith.constant 0 : index633    // CHECK: %[[mid:.+]] = builtin.unrealized_conversion_cast %[[c0]] : index to i64634    // CHECK: %[[S2:.+]] = gpu.thread_id  x635    // CHECK: %[[P:.+]] = arith.cmpi eq, %[[S2]], %[[c0]] : index636    %c0 = arith.constant 0 : index637    %tidx = gpu.thread_id x638    %pred = arith.cmpi eq, %tidx, %c0 : index639 640    // CHECK: %[[barMemref:.+]] = memref.get_global @__mbarrier{{.*}} : memref<1xi64, 3>641    %barrier = nvgpu.mbarrier.create -> !barrierType642 643    // CHECK: %[[barStr:.+]] =  builtin.unrealized_conversion_cast %[[barMemref]] : memref<1xi64, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>644    // CHECK: %[[base:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>645    // CHECK: %[[barPtr:.+]] = llvm.getelementptr %[[base]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64646    // CHECK: nvvm.mbarrier.init %[[barPtr]], {{.*}}, predicate = %[[P]]647    nvgpu.mbarrier.init %barrier[%c0], %mine, predicate = %pred : !barrierType648 649    %txcount = arith.constant 256 : index650    // CHECK: %[[base2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>651    // CHECK: %[[barPtr2:.+]] = llvm.getelementptr %[[base2]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64652    // CHECK: nvvm.mbarrier.arrive.expect_tx %[[barPtr2]], {{.*}}, predicate = %[[P]]653    nvgpu.mbarrier.arrive.expect_tx %barrier[%c0], %txcount, predicate = %pred : !barrierType654 655    %phase_c0 = arith.constant 0 : i1656    %ticks = arith.constant 10000000 : index657    // CHECK: %[[base3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>658    // CHECK: %[[barPtr3:.+]] = llvm.getelementptr %[[base3]][%[[mid]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i64659    // CHECK: nvvm.mbarrier.try_wait.parity %[[barPtr3]]660    nvgpu.mbarrier.try_wait.parity %barrier[%c0], %phase_c0, %ticks : !barrierType661 662    func.return663}664 665// CHECK-LABEL: func @async_tma_load666!tensorMap1d = !nvgpu.tensormap.descriptor<tensor = memref<128xf32,3>,         swizzle=none,        l2promo = none,        oob = nan,  interleave = none>667!tensorMap2d = !nvgpu.tensormap.descriptor<tensor = memref<32x8xf32,3>,       swizzle=swizzle_32b, l2promo = none,        oob = zero, interleave = none>668!tensorMap3d = !nvgpu.tensormap.descriptor<tensor = memref<2x32x16xf32,3>,     swizzle=swizzle_64b, l2promo = l2promo_64b, oob = zero, interleave = none>669!tensorMap4d = !nvgpu.tensormap.descriptor<tensor = memref<2x2x32x32xf32,3>,   swizzle=swizzle_128b,l2promo = l2promo_128b,oob = zero, interleave = none>670!tensorMap5d = !nvgpu.tensormap.descriptor<tensor = memref<2x2x2x32x32xf32,3>, swizzle=none,        l2promo = none,        oob = zero, interleave = none>671!mbarrier = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>672func.func @async_tma_load(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d, %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d, %tensorMap5d: !tensorMap5d,673                              %buffer1d: memref<128xf32,3>,674                              %buffer2d: memref<32x8xf32,3>,675                              %buffer3d: memref<2x32x16xf32,3>,676                              %buffer4d: memref<2x2x32x32xf32,3>,677                              %buffer5d: memref<2x2x2x32x32xf32,3>,678                              %mbarrier: !mbarrier) {679  %c0 = arith.constant 0 : index680  %crd0 = arith.constant 0 : index681  %crd1 = arith.constant 0 : index682  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}]683  nvgpu.tma.async.load %tensorMap1d[%crd0], %mbarrier[%c0] to %buffer1d : !tensorMap1d, !mbarrier -> memref<128xf32,3>684  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}]685  nvgpu.tma.async.load %tensorMap2d[%crd0, %crd1], %mbarrier[%c0] to %buffer2d : !tensorMap2d, !mbarrier -> memref<32x8xf32,3>686  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}]687  nvgpu.tma.async.load %tensorMap3d[%crd0, %crd1, %crd0], %mbarrier[%c0] to %buffer3d : !tensorMap3d, !mbarrier -> memref<2x32x16xf32,3>688  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]689  nvgpu.tma.async.load %tensorMap4d[%crd0, %crd1, %crd1, %crd0], %mbarrier[%c0] to %buffer4d : !tensorMap4d, !mbarrier -> memref<2x2x32x32xf32,3>690  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]691  nvgpu.tma.async.load %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0], %mbarrier[%c0] to %buffer5d : !tensorMap5d, !mbarrier -> memref<2x2x2x32x32xf32,3>692  func.return693}694 695// CHECK-LABEL: func @async_tma_load696!tensorMap1dgpuspace = !nvgpu.tensormap.descriptor<tensor = memref<128xf32, #gpu.address_space<workgroup>>,         swizzle=none,        l2promo = none,        oob = nan,  interleave = none>697func.func @async_tma_load_gpu_address_space(%tensorMap1d: !tensorMap1dgpuspace,698                          %buffer1d: memref<128xf32, #gpu.address_space<workgroup>>,699                          %mbarrier: !mbarrier) {700  %c0 = arith.constant 0 : index701  %crd0 = arith.constant 0 : index702  %crd1 = arith.constant 0 : index703  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}]704  nvgpu.tma.async.load %tensorMap1d[%crd0], %mbarrier[%c0] to %buffer1d : !tensorMap1dgpuspace, !mbarrier -> memref<128xf32,#gpu.address_space<workgroup>>705   func.return706}707 708// CHECK-LABEL: func @async_tma_load_pred709func.func @async_tma_load_pred(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d, %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d, %tensorMap5d: !tensorMap5d,710                              %buffer1d: memref<128xf32,3>,711                              %buffer2d: memref<32x8xf32,3>,712                              %buffer3d: memref<2x32x16xf32,3>,713                              %buffer4d: memref<2x2x32x32xf32,3>,714                              %buffer5d: memref<2x2x2x32x32xf32,3>,715                              %mbarrier: !mbarrier,716                              %p: i1) {717  %c0 = arith.constant 0 : index718  %crd0 = arith.constant 0 : index719  %crd1 = arith.constant 0 : index720  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}] predicate = %{{.*}}721  nvgpu.tma.async.load %tensorMap1d[%crd0], %mbarrier[%c0] to %buffer1d, predicate = %p : !tensorMap1d, !mbarrier -> memref<128xf32,3>722  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}] predicate = %{{.*}}723  nvgpu.tma.async.load %tensorMap2d[%crd0, %crd1], %mbarrier[%c0] to %buffer2d, predicate = %p : !tensorMap2d, !mbarrier -> memref<32x8xf32,3>724  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}] predicate = %{{.*}}725  nvgpu.tma.async.load %tensorMap3d[%crd0, %crd1, %crd0], %mbarrier[%c0] to %buffer3d, predicate = %p : !tensorMap3d, !mbarrier -> memref<2x32x16xf32,3>726  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] predicate = %{{.*}}727  nvgpu.tma.async.load %tensorMap4d[%crd0, %crd1, %crd1, %crd0], %mbarrier[%c0] to %buffer4d, predicate = %p : !tensorMap4d, !mbarrier -> memref<2x2x32x32xf32,3>728  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] predicate = %{{.*}}729  nvgpu.tma.async.load %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0], %mbarrier[%c0] to %buffer5d, predicate = %p : !tensorMap5d, !mbarrier -> memref<2x2x2x32x32xf32,3>730  func.return731}732 733func.func @async_tma_load_multicast(734  %tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d,735  %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d,736  %tensorMap5d: !tensorMap5d, %buffer1d: memref<128xf32,3>,737  %buffer2d: memref<32x8xf32,3>, %buffer3d: memref<2x32x16xf32,3>,738  %buffer4d: memref<2x2x32x32xf32,3>, %buffer5d: memref<2x2x2x32x32xf32,3>,739  %mbarrier: !mbarrier,740  %multicastMask: i16) {741  %c0 = arith.constant 0 : index742  %crd0 = arith.constant 0 : index743  %crd1 = arith.constant 0 : index744  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}]745  nvgpu.tma.async.load %tensorMap1d[%crd0], %mbarrier[%c0] to %buffer1d multicast_mask = %multicastMask : !tensorMap1d, !mbarrier -> memref<128xf32,3>746  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}]747  nvgpu.tma.async.load %tensorMap2d[%crd0, %crd1], %mbarrier[%c0] to %buffer2d multicast_mask = %multicastMask : !tensorMap2d, !mbarrier -> memref<32x8xf32,3>748  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}]749  nvgpu.tma.async.load %tensorMap3d[%crd0, %crd1, %crd0], %mbarrier[%c0] to %buffer3d multicast_mask = %multicastMask : !tensorMap3d, !mbarrier -> memref<2x32x16xf32,3>750  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]751  nvgpu.tma.async.load %tensorMap4d[%crd0, %crd1, %crd1, %crd0], %mbarrier[%c0] to %buffer4d multicast_mask = %multicastMask : !tensorMap4d, !mbarrier -> memref<2x2x32x32xf32,3>752  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}} box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]753  nvgpu.tma.async.load %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0], %mbarrier[%c0] to %buffer5d multicast_mask = %multicastMask : !tensorMap5d, !mbarrier -> memref<2x2x2x32x32xf32,3>754  func.return755}756 757func.func @async_tma_store(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d, %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d, %tensorMap5d: !tensorMap5d,758                           %buffer1d: memref<128xf32,3>,759                           %buffer2d: memref<32x8xf32,3>,760                           %buffer3d: memref<2x32x16xf32,3>,761                           %buffer4d: memref<2x2x32x32xf32,3>,762                           %buffer5d: memref<2x2x2x32x32xf32,3>) {763  %c0 = arith.constant 0 : index764  %crd0 = arith.constant 0 : index765  %crd1 = arith.constant 0 : index766  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}]767  nvgpu.tma.async.store %buffer1d to %tensorMap1d[%crd0] : memref<128xf32,3> -> !tensorMap1d768  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}]769  nvgpu.tma.async.store %buffer2d to %tensorMap2d[%crd0, %crd1]  : memref<32x8xf32,3> -> !tensorMap2d770  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}]771  nvgpu.tma.async.store %buffer3d to %tensorMap3d[%crd0, %crd1, %crd0]  : memref<2x32x16xf32,3> -> !tensorMap3d772  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]773  nvgpu.tma.async.store %buffer4d to %tensorMap4d[%crd0, %crd1, %crd1, %crd0]  : memref<2x2x32x32xf32,3> -> !tensorMap4d774  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]775  nvgpu.tma.async.store %buffer5d to %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0]  : memref<2x2x2x32x32xf32,3> -> !tensorMap5d776  func.return777}778 779 780func.func @async_tma_store_predicate(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d, %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d, %tensorMap5d: !tensorMap5d,781                           %buffer1d: memref<128xf32,3>,782                           %buffer2d: memref<32x8xf32,3>,783                           %buffer3d: memref<2x32x16xf32,3>,784                           %buffer4d: memref<2x2x32x32xf32,3>,785                           %buffer5d: memref<2x2x2x32x32xf32,3>,786                           %p: i1) {787  %c0 = arith.constant 0 : index788  %crd0 = arith.constant 0 : index789  %crd1 = arith.constant 0 : index790  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}], predicate = %{{.*}}791  nvgpu.tma.async.store %buffer1d to %tensorMap1d[%crd0], predicate = %p : memref<128xf32,3> -> !tensorMap1d792  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}], predicate = %{{.*}}793  nvgpu.tma.async.store %buffer2d to %tensorMap2d[%crd0, %crd1], predicate = %p  : memref<32x8xf32,3> -> !tensorMap2d794  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}], predicate = %{{.*}}795  nvgpu.tma.async.store %buffer3d to %tensorMap3d[%crd0, %crd1, %crd0], predicate = %p  : memref<2x32x16xf32,3> -> !tensorMap3d796  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], predicate = %{{.*}}797  nvgpu.tma.async.store %buffer4d to %tensorMap4d[%crd0, %crd1, %crd1, %crd0], predicate = %p  : memref<2x2x32x32xf32,3> -> !tensorMap4d798  // CHECK: nvvm.cp.async.bulk.tensor.global.shared.cta %{{.*}} %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], predicate = %{{.*}}799  nvgpu.tma.async.store %buffer5d to %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0], predicate = %p  : memref<2x2x2x32x32xf32,3> -> !tensorMap5d800  func.return801}802 803func.func @create_tensor_map(%devicePtr2d : memref<64x128xf32>, %devicePtr1d : memref<128xf32>) {804  %crd0 = arith.constant 64 : index805  %crd1 = arith.constant 128 : index806  %devicePtr2d_unranked = memref.cast %devicePtr2d : memref<64x128xf32> to memref<*xf32>807  // CHECK: llvm.call @mgpuTensorMapEncodeTiledMemref808  %tensorMap2d = nvgpu.tma.create.descriptor %devicePtr2d_unranked box[%crd0, %crd1] : memref<*xf32> -> !tensorMap2d809 810  %devicePtr1d_unranked = memref.cast %devicePtr1d : memref<128xf32> to memref<*xf32>811  // CHECK: llvm.call @mgpuTensorMapEncodeTiledMemref812  %tensorMap1d = nvgpu.tma.create.descriptor %devicePtr1d_unranked box[%crd1] : memref<*xf32> -> !tensorMap1d813  func.return814}815 816// CHECK-LABEL: @tma_prefetch(817// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none>, %[[arg1:[a-zA-Z0-9_]+]]: i1818func.func @tma_prefetch(%tensorMap1d: !tensorMap1d, %p : i1) {819  // CHECK: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none> to !llvm.ptr820  // CHECK: nvvm.prefetch tensormap, %[[S0]] : !llvm.ptr821  nvgpu.tma.prefetch.descriptor %tensorMap1d: !tensorMap1d822  // CHECK: nvvm.prefetch tensormap, %[[S0]], predicate = %[[arg1]] : !llvm.ptr, i1823  nvgpu.tma.prefetch.descriptor %tensorMap1d, predicate = %p: !tensorMap1d824  func.return825}826 827 828// CHECK-LABEL: @tma_fence(829// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none>830func.func @tma_fence(%tensorMap1d: !tensorMap1d) {831  // CHECK: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none> to !llvm.ptr832  // CHECK: %[[S1:.+]] = llvm.mlir.constant(128 : i32) : i32833  // CHECK: nvvm.fence.proxy.acquire <sys> %[[S0]], %[[S1]]834  nvgpu.tma.fence.descriptor %tensorMap1d: !tensorMap1d835  func.return836}837 838!lhsTensorMap = !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = none, oob = zero, interleave = none>839!rhsTensorMap = !nvgpu.tensormap.descriptor<tensor = memref<64x64xf16, strided<[64, 1], offset: 8192>, 3>, swizzle = swizzle_128b, l2promo = none, oob = zero, interleave = none>840 841module @mymodule {842  // Dynamic Shared memory843  memref.global "private" @dynamicShmem : memref<0xf16,3>844 845  func.func @async_tma_load(%lhsTensorMap: !lhsTensorMap, %rhsTensorMap: !rhsTensorMap, %mbarrier: !barrierType) {846    %c0 = arith.constant 0 : index847    %dynamicMem = memref.get_global @dynamicShmem : memref<0xf16, 3>848    %lhsShmem = memref.reinterpret_cast %dynamicMem to offset: [0], sizes: [128,64], strides: [64,1] : memref<0xf16, 3> to memref<128x64xf16,3>849    %rhsShmem2 = memref.reinterpret_cast %dynamicMem to offset: [0], sizes: [4, 64, 64],  strides: [4096, 64, 1] : memref<0xf16, 3> to memref<4x64x64xf16,3>850    %rhsShmem3 = memref.subview %rhsShmem2[2, 0, 0][1, 64, 64][1, 1, 1] : memref<4x64x64xf16,3> to memref<1x64x64xf16, strided<[4096, 64, 1], offset: 8192>, 3>851    %rhsShmem = memref.subview %rhsShmem3[0, 0, 0][1, 64, 64][1, 1, 1]  : memref<1x64x64xf16, strided<[4096, 64, 1], offset: 8192>, 3> to memref<64x64xf16, strided<[64, 1], offset: 8192>, 3>852    // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global853    nvgpu.tma.async.load %lhsTensorMap[%c0, %c0], %mbarrier[%c0] to %lhsShmem : !lhsTensorMap, !barrierType -> memref<128x64xf16,3>854    // CHECK: %[[desc:.+]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>855    // CHECK: %[[c8192:.+]] = llvm.mlir.constant(8192 : index) : i64856    // CHECK: %[[shmemOfset:.+]] = llvm.getelementptr %[[desc]][%[[c8192]]] : (!llvm.ptr<3>, i64)857    // CHECK: %[[dest:.+]] = llvm.addrspacecast %[[shmemOfset]] : !llvm.ptr<3> to !llvm.ptr<7>858    // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %[[dest]], %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}]859    nvgpu.tma.async.load %rhsTensorMap[%c0, %c0], %mbarrier[%c0] to %rhsShmem : !rhsTensorMap, !barrierType -> memref<64x64xf16, strided<[64, 1], offset: 8192>, 3>860    return861  }862}863 864!tensorMap = !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16,3>, swizzle = swizzle_128b, l2promo=none, oob=zero, interleave=none>865memref.global "private" @dynamicShmem : memref<0xf16,3>866// CHECK-LABEL: func @create_wgmma_descriptor(867func.func @create_wgmma_descriptor(%tensorMap : !tensorMap) -> !nvgpu.warpgroup.descriptor<tensor=memref<128x64xf16,3>>{868  %dynamicMem = memref.get_global @dynamicShmem : memref<0xf16, 3>869  %lhsShmem = memref.reinterpret_cast %dynamicMem to offset: [0], sizes: [128,64], strides: [64,1] : memref<0xf16, 3> to memref<128x64xf16,3>870    // CHECK: %[[S0:.+]] = memref.get_global @dynamicShmem : memref<0xf16, 3>871    // CHECK: %[[Sre:.+]] = memref.reinterpret_cast %[[S0]] to offset: [0], sizes: [128, 64], strides: [64, 1] : memref<0xf16, 3> to memref<128x64xf16, 3>872    // CHECK: %[[S1:.+]] = builtin.unrealized_conversion_cast %[[Sre]] : memref<128x64xf16, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>873    // CHECK: %[[c64:.+]] =  llvm.mlir.constant(64 : i64) : i64874    // CHECK: %[[c1024:.+]] = llvm.mlir.constant(1024 : i64) : i64875    // CHECK: %[[S2:.+]] = llvm.extractvalue %[[S1]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>876    // CHECK: %[[S3:.+]] = llvm.ptrtoint %[[S2]] : !llvm.ptr<3> to i64877    // CHECK: %[[S4:.+]] = llvm.mlir.constant(46 : i64) : i64878    // CHECK: %[[S5:.+]] = llvm.shl %[[S3]], %[[S4]]  : i64879    // CHECK: %[[S6:.+]] = llvm.mlir.constant(50 : i64) : i64880    // CHECK: %[[S7:.+]] = llvm.lshr %[[S5]], %[[S6]]  : i64881    // CHECK: %[[S8:.+]] = llvm.mlir.constant(0 : i64) : i64882    // CHECK: %[[S9:.+]] = llvm.mlir.constant(1 : i64) : i64883    // CHECK: %[[S10:.+]] = llvm.mlir.constant(62 : i64) : i64884    // CHECK: %[[S11:.+]] = llvm.shl %[[S9]], %[[S10]]  : i64885    // CHECK: %[[S12:.+]] = llvm.or %[[S8]], %[[S11]]  : i64886    // CHECK: %[[S13:.+]] = llvm.mlir.constant(0 : i64) : i64887    // CHECK: %[[S14:.+]] = llvm.mlir.constant(49 : i64) : i64888    // CHECK: %[[S15:.+]] = llvm.shl %[[S13]], %[[S14]]  : i64889    // CHECK: %[[S16:.+]] = llvm.or %[[S12]], %[[S15]]  : i64890    // CHECK: %[[S18:.+]] = llvm.mlir.constant(32 : i64) : i64891    // CHECK: %[[S19:.+]] = llvm.shl %[[c64]], %[[S18]]  : i64892    // CHECK: %[[S20:.+]] = llvm.or %[[S16]], %[[S19]]  : i64893    // CHECK: %[[S22:.+]] = llvm.mlir.constant(16 : i64) : i64894    // CHECK: %[[S23:.+]] = llvm.shl %[[c1024]], %[[S22]]  : i64895    // CHECK: %[[S24:.+]] = llvm.or %[[S20]], %[[S23]]  : i64896    // CHECK: %[[S25:.+]] = llvm.mlir.constant(0 : i64) : i64897    // CHECK: %[[S26:.+]] = llvm.shl %[[S7]], %[[S25]]  : i64898    // CHECK: %[[S27:.+]] = llvm.or %[[S24]], %[[S26]]  : i64899    // CHECK: %[[ret:.+]] = builtin.unrealized_conversion_cast %[[S27]] : i64 to !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>900    // CHECK: return %[[ret]]901  %descA = nvgpu.warpgroup.generate.descriptor %lhsShmem, %tensorMap : memref<128x64xf16,3>, !tensorMap -> !nvgpu.warpgroup.descriptor<tensor=memref<128x64xf16,3>>902  func.return %descA : !nvgpu.warpgroup.descriptor<tensor=memref<128x64xf16,3>>903}904 905// CHECK-LABEL: @warpgroup_mma_128_128_64(906// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>, %[[arg1:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>, %[[arg2:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>)907func.func @warpgroup_mma_128_128_64(908      %descA: !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>,909      %descB: !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>,910      %acc: !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>)911{912// CHECK-DAG: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>> to i64913// CHECK-DAG: %[[S1:.+]] = builtin.unrealized_conversion_cast %[[arg1]] : !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>> to i64914// CHECK-DAG: %[[ARG:.+]] = builtin.unrealized_conversion_cast %[[arg2]] : !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>> to !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>915// CHECK: nvvm.wgmma.fence.aligned916// CHECK: %[[UD:.+]] =  llvm.mlir.poison : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>917// CHECK: %[[S2:.+]] = llvm.extractvalue %[[ARG]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>918// CHECK: %[[S4:.+]] = nvvm.wgmma.mma_async %[[S0]], %[[S1]], %[[S2]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>919// CHECK: %[[S5:.+]] = llvm.mlir.constant(2 : i32) : i64920// CHECK: %[[S6:.+]] = llvm.add %[[S0]], %[[S5]] : i64921// CHECK: %[[S7:.+]] = llvm.mlir.constant(128 : i32) : i64922// CHECK: %[[S8:.+]] = llvm.add %[[S1]], %[[S7]]  : i64923// CHECK: %[[S9:.+]] = nvvm.wgmma.mma_async %[[S6]], %[[S8]], %[[S4]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct924// CHECK: %[[S10:.+]] = llvm.mlir.constant(4 : i32) : i64925// CHECK: %[[S11:.+]] = llvm.add %[[S0]], %[[S10]]  : i64926// CHECK: %[[S12:.+]] = llvm.mlir.constant(256 : i32) : i64927// CHECK: %[[S13:.+]] = llvm.add %[[S1]], %[[S12]]  : i64928// CHECK: %[[S14:.+]] = nvvm.wgmma.mma_async %[[S11]], %[[S13]], %[[S9]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct929// CHECK: %[[S15:.+]] = llvm.mlir.constant(6 : i32) : i64930// CHECK: %[[S16:.+]] = llvm.add %[[S0]], %[[S15]]  : i64931// CHECK: %[[S17:.+]] = llvm.mlir.constant(384 : i32) : i64932// CHECK: %[[S18:.+]] = llvm.add %[[S1]], %[[S17]]  : i64933// CHECK: %[[S19:.+]] = nvvm.wgmma.mma_async %[[S16]], %[[S18]], %[[S14]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct934// CHECK: %[[S3:.+]] = llvm.extractvalue %[[ARG]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>935// CHECK: %[[S21:.+]] = llvm.mlir.constant(512 : i32) : i64936// CHECK: %[[S22:.+]] = llvm.add %[[S0]], %[[S21]]  : i64937// CHECK: %[[S23:.+]] = nvvm.wgmma.mma_async %[[S22]], %[[S1]], %[[S3]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct938// CHECK: %[[S24:.+]] = llvm.mlir.constant(514 : i32) : i64939// CHECK: %[[S25:.+]] = llvm.add %[[S0]], %[[S24]]  : i64940// CHECK: %[[S26:.+]] = llvm.mlir.constant(128 : i32) : i64941// CHECK: %[[S27:.+]] = llvm.add %[[S1]], %[[S26]]  : i64942// CHECK: %[[S28:.+]] = nvvm.wgmma.mma_async %[[S25]], %[[S27]], %[[S23]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct943// CHECK: %[[S29:.+]] = llvm.mlir.constant(516 : i32) : i64944// CHECK: %[[S30:.+]] = llvm.add %[[S0]], %[[S29]]  : i64945// CHECK: %[[S31:.+]] = llvm.mlir.constant(256 : i32) : i64946// CHECK: %[[S32:.+]] = llvm.add %[[S1]], %[[S31]]  : i64947// CHECK: %[[S33:.+]] = nvvm.wgmma.mma_async %[[S30]], %[[S32]], %[[S28]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct948// CHECK: %[[S34:.+]] = llvm.mlir.constant(518 : i32) : i64949// CHECK: %[[S35:.+]] = llvm.add %[[S0]], %[[S34]]  : i64950// CHECK: %[[S36:.+]] = llvm.mlir.constant(384 : i32) : i64951// CHECK: %[[S37:.+]] = llvm.add %[[S1]], %[[S36]]  : i64952// CHECK: %[[S38:.+]] = nvvm.wgmma.mma_async %[[S35]], %[[S37]], %[[S33]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct953// CHECK: %[[S40:.+]] = llvm.insertvalue %[[S19]], %[[UD]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>954// CHECK: %[[S41:.+]] = llvm.insertvalue %[[S38]], %[[S40]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>955// CHECK: nvvm.wgmma.commit.group.sync.aligned956// CHECK: nvvm.wgmma.wait.group.sync.aligned 1957  %wgmmaResult = nvgpu.warpgroup.mma %descA, %descB, %acc {transposeB}:958      !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>,959      !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>,960      !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>961      ->962      !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>963  return964}965 966// CHECK-LABEL: @warpgroup_mma_store(967// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>, %[[arg2:[a-zA-Z0-9_]+]]: memref<128x128xf32, 3>)968func.func @warpgroup_mma_store(969    %result : !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>,970    %matrixD: memref<128x128xf32,3>) {971// CHECK: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>> to !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>972// CHECK: %[[EX1:.+]] = llvm.extractvalue %[[S0]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>973// CHECK: %[[S6:.+]] = llvm.mlir.constant(1 : i32) : i32974// CHECK: %[[S5:.+]] = llvm.mlir.constant(2 : i32) : i32975// CHECK: %[[S2:.+]] = llvm.mlir.constant(4 : i32) : i32976// CHECK: %[[S4:.+]] = llvm.mlir.constant(8 : i32) : i32977// CHECK: %[[S7:.+]] = llvm.mlir.constant(16 : i32) : i32978// CHECK: %[[WarpSize:.+]] = llvm.mlir.constant(32 : i32) : i32979 980// ### Store {d0, d1} of each thread ###981 982// CHECK: %[[S8:.+]] = nvvm.read.ptx.sreg.tid.x : i32983// CHECK: %[[S9:.+]] = llvm.urem %[[S8]], %[[WarpSize]]  : i32984// CHECK: %[[S10:.+]] = llvm.udiv %[[S8]], %[[WarpSize]]  : i32985// CHECK: %[[S11:.+]] = llvm.udiv %[[S9]], %[[S2]]  : i32986// CHECK: %[[S12:.+]] = llvm.urem %[[S9]], %[[S2]]  : i32987// CHECK: %[[S13:.+]] = llvm.mul %[[S12]], %[[S5]]  : i32988// CHECK: %[[S14:.+]] = llvm.mul %[[S10]], %[[S7]]  : i32989// CHECK: %[[S15:.+]] = llvm.add %[[S11]], %[[S14]]  : i32990// CHECK: %[[S16:.+]] = llvm.mlir.constant(0 : i32) : i32991// CHECK: %[[S17:.+]] = llvm.mul %[[S16]], %[[S4]]  : i32992// CHECK: %[[S18:.+]] = llvm.add %[[S15]], %[[S17]]  : i32993// CHECK: %[[S19:.+]] = llvm.mlir.constant(0 : i32) : i32994// CHECK: %[[S20:.+]] = llvm.mul %[[S19]], %[[S4]]  : i32995// CHECK: %[[S21:.+]] = llvm.add %[[S13]], %[[S20]]  : i32996// CHECK: %[[S22:.+]] = arith.index_cast %[[S18]] : i32 to index997// CHECK: %[[S23:.+]] = arith.index_cast %[[S21]] : i32 to index998// CHECK: %[[S24:.+]] = llvm.add %[[S21]], %[[S6]]  : i32999// CHECK: %[[S25:.+]] = arith.index_cast %[[S24]] : i32 to index1000// CHECK: %[[S26:.+]] = llvm.extractvalue %[[EX1]][0] : !llvm.struct1001// CHECK: %[[S27:.+]] = llvm.extractvalue %[[EX1]][1] : !llvm.struct1002// CHECK: memref.store %[[S26]], %[[arg2]][%[[S22]], %[[S23]]] : memref<128x128xf32, 3>1003// CHECK: memref.store %[[S27]], %[[arg2]][%[[S22]], %[[S25]]] : memref<128x128xf32, 3>1004 1005// ### Store {d2, d3} of each thread ###1006 1007// CHECK: %[[S28:.+]] = llvm.mlir.constant(1 : i32) : i321008// CHECK: %[[S29:.+]] = llvm.mul %[[S28]], %[[S4]]  : i321009// CHECK: %[[S30:.+]] = llvm.add %[[S13]], %[[S29]]  : i321010// CHECK: %[[S31:.+]] = arith.index_cast %[[S18]] : i32 to index1011// CHECK: %[[S32:.+]] = arith.index_cast %[[S30]] : i32 to index1012// CHECK: %[[S33:.+]] = llvm.add %[[S30]], %[[S6]]  : i321013// CHECK: %[[S34:.+]] = arith.index_cast %[[S33]] : i32 to index1014// CHECK: %[[S35:.+]] = llvm.extractvalue %[[EX1]][4] : !llvm.struct<1015// CHECK: %[[S36:.+]] = llvm.extractvalue %[[EX1]][5] : !llvm.struct<1016// CHECK: memref.store %[[S35]], %[[arg2]][%[[S31]], %[[S32]]] : memref<128x128xf32, 3>1017// CHECK: memref.store %[[S36]], %[[arg2]][%[[S31]], %[[S34]]] : memref<128x128xf32, 3>1018 1019// ### Store {d4, d5} of each thread ###1020 1021// CHECK: %[[S37:.+]] = llvm.mlir.constant(2 : i32) : i321022// CHECK: %[[S38:.+]] = llvm.mul %[[S37]], %[[S4]]  : i321023// CHECK: %[[S39:.+]] = llvm.add %[[S13]], %[[S38]]  : i321024// CHECK: %[[S40:.+]] = arith.index_cast %[[S18]] : i32 to index1025// CHECK: %[[S41:.+]] = arith.index_cast %[[S39]] : i32 to index1026// CHECK: %[[S42:.+]] = llvm.add %[[S39]], %[[S6]]  : i321027// CHECK: %[[S43:.+]] = arith.index_cast %[[S42]] : i32 to index1028// CHECK: %[[S44:.+]] = llvm.extractvalue %[[EX1]][8] : !llvm.struct<1029// CHECK: %[[S45:.+]] = llvm.extractvalue %[[EX1]][9] : !llvm.struct<1030// CHECK: memref.store %[[S44]], %[[arg2]][%[[S40]], %[[S41]]] : memref<128x128xf32, 3>1031// CHECK: memref.store %[[S45]], %[[arg2]][%[[S40]], %[[S43]]] : memref<128x128xf32, 3>1032 1033// ### Store {d6, d7} of each thread ###1034 1035// CHECK: %[[S46:.+]] = llvm.mlir.constant(3 : i32) : i321036// CHECK: %[[S47:.+]] = llvm.mul %[[S46]], %[[S4]]  : i321037// CHECK: %[[S48:.+]] = llvm.add %[[S13]], %[[S47]]  : i321038// CHECK: %[[S49:.+]] = arith.index_cast %[[S18]] : i32 to index1039// CHECK: %[[S50:.+]] = arith.index_cast %[[S48]] : i32 to index1040// CHECK: %[[S51:.+]] = llvm.add %[[S48]], %[[S6]]  : i321041// CHECK: %[[S52:.+]] = arith.index_cast %[[S51]] : i32 to index1042// CHECK: %[[S53:.+]] = llvm.extractvalue %[[EX1]][12] : !llvm.struct<1043// CHECK: %[[S54:.+]] = llvm.extractvalue %[[EX1]][13] : !llvm.struct<1044// CHECK: memref.store %[[S53]], %[[arg2]][%[[S49]], %[[S50]]] : memref<128x128xf32, 3>1045// CHECK: memref.store %[[S54]], %[[arg2]][%[[S49]], %[[S52]]] : memref<128x128xf32, 3>1046 1047// Pattern continues similarly 28x times until {... d62, d63}1048 1049// CHECK: %[[c1:.+]] = llvm.mlir.constant(1 : i32) : i321050// CHECK: %[[c2:.+]] = llvm.mlir.constant(2 : i32) : i321051 1052// ### Store {d64, d65} of each thread ###1053// CHECK: %[[EX2:.+]] = llvm.extractvalue %[[S0]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1054// CHECK: %[[S315:.+]] = llvm.mlir.constant(1 : i32) : i321055// CHECK: %[[S312:.+]] = llvm.mlir.constant(2 : i32) : i321056// CHECK: %[[S311:.+]] = llvm.mlir.constant(4 : i32) : i321057// CHECK: %[[S313:.+]] = llvm.mlir.constant(8 : i32) : i321058// CHECK: %[[S316:.+]] = llvm.mlir.constant(16 : i32) : i321059// CHECK: %[[WS2:.+]] = llvm.mlir.constant(32 : i32) : i321060// CHECK: %[[S317:.+]] = nvvm.read.ptx.sreg.tid.x : i321061// CHECK: %[[S318:.+]] = llvm.urem %[[S317]], %[[WS2]]  : i321062// CHECK: %[[S319:.+]] = llvm.udiv %[[S317]], %[[WS2]]  : i321063// CHECK: %[[S320:.+]] = llvm.udiv %[[S318]], %[[S311]]  : i321064// CHECK: %[[S321:.+]] = llvm.urem %[[S318]], %[[S311]]  : i321065// CHECK: %[[S322:.+]] = llvm.mul %[[S321]], %[[S312]]  : i321066// CHECK: %[[S323:.+]] = llvm.mul %[[S319]], %[[S316]]  : i321067// CHECK: %[[S324:.+]] = llvm.add %[[S320]], %[[S323]]  : i321068// CHECK: %[[S325:.+]] = llvm.mlir.constant(64 : i32) : i321069// CHECK: %[[S326:.+]] = llvm.add %[[S324]], %[[S325]]  : i321070// CHECK: %[[S327:.+]] = llvm.mlir.constant(0 : i32) : i321071// CHECK: %[[S328:.+]] = llvm.mul %[[S327]], %[[S313]]  : i321072// CHECK: %[[S329:.+]] = llvm.add %[[S326]], %[[S328]]  : i321073// CHECK: %[[S330:.+]] = llvm.mlir.constant(0 : i32) : i321074// CHECK: %[[S331:.+]] = llvm.mul %[[S330]], %[[S313]]  : i321075// CHECK: %[[S332:.+]] = llvm.add %[[S322]], %[[S331]]  : i321076// CHECK: %[[S333:.+]] = arith.index_cast %[[S329]] : i32 to index1077// CHECK: %[[S334:.+]] = arith.index_cast %[[S332]] : i32 to index1078// CHECK: %[[S335:.+]] = llvm.add %[[S332]], %[[S315]]  : i321079// CHECK: %[[S336:.+]] = arith.index_cast %[[S335]] : i32 to index1080// CHECK: %[[S337:.+]] = llvm.extractvalue %[[EX2]][0]1081// CHECK: %[[S338:.+]] = llvm.extractvalue %[[EX2]][1]1082// CHECK: memref.store %[[S337]], %[[arg2]][%[[S333]], %[[S334]]] : memref<128x128xf32, 3>1083// CHECK: memref.store %[[S338]], %[[arg2]][%[[S333]], %[[S336]]] : memref<128x128xf32, 3>1084 1085// Pattern continues similarly 31x times until {... d126, d127}1086 1087  nvgpu.warpgroup.mma.store %result, %matrixD :1088    !nvgpu.warpgroup.accumulator< fragmented = vector<128x128xf32>>1089    to memref<128x128xf32,3>1090  return1091}1092 1093// CHECK-LABEL: @warpgroup_mma_store_multiple1094func.func @warpgroup_mma_store_multiple(1095    %shmem_m64n8k : memref<64x8xf32>,1096    %shmem_m64n16k : memref<64x16xf32>,1097    %shmem_m64n24k : memref<64x24xf32>,1098    %shmem_m64n32k : memref<64x32xf32>,1099    %shmem_m64n40k : memref<64x40xf32>,1100    %shmem_m64n48k : memref<64x48xf32>,1101    %shmem_m64n56k : memref<64x56xf32>,1102    %shmem_m64n64k : memref<64x64xf32>,1103    %shmem_m64n72k : memref<64x72xf32>,1104    %shmem_m64n80k : memref<64x80xf32>,1105    %shmem_m64n88k : memref<64x88xf32>,1106    %shmem_m64n96k : memref<64x96xf32>,1107    %shmem_m64n104k : memref<64x104xf32>,1108    %shmem_m64n112k : memref<64x112xf32>,1109    %shmem_m64n120k : memref<64x120xf32>,1110    %shmem_m64n128k : memref<64x128xf32>,1111    %shmem_m64n136k : memref<64x136xf32>,1112    %shmem_m64n144k : memref<64x144xf32>,1113    %shmem_m64n152k : memref<64x152xf32>,1114    %shmem_m64n160k : memref<64x160xf32>,1115    %shmem_m64n168k : memref<64x168xf32>,1116    %shmem_m64n176k : memref<64x176xf32>,1117    %shmem_m64n184k : memref<64x184xf32>,1118    %shmem_m64n192k : memref<64x192xf32>,1119    %shmem_m64n200k : memref<64x200xf32>,1120    %shmem_m64n208k : memref<64x208xf32>,1121    %shmem_m64n216k : memref<64x216xf32>,1122    %shmem_m64n224k : memref<64x224xf32>,1123    %shmem_m64n232k : memref<64x232xf32>,1124    %shmem_m64n240k : memref<64x240xf32>,1125    %shmem_m64n248k : memref<64x248xf32>,1126    %shmem_m64n256k : memref<64x256xf32>,1127    %res_m64n16k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x16xf32>>,1128    %res_m64n24k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x24xf32>>,1129    %res_m64n32k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x32xf32>>,1130    %res_m64n40k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x40xf32>>,1131    %res_m64n48k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x48xf32>>,1132    %res_m64n56k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x56xf32>>,1133    %res_m64n64k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x64xf32>>,1134    %res_m64n72k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x72xf32>>,1135    %res_m64n80k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x80xf32>>,1136    %res_m64n88k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x88xf32>>,1137    %res_m64n96k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x96xf32>>,1138    %res_m64n104k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x104xf32>>,1139    %res_m64n112k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x112xf32>>,1140    %res_m64n120k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x120xf32>>,1141    %res_m64n128k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>,1142    %res_m64n136k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x136xf32>>,1143    %res_m64n144k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x144xf32>>,1144    %res_m64n152k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x152xf32>>,1145    %res_m64n160k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x160xf32>>,1146    %res_m64n168k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x168xf32>>,1147    %res_m64n176k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x176xf32>>,1148    %res_m64n184k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x184xf32>>,1149    %res_m64n192k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x192xf32>>,1150    %res_m64n200k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x200xf32>>,1151    %res_m64n208k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x208xf32>>,1152    %res_m64n216k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x216xf32>>,1153    %res_m64n224k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x224xf32>>,1154    %res_m64n232k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x232xf32>>,1155    %res_m64n240k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x240xf32>>,1156    %res_m64n248k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x248xf32>>,1157    %res_m64n256k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x256xf32>>) {1158    // CHECK-COUNT-8: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x16xf32>1159    // CHECK-COUNT-12: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x24xf32>1160    // CHECK-COUNT-16: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x32xf32>1161    // CHECK-COUNT-20: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x40xf32>1162    // CHECK-COUNT-24: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x48xf32>1163    // CHECK-COUNT-28: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x56xf32>1164    // CHECK-COUNT-32: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x64xf32>1165    // CHECK-COUNT-36: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x72xf32>1166    // CHECK-COUNT-40: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x80xf32>1167    // CHECK-COUNT-44: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x88xf32>1168    // CHECK-COUNT-48: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x96xf32>1169    // CHECK-COUNT-52: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x104xf32>1170    // CHECK-COUNT-56: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x112xf32>1171    // CHECK-COUNT-60: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x120xf32>1172    // CHECK-COUNT-64: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x128xf32>1173    // CHECK-COUNT-68: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x136xf32>1174    // CHECK-COUNT-72: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x144xf32>1175    // CHECK-COUNT-76: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x152xf32>1176    // CHECK-COUNT-80: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x160xf32>1177    // CHECK-COUNT-84: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x168xf32>1178    // CHECK-COUNT-88: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x176xf32>1179    // CHECK-COUNT-92: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x184xf32>1180    // CHECK-COUNT-96: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x192xf32>1181    // CHECK-COUNT-100: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x200xf32>1182    // CHECK-COUNT-104: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x208xf32>1183    // CHECK-COUNT-108: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x216xf32>1184    // CHECK-COUNT-112: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x224xf32>1185    // CHECK-COUNT-116: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x232xf32>1186    // CHECK-COUNT-120: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x240xf32>1187    // CHECK-COUNT-124: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x248xf32>1188    // CHECK-COUNT-128: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x256xf32>1189    nvgpu.warpgroup.mma.store  %res_m64n16k, %shmem_m64n16k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x16xf32>> to memref<64x16xf32>1190    nvgpu.warpgroup.mma.store  %res_m64n24k, %shmem_m64n24k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x24xf32>> to memref<64x24xf32>1191    nvgpu.warpgroup.mma.store  %res_m64n32k, %shmem_m64n32k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x32xf32>> to memref<64x32xf32>1192    nvgpu.warpgroup.mma.store  %res_m64n40k, %shmem_m64n40k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x40xf32>> to memref<64x40xf32>1193    nvgpu.warpgroup.mma.store  %res_m64n48k, %shmem_m64n48k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x48xf32>> to memref<64x48xf32>1194    nvgpu.warpgroup.mma.store  %res_m64n56k, %shmem_m64n56k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x56xf32>> to memref<64x56xf32>1195    nvgpu.warpgroup.mma.store  %res_m64n64k, %shmem_m64n64k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x64xf32>> to memref<64x64xf32>1196    nvgpu.warpgroup.mma.store  %res_m64n72k, %shmem_m64n72k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x72xf32>> to memref<64x72xf32>1197    nvgpu.warpgroup.mma.store  %res_m64n80k, %shmem_m64n80k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x80xf32>> to memref<64x80xf32>1198    nvgpu.warpgroup.mma.store  %res_m64n88k, %shmem_m64n88k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x88xf32>> to memref<64x88xf32>1199    nvgpu.warpgroup.mma.store  %res_m64n96k, %shmem_m64n96k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x96xf32>> to memref<64x96xf32>1200    nvgpu.warpgroup.mma.store  %res_m64n104k, %shmem_m64n104k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x104xf32>> to memref<64x104xf32>1201    nvgpu.warpgroup.mma.store  %res_m64n112k, %shmem_m64n112k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x112xf32>> to memref<64x112xf32>1202    nvgpu.warpgroup.mma.store  %res_m64n120k, %shmem_m64n120k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x120xf32>> to memref<64x120xf32>1203    nvgpu.warpgroup.mma.store  %res_m64n128k, %shmem_m64n128k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>> to memref<64x128xf32>1204    nvgpu.warpgroup.mma.store  %res_m64n136k, %shmem_m64n136k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x136xf32>> to memref<64x136xf32>1205    nvgpu.warpgroup.mma.store  %res_m64n144k, %shmem_m64n144k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x144xf32>> to memref<64x144xf32>1206    nvgpu.warpgroup.mma.store  %res_m64n152k, %shmem_m64n152k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x152xf32>> to memref<64x152xf32>1207    nvgpu.warpgroup.mma.store  %res_m64n160k, %shmem_m64n160k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x160xf32>> to memref<64x160xf32>1208    nvgpu.warpgroup.mma.store  %res_m64n168k, %shmem_m64n168k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x168xf32>> to memref<64x168xf32>1209    nvgpu.warpgroup.mma.store  %res_m64n176k, %shmem_m64n176k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x176xf32>> to memref<64x176xf32>1210    nvgpu.warpgroup.mma.store  %res_m64n184k, %shmem_m64n184k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x184xf32>> to memref<64x184xf32>1211    nvgpu.warpgroup.mma.store  %res_m64n192k, %shmem_m64n192k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x192xf32>> to memref<64x192xf32>1212    nvgpu.warpgroup.mma.store  %res_m64n200k, %shmem_m64n200k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x200xf32>> to memref<64x200xf32>1213    nvgpu.warpgroup.mma.store  %res_m64n208k, %shmem_m64n208k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x208xf32>> to memref<64x208xf32>1214    nvgpu.warpgroup.mma.store  %res_m64n216k, %shmem_m64n216k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x216xf32>> to memref<64x216xf32>1215    nvgpu.warpgroup.mma.store  %res_m64n224k, %shmem_m64n224k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x224xf32>> to memref<64x224xf32>1216    nvgpu.warpgroup.mma.store  %res_m64n232k, %shmem_m64n232k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x232xf32>> to memref<64x232xf32>1217    nvgpu.warpgroup.mma.store  %res_m64n240k, %shmem_m64n240k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x240xf32>> to memref<64x240xf32>1218    nvgpu.warpgroup.mma.store  %res_m64n248k, %shmem_m64n248k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x248xf32>> to memref<64x248xf32>1219    nvgpu.warpgroup.mma.store  %res_m64n256k, %shmem_m64n256k : !nvgpu.warpgroup.accumulator<fragmented = vector<64x256xf32>> to memref<64x256xf32>1220  return1221}1222 1223func.func @warpgroup_mma_init() {1224  //CHECK: %[[S1:.+]] = llvm.mlir.constant(0.000000e+00 : f32) : f31225  //CHECK: %[[S0:.+]] = llvm.mlir.poison : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1226  //CHECK: %[[EX:.+]] = llvm.extractvalue %[[S0]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1227  //CHECK: %[[S2:.+]] = llvm.insertvalue %[[S1]], %[[EX]][0] : !llvm.struct1228  //CHECK: %[[S3:.+]] = llvm.insertvalue %[[S1]], %[[S2]][1] : !llvm.struct1229  //CHECK: %[[S4:.+]] = llvm.insertvalue %[[S1]], %[[S3]][2] : !llvm.struct1230  //CHECK: %[[S5:.+]] = llvm.insertvalue %[[S1]], %[[S4]][3] : !llvm.struct1231  //CHECK: %[[S6:.+]] = llvm.insertvalue %[[S1]], %[[S5]][4] : !llvm.struct1232  //CHECK: %[[S7:.+]] = llvm.insertvalue %[[S1]], %[[S6]][5] : !llvm.struct1233  //CHECK: %[[S8:.+]] = llvm.insertvalue %[[S1]], %[[S7]][6] : !llvm.struct1234  //CHECK: %[[S9:.+]] = llvm.insertvalue %[[S1]], %[[S8]][7] : !llvm.struct1235  //CHECK: %[[S10:.+]] = llvm.insertvalue %[[S1]], %[[S9]][8] : !llvm.struct1236  //CHECK: %[[S11:.+]] = llvm.insertvalue %[[S1]], %[[S10]][9] : !llvm.struct1237  //CHECK: %[[S12:.+]] = llvm.insertvalue %[[S1]], %[[S11]][10] : !llvm.struct1238  //CHECK: %[[S13:.+]] = llvm.insertvalue %[[S1]], %[[S12]][11] : !llvm.struct1239  //CHECK: %[[S14:.+]] = llvm.insertvalue %[[S1]], %[[S13]][12] : !llvm.struct1240  //CHECK: %[[S15:.+]] = llvm.insertvalue %[[S1]], %[[S14]][13] : !llvm.struct1241  //CHECK: %[[S16:.+]] = llvm.insertvalue %[[S1]], %[[S15]][14] : !llvm.struct1242  //CHECK: %[[S17:.+]] = llvm.insertvalue %[[S1]], %[[S16]][15] : !llvm.struct1243  //CHECK: %[[S18:.+]] = llvm.insertvalue %[[S1]], %[[S17]][16] : !llvm.struct1244  //CHECK: %[[S19:.+]] = llvm.insertvalue %[[S1]], %[[S18]][17] : !llvm.struct1245  //CHECK: %[[S20:.+]] = llvm.insertvalue %[[S1]], %[[S19]][18] : !llvm.struct1246  //CHECK: %[[S21:.+]] = llvm.insertvalue %[[S1]], %[[S20]][19] : !llvm.struct1247  //CHECK: %[[S22:.+]] = llvm.insertvalue %[[S1]], %[[S21]][20] : !llvm.struct1248  //CHECK: %[[S23:.+]] = llvm.insertvalue %[[S1]], %[[S22]][21] : !llvm.struct1249  //CHECK: %[[S24:.+]] = llvm.insertvalue %[[S1]], %[[S23]][22] : !llvm.struct1250  //CHECK: %[[S25:.+]] = llvm.insertvalue %[[S1]], %[[S24]][23] : !llvm.struct1251  //CHECK: %[[S26:.+]] = llvm.insertvalue %[[S1]], %[[S25]][24] : !llvm.struct1252  //CHECK: %[[S27:.+]] = llvm.insertvalue %[[S1]], %[[S26]][25] : !llvm.struct1253  //CHECK: %[[S28:.+]] = llvm.insertvalue %[[S1]], %[[S27]][26] : !llvm.struct1254  //CHECK: %[[S29:.+]] = llvm.insertvalue %[[S1]], %[[S28]][27] : !llvm.struct1255  //CHECK: %[[S30:.+]] = llvm.insertvalue %[[S1]], %[[S29]][28] : !llvm.struct1256  //CHECK: %[[S31:.+]] = llvm.insertvalue %[[S1]], %[[S30]][29] : !llvm.struct1257  //CHECK: %[[S32:.+]] = llvm.insertvalue %[[S1]], %[[S31]][30] : !llvm.struct1258  //CHECK: %[[S33:.+]] = llvm.insertvalue %[[S1]], %[[S32]][31] : !llvm.struct1259  //CHECK: %[[S34:.+]] = llvm.insertvalue %[[S1]], %[[S33]][32] : !llvm.struct1260  //CHECK: %[[S35:.+]] = llvm.insertvalue %[[S1]], %[[S34]][33] : !llvm.struct1261  //CHECK: %[[S36:.+]] = llvm.insertvalue %[[S1]], %[[S35]][34] : !llvm.struct1262  //CHECK: %[[S37:.+]] = llvm.insertvalue %[[S1]], %[[S36]][35] : !llvm.struct1263  //CHECK: %[[S38:.+]] = llvm.insertvalue %[[S1]], %[[S37]][36] : !llvm.struct1264  //CHECK: %[[S39:.+]] = llvm.insertvalue %[[S1]], %[[S38]][37] : !llvm.struct1265  //CHECK: %[[S40:.+]] = llvm.insertvalue %[[S1]], %[[S39]][38] : !llvm.struct1266  //CHECK: %[[S41:.+]] = llvm.insertvalue %[[S1]], %[[S40]][39] : !llvm.struct1267  //CHECK: %[[S42:.+]] = llvm.insertvalue %[[S1]], %[[S41]][40] : !llvm.struct1268  //CHECK: %[[S43:.+]] = llvm.insertvalue %[[S1]], %[[S42]][41] : !llvm.struct1269  //CHECK: %[[S44:.+]] = llvm.insertvalue %[[S1]], %[[S43]][42] : !llvm.struct1270  //CHECK: %[[S45:.+]] = llvm.insertvalue %[[S1]], %[[S44]][43] : !llvm.struct1271  //CHECK: %[[S46:.+]] = llvm.insertvalue %[[S1]], %[[S45]][44] : !llvm.struct1272  //CHECK: %[[S47:.+]] = llvm.insertvalue %[[S1]], %[[S46]][45] : !llvm.struct1273  //CHECK: %[[S48:.+]] = llvm.insertvalue %[[S1]], %[[S47]][46] : !llvm.struct1274  //CHECK: %[[S49:.+]] = llvm.insertvalue %[[S1]], %[[S48]][47] : !llvm.struct1275  //CHECK: %[[S50:.+]] = llvm.insertvalue %[[S1]], %[[S49]][48] : !llvm.struct1276  //CHECK: %[[S51:.+]] = llvm.insertvalue %[[S1]], %[[S50]][49] : !llvm.struct1277  //CHECK: %[[S52:.+]] = llvm.insertvalue %[[S1]], %[[S51]][50] : !llvm.struct1278  //CHECK: %[[S53:.+]] = llvm.insertvalue %[[S1]], %[[S52]][51] : !llvm.struct1279  //CHECK: %[[S54:.+]] = llvm.insertvalue %[[S1]], %[[S53]][52] : !llvm.struct1280  //CHECK: %[[S55:.+]] = llvm.insertvalue %[[S1]], %[[S54]][53] : !llvm.struct1281  //CHECK: %[[S56:.+]] = llvm.insertvalue %[[S1]], %[[S55]][54] : !llvm.struct1282  //CHECK: %[[S57:.+]] = llvm.insertvalue %[[S1]], %[[S56]][55] : !llvm.struct1283  //CHECK: %[[S58:.+]] = llvm.insertvalue %[[S1]], %[[S57]][56] : !llvm.struct1284  //CHECK: %[[S59:.+]] = llvm.insertvalue %[[S1]], %[[S58]][57] : !llvm.struct1285  //CHECK: %[[S60:.+]] = llvm.insertvalue %[[S1]], %[[S59]][58] : !llvm.struct1286  //CHECK: %[[S61:.+]] = llvm.insertvalue %[[S1]], %[[S60]][59] : !llvm.struct1287  //CHECK: %[[S62:.+]] = llvm.insertvalue %[[S1]], %[[S61]][60] : !llvm.struct1288  //CHECK: %[[S63:.+]] = llvm.insertvalue %[[S1]], %[[S62]][61] : !llvm.struct1289  //CHECK: %[[S64:.+]] = llvm.insertvalue %[[S1]], %[[S63]][62] : !llvm.struct1290  //CHECK: %[[S65:.+]] = llvm.insertvalue %[[S1]], %[[S64]][63] : !llvm.struct1291  %matrixC = nvgpu.warpgroup.mma.init.accumulator -> !nvgpu.warpgroup.accumulator< fragmented = vector<128x128xf32>>1292  return1293}1294 1295// CHECK-LABEL: @warpgroup_matrix_multiply_m128n128k64(1296// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>, %[[arg1:[a-zA-Z0-9_]+]]: !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>, %[[arg2:[a-zA-Z0-9_]+]]: memref<128x128xf32, 3>)1297func.func @warpgroup_matrix_multiply_m128n128k64(1298      %descA: !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>,1299      %descB: !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>,1300      %shmemD: memref<128x128xf32, 3>)1301{1302  // Init1303  %matrixC = nvgpu.warpgroup.mma.init.accumulator ->1304                      !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>1305 1306  // GEMM1307  %matrixD = nvgpu.warpgroup.mma %descA, %descB, %matrixC {transposeB}:1308        !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>,1309        !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>,1310        !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>1311        ->1312        !nvgpu.warpgroup.accumulator<fragmented = vector<128x128xf32>>1313 1314 1315  // Epilogue1316  nvgpu.warpgroup.mma.store %matrixD, %shmemD :1317    !nvgpu.warpgroup.accumulator< fragmented = vector<128x128xf32>>1318    to memref<128x128xf32,3>1319 1320 1321// CHECK-DAG: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>> to i641322// CHECK-DAG: %[[S1:.+]] = builtin.unrealized_conversion_cast %[[arg1]] : !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>> to i641323// CHECK: %[[S3:.+]] = llvm.mlir.constant(0.000000e+00 : f32) : f321324// CHECK: %[[S4:.+]] = llvm.mlir.poison : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1325// CHECK: %[[S5:.+]] = llvm.extractvalue %[[S4]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1326// CHECK: %[[S6:.+]] = llvm.insertvalue %[[S3]], %[[S5]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1327// CHECK: %[[S68:.+]] = llvm.insertvalue %[[S3]], %{{.*}}[63] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1328// CHECK: %[[S69:.+]] = llvm.extractvalue %[[S4]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1329// CHECK: %[[S71:.+]] = llvm.insertvalue %[[S3]], %[[S69]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1330// CHECK: %[[S134:.+]] = llvm.insertvalue %[[S3]], %{{.*}}[63] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1331// CHECK: %[[S135:.+]] = llvm.insertvalue %[[S68]], %[[S4]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1332// CHECK: %[[S136:.+]] = llvm.insertvalue %[[S134]], %[[S135]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1333// CHECK: nvvm.wgmma.fence.aligned1334// CHECK: %[[S137:.+]] = llvm.mlir.poison : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1335// CHECK: %[[S138:.+]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1336// CHECK: %[[S139:.+]] = nvvm.wgmma.mma_async %[[S0]], %[[S1]], %[[S138]], <m = 64, n = 128, k = 16>, D[<f32>, <one>, <wrapped>], A[<f16>, <one>, <row>], B[<f16>, <one>, <row>] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)> -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1337// CHECK: nvvm.wgmma.mma_async1338// CHECK: nvvm.wgmma.mma_async1339// CHECK: %[[S154:.+]] = nvvm.wgmma.mma_async1340// CHECK: nvvm.wgmma.mma_async1341// CHECK: nvvm.wgmma.mma_async1342// CHECK: nvvm.wgmma.mma_async1343// CHECK: %[[S173:.+]] = nvvm.wgmma.mma_async1344// CHECK: %[[S174:.+]] = llvm.insertvalue %[[S154]], %[[S137]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1345// CHECK: %[[S175:.+]] = llvm.insertvalue %[[S173]], %[[S174]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1346// CHECK: nvvm.wgmma.commit.group.sync.aligned1347// CHECK: nvvm.wgmma.wait.group.sync.aligned 11348// CHECK: %[[S176:.+]] = llvm.extractvalue %[[S175]][0] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1349// CHECK: %[[S193:.+]] = llvm.extractvalue %[[S176]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1350// CHECK: %[[S194:.+]] = llvm.extractvalue %[[S176]][1] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1351// CHECK: memref.store %[[S193]], %[[arg2]][%{{.*}}, %{{.*}}] : memref<128x128xf32, 3>1352// CHECK: memref.store %[[S194]], %[[arg2]][%{{.*}}, %{{.*}}] : memref<128x128xf32, 3>1353// CHECK: %[[S485:.+]] = llvm.extractvalue %[[S175]][1] : !llvm.struct<(struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>, struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>)>1354// CHECK: %[[S503:.+]] = llvm.extractvalue %[[S485]][0] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1355// CHECK: %[[S504:.+]] = llvm.extractvalue %[[S485]][1] : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32)>1356// CHECK: memref.store %[[S503]], %[[arg2]][%{{.*}}, %{{.*}}] : memref<128x128xf32, 3>1357// CHECK: memref.store %[[S504]], %[[arg2]][%{{.*}}, %{{.*}}] : memref<128x128xf32, 3>1358  return1359}1360 1361 1362module attributes {transform.with_named_sequence} {1363  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1364    %0 = transform.structured.match ops{["func.func"]} in %arg11365      : (!transform.any_op) -> !transform.any_op1366    transform.apply_conversion_patterns to %0 {1367      transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm1368    } with type_converter {1369      transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter1370        {use_opaque_pointers = true}1371    } {legal_dialects = ["arith", "func", "llvm", "memref", "nvvm", "vector", "scf"], partial_conversion} : !transform.any_op1372    transform.yield1373  }1374}1375 1376// CHECK-LABEL: @rcp_approx_ftz_f321377// CHECK-SAME:  %[[IN:.*]]: vector<32x16xf32>1378func.func @rcp_approx_ftz_f32(%in: vector<32x16xf32>) {1379  // CHECK: %[[IN_LLVM:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<32x16xf32> to !llvm.array<32 x vector<16xf32>>1380  // CHECK: %[[IN1DVEC:.*]] = llvm.extractvalue %[[IN_LLVM]][0] : !llvm.array<32 x vector<16xf32>>1381  // CHECK: %[[OUT1DVEC:.*]] = llvm.mlir.poison : vector<16xf32>1382  // CHECK: %[[IDX_0:.+]] = llvm.mlir.constant(0 : i64) : i641383  // CHECK: %[[ELEM_0:.*]] = llvm.extractelement %[[IN1DVEC]][%[[IDX_0]] : i64]1384  // CHECK: %[[ELEM_RCP0:.*]] = nvvm.rcp.approx.ftz.f %[[ELEM_0]] : f321385  // CHECK: llvm.insertelement %[[ELEM_RCP0]], %[[OUT1DVEC]][%[[IDX_0]] : i64] : vector<16xf32>1386  // CHECK-COUNT-511: nvvm.rcp.approx.ftz.f1387  %out = nvgpu.rcp %in {rounding = approx, ftz} : vector<32x16xf32>1388  return1389}1390