1// RUN: mlir-opt %s -convert-gpu-to-rocdl -split-input-file | FileCheck %s
2// RUN: mlir-opt %s -convert-gpu-to-rocdl='index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s
3
4gpu.module @test_module {
5  // CHECK-LABEL: func @gpu_index_ops()
6  // CHECK32-LABEL: func @gpu_index_ops()
7  func @gpu_index_ops()
8      -> (index, index, index, index, index, index,
9          index, index, index, index, index, index) {
10    // CHECK32-NOT: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
11
12    // CHECK: rocdl.workitem.id.x : !llvm.i32
13    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
14    %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
15    // CHECK: rocdl.workitem.id.y : !llvm.i32
16    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
17    %tIdY = "gpu.thread_id"() {dimension = "y"} : () -> (index)
18    // CHECK: rocdl.workitem.id.z : !llvm.i32
19    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
20    %tIdZ = "gpu.thread_id"() {dimension = "z"} : () -> (index)
21
22    // CHECK: rocdl.workgroup.dim.x : !llvm.i32
23    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
24    %bDimX = "gpu.block_dim"() {dimension = "x"} : () -> (index)
25    // CHECK: rocdl.workgroup.dim.y : !llvm.i32
26    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
27    %bDimY = "gpu.block_dim"() {dimension = "y"} : () -> (index)
28    // CHECK: rocdl.workgroup.dim.z : !llvm.i32
29    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
30    %bDimZ = "gpu.block_dim"() {dimension = "z"} : () -> (index)
31
32    // CHECK: rocdl.workgroup.id.x : !llvm.i32
33    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
34    %bIdX = "gpu.block_id"() {dimension = "x"} : () -> (index)
35    // CHECK: rocdl.workgroup.id.y : !llvm.i32
36    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
37    %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index)
38    // CHECK: rocdl.workgroup.id.z : !llvm.i32
39    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
40    %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index)
41
42    // CHECK: rocdl.grid.dim.x : !llvm.i32
43    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
44    %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index)
45    // CHECK: rocdl.grid.dim.y : !llvm.i32
46    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
47    %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index)
48    // CHECK: rocdl.grid.dim.z : !llvm.i32
49    // CHECK: = llvm.sext %{{.*}} : !llvm.i32 to !llvm.i64
50    %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
51
52    std.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
53               %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ
54        : index, index, index, index, index, index,
55          index, index, index, index, index, index
56  }
57}
58
59// -----
60
61gpu.module @test_module {
62  // CHECK-LABEL: func @gpu_index_comp
63  // CHECK32-LABEL: func @gpu_index_comp
64  func @gpu_index_comp(%idx : index) -> index {
65    // CHECK: = llvm.add %{{.*}}, %{{.*}} : !llvm.i64
66    // CHECK32: = llvm.add %{{.*}}, %{{.*}} : !llvm.i32
67    %0 = addi %idx, %idx : index
68    // CHECK: llvm.return %{{.*}} : !llvm.i64
69    // CHECK32: llvm.return %{{.*}} : !llvm.i32
70    std.return %0 : index
71  }
72}
73
74// -----
75
76gpu.module @test_module {
77  // CHECK-LABEL: func @gpu_sync()
78  func @gpu_sync() {
79    // CHECK: rocdl.barrier
80    gpu.barrier
81    std.return
82  }
83}
84
85// -----
86
87gpu.module @test_module {
88  // CHECK: llvm.func @__ocml_fabs_f32(!llvm.float) -> !llvm.float
89  // CHECK: llvm.func @__ocml_fabs_f64(!llvm.double) -> !llvm.double
90  // CHECK-LABEL: func @gpu_fabs
91  func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
92    %result32 = std.absf %arg_f32 : f32
93    // CHECK: llvm.call @__ocml_fabs_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
94    %result64 = std.absf %arg_f64 : f64
95    // CHECK: llvm.call @__ocml_fabs_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
96    std.return %result32, %result64 : f32, f64
97  }
98}
99
100// -----
101
102gpu.module @test_module {
103  // CHECK: llvm.func @__ocml_ceil_f32(!llvm.float) -> !llvm.float
104  // CHECK: llvm.func @__ocml_ceil_f64(!llvm.double) -> !llvm.double
105  // CHECK-LABEL: func @gpu_ceil
106  func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
107    %result32 = std.ceilf %arg_f32 : f32
108    // CHECK: llvm.call @__ocml_ceil_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
109    %result64 = std.ceilf %arg_f64 : f64
110    // CHECK: llvm.call @__ocml_ceil_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
111    std.return %result32, %result64 : f32, f64
112  }
113}
114
115// -----
116
117gpu.module @test_module {
118  // CHECK: llvm.func @__ocml_cos_f32(!llvm.float) -> !llvm.float
119  // CHECK: llvm.func @__ocml_cos_f64(!llvm.double) -> !llvm.double
120  // CHECK-LABEL: func @gpu_cos
121  func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
122    %result32 = std.cos %arg_f32 : f32
123    // CHECK: llvm.call @__ocml_cos_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
124    %result64 = std.cos %arg_f64 : f64
125    // CHECK: llvm.call @__ocml_cos_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
126    std.return %result32, %result64 : f32, f64
127  }
128}
129
130// -----
131gpu.module @test_module {
132  // CHECK: llvm.func @__ocml_exp_f32(!llvm.float) -> !llvm.float
133  // CHECK: llvm.func @__ocml_exp_f64(!llvm.double) -> !llvm.double
134  // CHECK-LABEL: func @gpu_exp
135  func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
136    %exp_f32 = std.exp %arg_f32 : f32
137    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
138    %result32 = std.exp %exp_f32 : f32
139    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
140    %result64 = std.exp %arg_f64 : f64
141    // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
142    std.return %result32, %result64 : f32, f64
143  }
144}
145
146
147// -----
148
149// Test that we handled properly operation with SymbolTable other than module op
150gpu.module @test_module {
151  "test.symbol_scope"() ({
152    // CHECK: test.symbol_scope
153    // CHECK: llvm.func @__ocml_exp_f32(!llvm.float) -> !llvm.float
154    // CHECK: llvm.func @__ocml_exp_f64(!llvm.double) -> !llvm.double
155    // CHECK-LABEL: func @gpu_exp
156    func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
157      %exp_f32 = std.exp %arg_f32 : f32
158      // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
159      %result32 = std.exp %exp_f32 : f32
160      // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
161      %result64 = std.exp %arg_f64 : f64
162      // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
163      std.return %result32, %result64 : f32, f64
164    }
165    "test.finish" () : () -> ()
166  }) : () -> ()
167}
168
169// -----
170
171gpu.module @test_module {
172  // CHECK: llvm.func @__ocml_log_f32(!llvm.float) -> !llvm.float
173  // CHECK: llvm.func @__ocml_log_f64(!llvm.double) -> !llvm.double
174  // CHECK-LABEL: func @gpu_log
175  func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
176    %result32 = std.log %arg_f32 : f32
177    // CHECK: llvm.call @__ocml_log_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
178    %result64 = std.log %arg_f64 : f64
179    // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
180    std.return %result32, %result64 : f32, f64
181  }
182}
183
184// -----
185
186gpu.module @test_module {
187  // CHECK: llvm.func @__ocml_log10_f32(!llvm.float) -> !llvm.float
188  // CHECK: llvm.func @__ocml_log10_f64(!llvm.double) -> !llvm.double
189  // CHECK-LABEL: func @gpu_log10
190  func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
191    %result32 = std.log10 %arg_f32 : f32
192    // CHECK: llvm.call @__ocml_log10_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
193    %result64 = std.log10 %arg_f64 : f64
194    // CHECK: llvm.call @__ocml_log10_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
195    std.return %result32, %result64 : f32, f64
196  }
197}
198
199// -----
200
201gpu.module @test_module {
202  // CHECK: llvm.func @__ocml_log2_f32(!llvm.float) -> !llvm.float
203  // CHECK: llvm.func @__ocml_log2_f64(!llvm.double) -> !llvm.double
204  // CHECK-LABEL: func @gpu_log2
205  func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
206    %result32 = std.log2 %arg_f32 : f32
207    // CHECK: llvm.call @__ocml_log2_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
208    %result64 = std.log2 %arg_f64 : f64
209    // CHECK: llvm.call @__ocml_log2_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
210    std.return %result32, %result64 : f32, f64
211  }
212}
213
214// -----
215
216gpu.module @test_module {
217  // CHECK: llvm.func @__ocml_tanh_f32(!llvm.float) -> !llvm.float
218  // CHECK: llvm.func @__ocml_tanh_f64(!llvm.double) -> !llvm.double
219  // CHECK-LABEL: func @gpu_tanh
220  func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
221    %result32 = std.tanh %arg_f32 : f32
222    // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
223    %result64 = std.tanh %arg_f64 : f64
224    // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
225    std.return %result32, %result64 : f32, f64
226  }
227}
228