mlir/test/Dialect/LLVMIR/nvvm.mlir

   1 // RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
   2
   3 // CHECK-LABEL: @nvvm_special_regs
   4 func.func @nvvm_special_regs() -> i32 {
   5   // CHECK: nvvm.read.ptx.sreg.tid.x : i32
   6   %0 = nvvm.read.ptx.sreg.tid.x : i32
   7   // CHECK: nvvm.read.ptx.sreg.tid.y : i32
   8   %1 = nvvm.read.ptx.sreg.tid.y : i32
   9   // CHECK: nvvm.read.ptx.sreg.tid.z : i32
  10   %2 = nvvm.read.ptx.sreg.tid.z : i32
  11   // CHECK: nvvm.read.ptx.sreg.ntid.x : i32
  12   %3 = nvvm.read.ptx.sreg.ntid.x : i32
  13   // CHECK: nvvm.read.ptx.sreg.ntid.y : i32
  14   %4 = nvvm.read.ptx.sreg.ntid.y : i32
  15   // CHECK: nvvm.read.ptx.sreg.ntid.z : i32
  16   %5 = nvvm.read.ptx.sreg.ntid.z : i32
  17   // CHECK: nvvm.read.ptx.sreg.ctaid.x : i32
  18   %6 = nvvm.read.ptx.sreg.ctaid.x : i32
  19   // CHECK: nvvm.read.ptx.sreg.ctaid.y : i32
  20   %7 = nvvm.read.ptx.sreg.ctaid.y : i32
  21   // CHECK: nvvm.read.ptx.sreg.ctaid.z : i32
  22   %8 = nvvm.read.ptx.sreg.ctaid.z : i32
  23   // CHECK: nvvm.read.ptx.sreg.nctaid.x : i32
  24   %9 = nvvm.read.ptx.sreg.nctaid.x : i32
  25   // CHECK: nvvm.read.ptx.sreg.nctaid.y : i32
  26   %10 = nvvm.read.ptx.sreg.nctaid.y : i32
  27   // CHECK: nvvm.read.ptx.sreg.nctaid.z : i32
  28   %11 = nvvm.read.ptx.sreg.nctaid.z : i32
  29   llvm.return %0 : i32
  30 }
  31
  32 // CHECK-LABEL: @nvvm_rcp
  33 func.func @nvvm_rcp(%arg0: f32) -> f32 {
  34   // CHECK: nvvm.rcp.approx.ftz.f %arg0 : f32
  35   %0 = nvvm.rcp.approx.ftz.f %arg0 : f32
  36   llvm.return %0 : f32
  37 }
  38
  39 // CHECK-LABEL: @llvm_nvvm_barrier0
  40 func.func @llvm_nvvm_barrier0() {
  41   // CHECK: nvvm.barrier0
  42   nvvm.barrier0
  43   llvm.return
  44 }
  45
  46 // CHECK-LABEL: @llvm_nvvm_barrier
  47 // CHECK-SAME: (%[[barId:.*]]: i32, %[[numberOfThreads:.*]]: i32)
  48 llvm.func @llvm_nvvm_barrier(%barId : i32, %numberOfThreads : i32) {
  49   // CHECK: nvvm.barrier
  50   nvvm.barrier
  51   // CHECK: nvvm.barrier id = %[[barId]]
  52   nvvm.barrier id = %barId
  53   // CHECK: nvvm.barrier id = %[[barId]] number_of_threads = %[[numberOfThreads]]
  54   nvvm.barrier id = %barId number_of_threads = %numberOfThreads
  55   llvm.return
  56 }
  57
  58 // CHECK-LABEL: @llvm_nvvm_barrier_arrive
  59 // CHECK-SAME: (%[[barId:.*]]: i32, %[[numberOfThreads:.*]]: i32)
  60 llvm.func @llvm_nvvm_barrier_arrive(%barId : i32, %numberOfThreads : i32) {
  61   // CHECK: nvvm.barrier.arrive number_of_threads = %[[numberOfThreads]]
  62   nvvm.barrier.arrive number_of_threads = %numberOfThreads
  63   // CHECK: nvvm.barrier.arrive id = %[[barId]] number_of_threads = %[[numberOfThreads]]
  64   nvvm.barrier.arrive id = %barId number_of_threads = %numberOfThreads
  65   llvm.return
  66 }
  67
  68 // CHECK-LABEL: @llvm_nvvm_cluster_arrive
  69 func.func @llvm_nvvm_cluster_arrive() {
  70   // CHECK: nvvm.cluster.arrive
  71   nvvm.cluster.arrive
  72   // CHECK: nvvm.cluster.arrive {aligned}
  73   nvvm.cluster.arrive {aligned}
  74   llvm.return
  75 }
  76
  77 // CHECK-LABEL: @llvm_nvvm_cluster_arrive_relaxed
  78 func.func @llvm_nvvm_cluster_arrive_relaxed() {
  79   // CHECK: nvvm.cluster.arrive.relaxed
  80   nvvm.cluster.arrive.relaxed
  81   // CHECK: nvvm.cluster.arrive.relaxed {aligned}
  82   nvvm.cluster.arrive.relaxed {aligned}
  83   llvm.return
  84 }
  85
  86 // CHECK-LABEL: @llvm_nvvm_cluster_wait
  87 func.func @llvm_nvvm_cluster_wait() {
  88   // CHECK: nvvm.cluster.wait
  89   nvvm.cluster.wait
  90   // CHECK: nvvm.cluster.wait {aligned}
  91   nvvm.cluster.wait {aligned}
  92   llvm.return
  93 }
  94
  95 // CHECK-LABEL: @llvm_nvvm_fence_sc_cluster
  96 func.func @llvm_nvvm_fence_sc_cluster() {
  97   // CHECK: nvvm.fence.sc.cluster
  98   nvvm.fence.sc.cluster
  99   llvm.return
 100 }
 101
 102 // CHECK-LABEL: @nvvm_shfl
 103 func.func @nvvm_shfl(
 104     %arg0 : i32, %arg1 : i32, %arg2 : i32,
 105     %arg3 : i32, %arg4 : f32) -> i32 {
 106   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32 -> i32
 107   %0 = nvvm.shfl.sync bfly %arg0, %arg3, %arg1, %arg2 : i32 -> i32
 108   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 -> f32
 109   %1 = nvvm.shfl.sync bfly %arg0, %arg4, %arg1, %arg2 : f32 -> f32
 110   // CHECK: nvvm.shfl.sync up %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 -> f32
 111   %2 = nvvm.shfl.sync up %arg0, %arg4, %arg1, %arg2 : f32 -> f32
 112   // CHECK: nvvm.shfl.sync down %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 -> f32
 113   %3 = nvvm.shfl.sync down %arg0, %arg4, %arg1, %arg2 : f32 -> f32
 114   // CHECK: nvvm.shfl.sync idx %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : f32 -> f32
 115   %4 = nvvm.shfl.sync idx %arg0, %arg4, %arg1, %arg2 : f32 -> f32
 116   llvm.return %0 : i32
 117 }
 118
 119 // CHECK-LABEL: @nvvm_shfl_pred
 120 func.func @nvvm_shfl_pred(
 121     %arg0 : i32, %arg1 : i32, %arg2 : i32,
 122     %arg3 : i32, %arg4 : f32) -> !llvm.struct<(i32, i1)> {
 123   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {return_value_and_is_valid} : i32 -> !llvm.struct<(i32, i1)>
 124   %0 = nvvm.shfl.sync bfly %arg0, %arg3, %arg1, %arg2 {return_value_and_is_valid} : i32 -> !llvm.struct<(i32, i1)>
 125   // CHECK: nvvm.shfl.sync bfly %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>
 126   %1 = nvvm.shfl.sync bfly %arg0, %arg4, %arg1, %arg2 {return_value_and_is_valid} : f32 -> !llvm.struct<(f32, i1)>
 127   llvm.return %0 : !llvm.struct<(i32, i1)>
 128 }
 129
 130 // CHECK-LABEL: @nvvm_vote(
 131 func.func @nvvm_vote(%arg0 : i32, %arg1 : i1) -> i32 {
 132   // CHECK: nvvm.vote.ballot.sync %{{.*}}, %{{.*}} : i32
 133   %0 = nvvm.vote.ballot.sync %arg0, %arg1 : i32
 134   llvm.return %0 : i32
 135 }
 136
 137 // CHECK-LABEL: @llvm_nvvm_bar_warp_sync
 138 func.func @llvm_nvvm_bar_warp_sync(%mask : i32) {
 139   // CHECK: nvvm.bar.warp.sync %{{.*}}
 140   nvvm.bar.warp.sync %mask : i32
 141   llvm.return
 142 }
 143
 144 // CHECK-LABEL: @nvvm_mma_m8n8k4_row_col_f32_f32
 145 func.func @nvvm_mma_m8n8k4_row_col_f32_f32(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 146                %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 147                %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32, %c4 : f32, %c5 : f32, %c6 : f32, %c7 : f32) {
 148   // CHECK: nvvm.mma.sync
 149   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7]
 150     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 151      shape = #nvvm.shape<m = 8, n = 8, k = 4>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
 152   llvm.return %0 : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
 153 }
 154
 155 // CHECK-LABEL: @nvvm_mma_m8n8k4_f16_f16
 156 func.func @nvvm_mma_m8n8k4_f16_f16(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 157                               %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 158                               %c0 : vector<2xf16>, %c1 : vector<2xf16>, %c2 : vector<2xf16>, %c3 : vector<2xf16>) {
 159   // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}]
 160   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
 161     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 162      shape = #nvvm.shape<m = 8, n = 8, k = 4>} : (vector<2xf16>,vector<2xf16>,vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
 163   llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
 164 }
 165
 166 // CHECK-LABEL: @nvvm_mma_m8n8k16_s8_s8
 167 func.func @nvvm_mma_m8n8k16_s8_s8(%a0 : i32, %b0 : i32,
 168                              %c0 : i32, %c1 : i32) {
 169   // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>, shape = #nvvm.shape<m = 8, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32, i32)>
 170   %0 = nvvm.mma.sync A[%a0] B[%b0] C[%c0, %c1]
 171     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 172      multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>,
 173      intOverflowBehavior=#nvvm.mma_int_overflow<wrapped>,
 174      shape = #nvvm.shape<m = 8, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32, i32)>
 175   llvm.return %0 : !llvm.struct<(i32, i32)>
 176 }
 177
 178 // CHECK-LABEL: @nvvm_mma_m16n8k8_f16_f16
 179 func.func @nvvm_mma_m16n8k8_f16_f16(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 180                                %b0 : vector<2xf16>,
 181                                %c0 : vector<2xf16>, %c1 : vector<2xf16>) {
 182   // CHECK: nvvm.mma.sync A[%{{.*}}, %{{.*}}] B[%{{.*}}] C[%{{.*}}, %{{.*}}] {{{.*}}} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 183   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1]
 184     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 185      shape = #nvvm.shape<m = 16, n = 8, k = 8>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 186   llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 187 }
 188
 189 // CHECK-LABEL: @nvvm_mma_m16n8k16_f16_f16
 190 func.func @nvvm_mma_m16n8k16_f16_f16(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 191                                 %a2 : vector<2xf16>, %a3 : vector<2xf16>,
 192                                 %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 193                                 %c0 : vector<2xf16>, %c1 : vector<2xf16>) {
 194   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 195   %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1]
 196     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 197      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 198   llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 199 }
 200
 201 // CHECK-LABEL: @nvvm_mma_m16n8k16_f32_f16
 202 func.func @nvvm_mma_m16n8k16_f32_f16(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 203                                 %a2 : vector<2xf16>, %a3 : vector<2xf16>,
 204                                 %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 205                                 %c0 : vector<2xf16>, %c1 : vector<2xf16>) {
 206   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(f32, f32, f32, f32)>
 207   %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1]
 208     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 209      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>,vector<2xf16>,vector<2xf16>) -> !llvm.struct<(f32, f32, f32, f32)>
 210   llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
 211 }
 212
 213 // CHECK-LABEL: @nvvm_mma_m16n8k16_f16_f32
 214 func.func @nvvm_mma_m16n8k16_f16_f32(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 215                                 %a2 : vector<2xf16>, %a3 : vector<2xf16>,
 216                                 %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 217                                 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) {
 218   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 219   %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
 220     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 221      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 222   llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
 223 }
 224
 225 // CHECK-LABEL: @nvvm_mma_m16n8k16_f32_f32
 226 func.func @nvvm_mma_m16n8k16_f32_f32(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
 227                                 %a2 : vector<2xf16>, %a3 : vector<2xf16>,
 228                                 %b0 : vector<2xf16>, %b1 : vector<2xf16>,
 229                                 %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) {
 230   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}, {{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>
 231   %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
 232     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 233      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32) -> !llvm.struct<(f32, f32, f32, f32)>
 234   llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
 235 }
 236
 237 // CHECK-LABEL: @nvvm_mma_m16n8k4_tf32_f32
 238 func.func @nvvm_mma_m16n8k4_tf32_f32(%a0 : i32, %a1 : i32,
 239                                      %b0 : i32,
 240                                      %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) {
 241   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>, shape = #nvvm.shape<m = 16, n = 8, k = 4>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
 242   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
 243     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 244      multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>,
 245      shape = #nvvm.shape<m = 16, n = 8, k = 4>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
 246   llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
 247 }
 248
 249 // CHECK-LABEL: @nvvm_mma_m16n8k16_s8_s8
 250 func.func @nvvm_mma_m16n8k16_s8_s8(%a0 : i32, %a1 : i32, %b0 : i32,
 251                               %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32) {
 252   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>
 253   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
 254     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 255      multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>,
 256      intOverflowBehavior=#nvvm.mma_int_overflow<wrapped>,
 257      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32,i32,i32,i32)>
 258   llvm.return %0 : !llvm.struct<(i32,i32,i32,i32)>
 259 }
 260
 261 // CHECK-LABEL: @nvvm_mma_m16n8k16_s8_u8
 262 func.func @nvvm_mma_m16n8k16_s8_u8(%a0 : i32, %a1 : i32,
 263                                 %b0 : i32,
 264                                 %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32) {
 265   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<u8>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>
 266   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
 267     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 268      multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<u8>,
 269      intOverflowBehavior=#nvvm.mma_int_overflow<satfinite>,
 270      shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, i32) -> !llvm.struct<(i32,i32,i32,i32)>
 271   llvm.return %0 : !llvm.struct<(i32,i32,i32,i32)>
 272 }
 273
 274 // CHECK-LABEL: @nvvm_mma_m16n8k256_b1_b1
 275 func.func @nvvm_mma_m16n8k256_b1_b1(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,
 276                                %b0 : i32, %b1 : i32,
 277                                %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32) {
 278   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}, {{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {b1Op = #nvvm.mma_b1op<xor_popc>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>, shape = #nvvm.shape<m = 16, n = 8, k = 256>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>
 279   %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
 280     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 281      multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>,
 282      b1Op = #nvvm.mma_b1op<xor_popc>, shape = #nvvm.shape<m = 16, n = 8, k = 256>} : (i32, i32, i32) -> !llvm.struct<(i32,i32,i32,i32)>
 283   llvm.return %0 : !llvm.struct<(i32,i32,i32,i32)>
 284 }
 285
 286 // CHECK-LABEL: @nvvm_mma_m16n8k128_b1_b1
 287 func.func @nvvm_mma_m16n8k128_b1_b1(%a0 : i32, %a1 : i32,
 288                                %b0 : i32,
 289                                %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32) {
 290   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {b1Op = #nvvm.mma_b1op<xor_popc>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>, shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>
 291   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
 292     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 293      multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>,
 294      b1Op = #nvvm.mma_b1op<xor_popc>,
 295      shape = #nvvm.shape<m = 16, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32,i32,i32,i32)>
 296   llvm.return %0 : !llvm.struct<(i32,i32,i32,i32)>
 297 }
 298
 299 // CHECK-LABEL: @nvvm_mma_m8n8k128_b1_b1
 300 func.func @nvvm_mma_m8n8k128_b1_b1(%a0 : i32,
 301                               %b0 : i32,
 302                               %c0 : i32, %c1 : i32) {
 303   // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}] {b1Op = #nvvm.mma_b1op<xor_popc>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>, shape = #nvvm.shape<m = 8, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32, i32)>
 304   %0 = nvvm.mma.sync A[%a0] B[%b0] C[%c0, %c1]
 305     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 306      multiplicandAPtxType = #nvvm.mma_type<b1>, multiplicandBPtxType = #nvvm.mma_type<b1>,
 307      b1Op = #nvvm.mma_b1op<xor_popc>, shape = #nvvm.shape<m = 8, n = 8, k = 128>} : (i32, i32, i32) -> !llvm.struct<(i32,i32)>
 308   llvm.return %0 : !llvm.struct<(i32,i32)>
 309 }
 310
 311 // CHECK-LABEL: @nvvm_mma_m16n8k32_s4_s4
 312 func.func @nvvm_mma_m16n8k32_s4_s4(%a0 : i32, %a1 : i32,
 313                                %b0 : i32,
 314                                %c0 : i32, %c1 : i32, %c2 : i32, %c3 : i32) {
 315   // CHECK: nvvm.mma.sync A[{{.*}}, {{.*}}] B[{{.*}}] C[{{.*}}, {{.*}}, {{.*}}, {{.*}}] {intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<s4>, multiplicandBPtxType = #nvvm.mma_type<s4>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32, i32, i32, i32)>
 316   %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
 317     {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
 318      multiplicandAPtxType = #nvvm.mma_type<s4>, multiplicandBPtxType = #nvvm.mma_type<s4>,
 319      intOverflowBehavior=#nvvm.mma_int_overflow<wrapped>,
 320      shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32) -> !llvm.struct<(i32,i32,i32,i32)>
 321   llvm.return %0 : !llvm.struct<(i32,i32,i32,i32)>
 322 }
 323
 324 // CHECK-LABEL: @nvvm_wmma_load_tf32
 325 func.func @nvvm_wmma_load_tf32(%arg0: !llvm.ptr, %arg1 : i32) -> !llvm.struct<(i32, i32, i32, i32)> {
 326   // CHECK: nvvm.wmma.load {{.*}} {eltype = #nvvm.mma_type<tf32>, frag = #nvvm.mma_frag<a>, k = 8 : i32, layout = #nvvm.mma_layout<row>, m = 16 : i32, n = 16 : i32}
 327   %0 = nvvm.wmma.load %arg0, %arg1
 328     {eltype = #nvvm.mma_type<tf32>, frag = #nvvm.mma_frag<a>, k = 8 : i32, layout = #nvvm.mma_layout<row>, m = 16 : i32, n = 16 : i32}
 329     : (!llvm.ptr) -> !llvm.struct<(i32, i32, i32, i32)>
 330   llvm.return %0 : !llvm.struct<(i32, i32, i32, i32)>
 331 }
 332
 333 // CHECK-LABEL: @nvvm_wmma_mma
 334 func.func @nvvm_wmma_mma(%0 : i32, %1 : i32, %2 : i32, %3 : i32, %4 : i32, %5 : i32,
 335                     %6 : i32, %7 : i32, %8 : f32, %9 : f32, %10 : f32,
 336                     %11 : f32, %12 : f32, %13 : f32, %14 : f32, %15 : f32)
 337                    -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)> {
 338   // CHECK: nvvm.wmma.mma {{.*}} {eltypeA = #nvvm.mma_type<tf32>, eltypeB = #nvvm.mma_type<f32>, k = 8 : i32, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<row>, m = 16 : i32, n = 16 : i32}
 339   %r = nvvm.wmma.mma %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15
 340     {eltypeA = #nvvm.mma_type<tf32>, eltypeB = #nvvm.mma_type<f32>, k = 8 : i32, layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<row>, m = 16 : i32, n = 16 : i32}
 341     : (i32, i32, i32, i32, i32, i32, i32, i32, f32, f32, f32, f32, f32, f32, f32, f32)
 342     -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
 343   llvm.return %r : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
 344 }
 345
 346 // CHECK-LABEL: @cp_async
 347 llvm.func @cp_async(%arg0: !llvm.ptr<3>, %arg1: !llvm.ptr<1>) {
 348 // CHECK:  nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = ca
 349   nvvm.cp.async.shared.global %arg0, %arg1, 16, cache = ca : !llvm.ptr<3>, !llvm.ptr<1>
 350 // CHECK:  nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache = cg
 351   nvvm.cp.async.shared.global %arg0, %arg1, 16, cache = cg : !llvm.ptr<3>, !llvm.ptr<1>
 352 // CHECK: nvvm.cp.async.commit.group
 353   nvvm.cp.async.commit.group
 354 // CHECK: nvvm.cp.async.wait.group 0
 355   nvvm.cp.async.wait.group 0
 356   llvm.return
 357 }
 358
 359 // CHECK-LABEL: llvm.func @ld_matrix
 360 llvm.func @ld_matrix(%arg0: !llvm.ptr<3>) {
 361   // CHECK: nvvm.ldmatrix %{{.*}} {layout = #nvvm.mma_layout<row>, num = 1 : i32} : (!llvm.ptr<3>) -> i32
 362   %l1 = nvvm.ldmatrix %arg0 {num = 1 : i32, layout = #nvvm.mma_layout<row>} : (!llvm.ptr<3>) -> i32
 363   // CHECK: nvvm.ldmatrix %{{.*}} {layout = #nvvm.mma_layout<row>, num = 2 : i32} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32)>
 364   %l2 = nvvm.ldmatrix %arg0 {num = 2 : i32, layout = #nvvm.mma_layout<row>} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32)>
 365   // CHECK: nvvm.ldmatrix %{{.*}} {layout = #nvvm.mma_layout<row>, num = 4 : i32} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
 366   %l4 = nvvm.ldmatrix %arg0 {num = 4 : i32, layout = #nvvm.mma_layout<row>} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
 367   llvm.return
 368 }
 369
 370 // CHECK-LABEL: llvm.func @redux_sync
 371 llvm.func @redux_sync(%value : i32, %offset : i32) -> i32 {
 372   // CHECK: nvvm.redux.sync  add %{{.*}}
 373   %r1 = nvvm.redux.sync add %value, %offset : i32 -> i32
 374   // CHECK: nvvm.redux.sync  max %{{.*}}
 375   %r2 = nvvm.redux.sync max %value, %offset : i32 -> i32
 376   // CHECK: nvvm.redux.sync  min %{{.*}}
 377   %r3 = nvvm.redux.sync min %value, %offset : i32 -> i32
 378   // CHECK: nvvm.redux.sync  umax %{{.*}}
 379   %r5 = nvvm.redux.sync umax %value, %offset : i32 -> i32
 380   // CHECK: nvvm.redux.sync  umin %{{.*}}
 381   %r6 = nvvm.redux.sync umin %value, %offset : i32 -> i32
 382   // CHECK: nvvm.redux.sync  and %{{.*}}
 383   %r7 = nvvm.redux.sync and %value, %offset : i32 -> i32
 384   // CHECK: nvvm.redux.sync  or %{{.*}}
 385   %r8 = nvvm.redux.sync or %value, %offset : i32 -> i32
 386   // CHECK: nvvm.redux.sync  xor %{{.*}}
 387   %r9 = nvvm.redux.sync xor %value, %offset : i32 -> i32
 388   llvm.return %r1 : i32
 389 }
 390
 391
 392 // -----
 393
 394 // expected-error@below {{attribute attached to unexpected op}}
 395 func.func private @expected_llvm_func() attributes { nvvm.kernel }
 396
 397 // -----
 398
 399 llvm.func private @mbarrier_init_generic(%barrier: !llvm.ptr) {
 400   %count = nvvm.read.ptx.sreg.ntid.x : i32
 401   // CHECK:   nvvm.mbarrier.init %{{.*}}, %{{.*}} : !llvm.ptr, i32
 402   nvvm.mbarrier.init %barrier, %count : !llvm.ptr, i32
 403   llvm.return
 404 }
 405
 406
 407 llvm.func private @mbarrier_init_shared(%barrier: !llvm.ptr<3>) {
 408   %count = nvvm.read.ptx.sreg.ntid.x : i32
 409   // CHECK:   nvvm.mbarrier.init.shared %{{.*}}, %{{.*}} : !llvm.ptr<3>, i32
 410   nvvm.mbarrier.init.shared %barrier, %count : !llvm.ptr<3>, i32
 411   llvm.return
 412 }
 413
 414
 415 llvm.func private @mbarrier_inval_generic(%barrier: !llvm.ptr) {
 416   // CHECK:   nvvm.mbarrier.inval %{{.*}} : !llvm.ptr
 417   nvvm.mbarrier.inval %barrier : !llvm.ptr
 418   llvm.return
 419 }
 420
 421
 422 llvm.func private @mbarrier_inval_shared(%barrier: !llvm.ptr<3>) {
 423   // CHECK:   nvvm.mbarrier.inval.shared %{{.*}} : !llvm.ptr<3>
 424   nvvm.mbarrier.inval.shared %barrier : !llvm.ptr<3>
 425   llvm.return
 426 }
 427
 428 llvm.func private @mbarrier_arrive(%barrier: !llvm.ptr) {
 429   // CHECK:   nvvm.mbarrier.arrive %{{.*}} : !llvm.ptr
 430   %0 = nvvm.mbarrier.arrive %barrier : !llvm.ptr  -> i64
 431   llvm.return
 432 }
 433
 434 llvm.func private @mbarrier_arrive_shared(%barrier: !llvm.ptr<3>) {
 435   // CHECK:   nvvm.mbarrier.arrive.shared %{{.*}} : !llvm.ptr<3>
 436   %0 = nvvm.mbarrier.arrive.shared %barrier : !llvm.ptr<3> -> i64
 437   llvm.return
 438 }
 439
 440 llvm.func private @mbarrier_arrive_nocomplete(%barrier: !llvm.ptr) {
 441   %count = nvvm.read.ptx.sreg.ntid.x : i32
 442   // CHECK:   nvvm.mbarrier.arrive.nocomplete %{{.*}} : !llvm.ptr
 443   %0 = nvvm.mbarrier.arrive.nocomplete %barrier, %count : !llvm.ptr, i32 -> i64
 444   llvm.return
 445 }
 446
 447 llvm.func private @mbarrier_arrive_nocomplete_shared(%barrier: !llvm.ptr<3>) {
 448   %count = nvvm.read.ptx.sreg.ntid.x : i32
 449   // CHECK:   nvvm.mbarrier.arrive.nocomplete.shared %{{.*}} : !llvm.ptr<3>
 450   %0 = nvvm.mbarrier.arrive.nocomplete.shared %barrier, %count : !llvm.ptr<3>, i32  -> i64
 451   llvm.return
 452 }
 453
 454 llvm.func private @mbarrier_test_wait(%barrier: !llvm.ptr, %token : i64) -> i1 {
 455   // CHECK:   nvvm.mbarrier.test.wait %{{.*}}
 456   %isComplete = nvvm.mbarrier.test.wait %barrier, %token : !llvm.ptr, i64 -> i1
 457   llvm.return %isComplete : i1
 458 }
 459
 460 llvm.func private @mbarrier_test_wait_shared(%barrier: !llvm.ptr<3>, %token : i64) {
 461   %count = nvvm.read.ptx.sreg.ntid.x : i32
 462   // CHECK:   nvvm.mbarrier.test.wait.shared %{{.*}}
 463   %isComplete = nvvm.mbarrier.test.wait.shared %barrier, %token : !llvm.ptr<3>, i64 -> i1
 464   llvm.return
 465 }
 466
 467 // CHECK-LABEL: @wgmma_fence_aligned
 468 func.func @wgmma_fence_aligned() {
 469   // CHECK: nvvm.wgmma.fence.aligned
 470   nvvm.wgmma.fence.aligned
 471   return
 472 }
 473
 474 // CHECK-LABEL: @wgmma_commit_group_sync_aligned
 475 func.func @wgmma_commit_group_sync_aligned() {
 476   // CHECK: nvvm.wgmma.commit.group.sync.aligned
 477   nvvm.wgmma.commit.group.sync.aligned
 478   return
 479 }
 480
 481
 482 // CHECK-LABEL: @wgmma_wait_group_sync_aligned
 483 func.func @wgmma_wait_group_sync_aligned() {
 484   // CHECK: nvvm.wgmma.wait.group.sync.aligned
 485   nvvm.wgmma.wait.group.sync.aligned 0
 486   return
 487 }
 488
 489 // -----
 490
 491 // Just check these don't emit errors.
 492 gpu.module @module_1 [#nvvm.target<chip = "sm_90", features = "+ptx70", link = ["my_device_lib.bc"], flags = {fast, ftz}>] {
 493 }
 494
 495 gpu.module @module_2 [#nvvm.target<chip = "sm_90">, #nvvm.target<chip = "sm_80">, #nvvm.target<chip = "sm_70">] {
 496 }
 497
 498 // CHECK-LABEL: nvvm.grid_constant
 499 llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant}) attributes {nvvm.kernel} {
 500   llvm.return
 501 }
 502
 503 // -----
 504
 505 // expected-error @below {{'"nvvm.grid_constant"' attribute must be present only on kernel arguments}}
 506 llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant}) {
 507   llvm.return
 508 }
 509
 510 // -----
 511
 512 // expected-error @below {{'"nvvm.grid_constant"' attribute requires the argument to also have attribute 'llvm.byval'}}
 513 llvm.func @kernel_func(%arg0: !llvm.ptr {nvvm.grid_constant}) attributes {nvvm.kernel} {
 514   llvm.return
 515 }
 516
 517 // -----
 518
 519 // expected-error @below {{'"nvvm.grid_constant"' must be a unit attribute}}
 520 llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant = true}) attributes {nvvm.kernel} {
 521   llvm.return
 522 }