• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// RUN: mlir-translate -mlir-to-nvvmir %s | FileCheck %s
2
3llvm.func @nvvm_special_regs() -> !llvm.i32 {
4  // CHECK: %1 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
5  %1 = nvvm.read.ptx.sreg.tid.x : !llvm.i32
6  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.tid.y()
7  %2 = nvvm.read.ptx.sreg.tid.y : !llvm.i32
8  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.tid.z()
9  %3 = nvvm.read.ptx.sreg.tid.z : !llvm.i32
10  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
11  %4 = nvvm.read.ptx.sreg.ntid.x : !llvm.i32
12  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ntid.y()
13  %5 = nvvm.read.ptx.sreg.ntid.y : !llvm.i32
14  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ntid.z()
15  %6 = nvvm.read.ptx.sreg.ntid.z : !llvm.i32
16  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
17  %7 = nvvm.read.ptx.sreg.ctaid.x : !llvm.i32
18  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ctaid.y()
19  %8 = nvvm.read.ptx.sreg.ctaid.y : !llvm.i32
20  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.ctaid.z()
21  %9 = nvvm.read.ptx.sreg.ctaid.z : !llvm.i32
22  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nctaid.x()
23  %10 = nvvm.read.ptx.sreg.nctaid.x : !llvm.i32
24  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nctaid.y()
25  %11 = nvvm.read.ptx.sreg.nctaid.y : !llvm.i32
26  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nctaid.z()
27  %12 = nvvm.read.ptx.sreg.nctaid.z : !llvm.i32
28  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.warpsize()
29  %13 = nvvm.read.ptx.sreg.warpsize : !llvm.i32
30  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.laneid()
31  %14 = nvvm.read.ptx.sreg.laneid : !llvm.i32
32  llvm.return %1 : !llvm.i32
33}
34
35llvm.func @llvm.nvvm.barrier0() {
36  // CHECK: call void @llvm.nvvm.barrier0()
37  nvvm.barrier0
38  llvm.return
39}
40
41llvm.func @nvvm_shfl(
42    %0 : !llvm.i32, %1 : !llvm.i32, %2 : !llvm.i32,
43    %3 : !llvm.i32, %4 : !llvm.float) -> !llvm.i32 {
44  // CHECK: call i32 @llvm.nvvm.shfl.sync.bfly.i32(i32 %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i32 %{{.*}})
45  %6 = nvvm.shfl.sync.bfly %0, %3, %1, %2 : !llvm.i32
46  // CHECK: call float @llvm.nvvm.shfl.sync.bfly.f32(i32 %{{.*}}, float %{{.*}}, i32 %{{.*}}, i32 %{{.*}})
47  %7 = nvvm.shfl.sync.bfly %0, %4, %1, %2 : !llvm.float
48  llvm.return %6 : !llvm.i32
49}
50
51llvm.func @nvvm_shfl_pred(
52    %0 : !llvm.i32, %1 : !llvm.i32, %2 : !llvm.i32,
53    %3 : !llvm.i32, %4 : !llvm.float) -> !llvm.struct<(i32, i1)> {
54  // CHECK: call { i32, i1 } @llvm.nvvm.shfl.sync.bfly.i32p(i32 %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i32 %{{.*}})
55  %6 = nvvm.shfl.sync.bfly %0, %3, %1, %2 {return_value_and_is_valid} : !llvm.struct<(i32, i1)>
56  // CHECK: call { float, i1 } @llvm.nvvm.shfl.sync.bfly.f32p(i32 %{{.*}}, float %{{.*}}, i32 %{{.*}}, i32 %{{.*}})
57  %7 = nvvm.shfl.sync.bfly %0, %4, %1, %2 {return_value_and_is_valid} : !llvm.struct<(float, i1)>
58  llvm.return %6 : !llvm.struct<(i32, i1)>
59}
60
61llvm.func @nvvm_vote(%0 : !llvm.i32, %1 : !llvm.i1) -> !llvm.i32 {
62  // CHECK: call i32 @llvm.nvvm.vote.ballot.sync(i32 %{{.*}}, i1 %{{.*}})
63  %3 = nvvm.vote.ballot.sync %0, %1 : !llvm.i32
64  llvm.return %3 : !llvm.i32
65}
66
67llvm.func @nvvm_mma(%a0 : !llvm.vec<2 x half>, %a1 : !llvm.vec<2 x half>,
68                    %b0 : !llvm.vec<2 x half>, %b1 : !llvm.vec<2 x half>,
69                    %c0 : !llvm.float, %c1 : !llvm.float, %c2 : !llvm.float, %c3 : !llvm.float,
70                    %c4 : !llvm.float, %c5 : !llvm.float, %c6 : !llvm.float, %c7 : !llvm.float) {
71  // CHECK: call { float, float, float, float, float, float, float, float } @llvm.nvvm.mma.m8n8k4.row.col.f32.f32
72  %0 = nvvm.mma.sync %a0, %a1, %b0, %b1, %c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7 {alayout="row", blayout="col"} : (!llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float) -> !llvm.struct<(float, float, float, float, float, float, float, float)>
73  llvm.return %0 : !llvm.struct<(float, float, float, float, float, float, float, float)>
74}
75
76// This function has the "kernel" attribute attached and should appear in the
77// NVVM annotations after conversion.
78llvm.func @kernel_func() attributes {gpu.kernel} {
79  llvm.return
80}
81
82// CHECK:     !nvvm.annotations =
83// CHECK-NOT: {i32 ()* @nvvm_special_regs, !"kernel", i32 1}
84// CHECK:     {void ()* @kernel_func, !"kernel", i32 1}
85