[mlir] Added missing GPU lowering ops.

Summary:
This diff adds missing GPU lowering ops to MLIR.

Reviewers: herhut, pifon2a, ftynse

Tags: #pre-merge_beta_testing, #llvm

Differential Revision: https://reviews.llvm.org/D72439

2 files changed, 90 insertions, 0 deletions

diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index b1820cb778f..24bf56557c3 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -112,6 +112,51 @@ module attributes {gpu.kernel_module} {
 // -----
 
 module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @__nv_fabsf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_fabs(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_fabs
+  func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.absf %arg_f32 : f32
+    // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.absf %arg_f64 : f64
+    // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @__nv_ceilf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_ceil(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_ceil
+  func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.ceilf %arg_f32 : f32
+    // CHECK: llvm.call @__nv_ceilf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.ceilf %arg_f64 : f64
+    // CHECK: llvm.call @__nv_ceil(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @__nv_cosf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_cos(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_cos
+  func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.cos %arg_f32 : f32
+    // CHECK: llvm.call @__nv_cosf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.cos %arg_f64 : f64
+    // CHECK: llvm.call @__nv_cos(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
   // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index f803266ad28..372fe2784c3 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -39,6 +39,51 @@ module attributes {gpu.kernel_module} {
 // -----
 
 module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @_ocml_fabs_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @_ocml_fabs_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_fabs
+  func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.absf %arg_f32 : f32
+    // CHECK: llvm.call @_ocml_fabs_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.absf %arg_f64 : f64
+    // CHECK: llvm.call @_ocml_fabs_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @_ocml_ceil_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @_ocml_ceil_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_ceil
+  func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.ceilf %arg_f32 : f32
+    // CHECK: llvm.call @_ocml_ceil_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.ceilf %arg_f64 : f64
+    // CHECK: llvm.call @_ocml_ceil_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
+  // CHECK: llvm.func @_ocml_cos_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @_ocml_cos_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_cos
+  func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.cos %arg_f32 : f32
+    // CHECK: llvm.call @_ocml_cos_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.cos %arg_f64 : f64
+    // CHECK: llvm.call @_ocml_cos_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+module attributes {gpu.kernel_module} {
   // CHECK: llvm.func @_ocml_exp_f32(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @_ocml_exp_f64(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp