Lower math.expm1 to intrinsics in the GPUToNVVM and GPUToROCDL conversions.

[akuegel]

Summary

This adds the lowering for expm1 for GPU backends.

[herhut]

This looks good. Is there a corresponding diff that adds the expm1 operation to math?

[akuegel]

Already landed as 9f581815ae4d5df7453f0b3bea845c0ba508c91e (Mehdi approved the review in https://reviews.llvm.org/D96704)

[herhut]

Thanks. Phabricator keeps hiding these from me (might be my fault) and I have to find them in my email.

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

void mlir::populateGpuToNVVMConversionPatterns(
   patterns.insert<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f",
                                                        "__nv_atan2");
   patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf",
                                                  "__nv_ceil");
   patterns.insert<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf",
                                                      "__nv_cos");
   patterns.insert<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf",
                                                      "__nv_exp");
+  patterns.insert<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f",
+                                                       "__nv_expm1");
   patterns.insert<OpToFuncCallLowering<FloorFOp>>(converter, "__nv_floorf",
                                                   "__nv_floor");
   patterns.insert<OpToFuncCallLowering<math::LogOp>>(converter, "__nv_logf",
                                                      "__nv_log");
   patterns.insert<OpToFuncCallLowering<math::Log1pOp>>(converter, "__nv_log1pf",
                                                        "__nv_log1p");
   patterns.insert<OpToFuncCallLowering<math::Log10Op>>(converter, "__nv_log10f",
                                                        "__nv_log10");

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

void mlir::populateGpuToROCDLConversionPatterns(
   patterns.insert<OpToFuncCallLowering<math::Atan2Op>>(
       converter, "__ocml_atan2_f32", "__ocml_atan2_f64");
   patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32",
                                                  "__ocml_ceil_f64");
   patterns.insert<OpToFuncCallLowering<math::CosOp>>(
       converter, "__ocml_cos_f32", "__ocml_cos_f64");
   patterns.insert<OpToFuncCallLowering<math::ExpOp>>(
       converter, "__ocml_exp_f32", "__ocml_exp_f64");
+  patterns.insert<OpToFuncCallLowering<math::ExpM1Op>>(
+      converter, "__ocml_expm1_f32", "__ocml_expm1_f64");
   patterns.insert<OpToFuncCallLowering<FloorFOp>>(converter, "__ocml_floor_f32",
                                                   "__ocml_floor_f64");
   patterns.insert<OpToFuncCallLowering<math::LogOp>>(
       converter, "__ocml_log_f32", "__ocml_log_f64");
   patterns.insert<OpToFuncCallLowering<math::Log10Op>>(
       converter, "__ocml_log10_f32", "__ocml_log10_f64");
   patterns.insert<OpToFuncCallLowering<math::Log1pOp>>(
       converter, "__ocml_log1p_f32", "__ocml_log1p_f64");

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

// CHECK-LABEL: func @gpu_exp
     }
     "test.finish" () : () -> ()
   }) : () -> ()
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK: llvm.func @__nv_expm1f(f32) -> f32
+  // CHECK: llvm.func @__nv_expm1(f64) -> f64
+  // CHECK-LABEL: func @gpu_expm1
+  func @gpu_expm1(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
+    %result32 = math.expm1 %arg_f32 : f32
+    // CHECK: llvm.call @__nv_expm1f(%{{.*}}) : (f32) -> f32
+    %result64 = math.expm1 %arg_f64 : f64
+    // CHECK: llvm.call @__nv_expm1(%{{.*}}) : (f64) -> f64
+    std.return %result32, %result64 : f32, f64
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_powf(f32, f32) -> f32
   // CHECK: llvm.func @__nv_pow(f64, f64) -> f64
   // CHECK-LABEL: func @gpu_pow
   func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %result32 = math.powf %arg_f32, %arg_f32 : f32
     // CHECK: llvm.call @__nv_powf(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
     %result64 = math.powf %arg_f64, %arg_f64 : f64
     // CHECK: llvm.call @__nv_pow(%{{.*}}, %{{.*}}) : (f64, f64) -> f64

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     // CHECK: llvm.call @__ocml_cos_f32(%{{.*}}) : (f32) -> f32
     %result64 = math.cos %arg_f64 : f64
     // CHECK: llvm.call @__ocml_cos_f64(%{{.*}}) : (f64) -> f64
     std.return %result32, %result64 : f32, f64
   }
 }
 
 // -----
+
 gpu.module @test_module {
   // CHECK: llvm.func @__ocml_exp_f32(f32) -> f32
   // CHECK: llvm.func @__ocml_exp_f64(f64) -> f64
   // CHECK-LABEL: func @gpu_exp
   func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %exp_f32 = math.exp %arg_f32 : f32
     // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
     %result32 = math.exp %exp_f32 : f32
"test.symbol_scope"() ({
     }
     "test.finish" () : () -> ()
   }) : () -> ()
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK: llvm.func @__ocml_expm1_f32(f32) -> f32
+  // CHECK: llvm.func @__ocml_expm1_f64(f64) -> f64
+  // CHECK-LABEL: func @gpu_expm1
+  func @gpu_expm1(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
+    %expm1_f32 = math.expm1 %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32
+    %result32 = math.expm1 %expm1_f32 : f32
+    // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32
+    %result64 = math.expm1 %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_expm1_f64(%{{.*}}) : (f64) -> f64
+    std.return %result32, %result64 : f32, f64
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK: llvm.func @__ocml_log_f32(f32) -> f32
   // CHECK: llvm.func @__ocml_log_f64(f64) -> f64
   // CHECK-LABEL: func @gpu_log
   func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %result32 = math.log %arg_f32 : f32
     // CHECK: llvm.call @__ocml_log_f32(%{{.*}}) : (f32) -> f32
     %result64 = math.log %arg_f64 : f64
     // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64