[llvm-branch-commits] [mlir] bb7d43e - Add rsqrt lowering from standard to NVVM

[Frederik Gossen via llvm-branch-commits]

Author: Frederik Gossen
Date: 2020-12-08T14:33:58+01:00
New Revision: bb7d43e7d5f683fc3e7109072610dc8d176a3bf8

URL: https://github.com/llvm/llvm-project/commit/bb7d43e7d5f683fc3e7109072610dc8d176a3bf8
DIFF: https://github.com/llvm/llvm-project/commit/bb7d43e7d5f683fc3e7109072610dc8d176a3bf8.diff

LOG: Add rsqrt lowering from standard to NVVM

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

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
    mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 61f97d44e125..4b7bb6e193a4 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -183,6 +183,8 @@ void mlir::populateGpuToNVVMConversionPatterns(
                                                  "__nv_log10");
   patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__nv_log2f",
                                                 "__nv_log2");
+  patterns.insert<OpToFuncCallLowering<RsqrtOp>>(converter, "__nv_rsqrtf",
+                                                 "__nv_rsqrt");
   patterns.insert<OpToFuncCallLowering<SinOp>>(converter, "__nv_sinf",
                                                "__nv_sin");
   patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__nv_tanhf",

diff  --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index 8e4e4515c638..347cc48daa20 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -295,6 +295,26 @@ gpu.module @test_module {
 
 // -----
 
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_rsqrtf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_rsqrt(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_rsqrt
+  func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64)
+      -> (f16, f32, f64) {
+    %result16 = std.rsqrt %arg_f16 : f16
+    // CHECK: llvm.fpext %{{.*}} : !llvm.half to !llvm.float
+    // CHECK-NEXT: llvm.call @__nv_rsqrtf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    // CHECK-NEXT: llvm.fptrunc %{{.*}} : !llvm.float to !llvm.half
+    %result32 = std.rsqrt %arg_f32 : f32
+    // CHECK: llvm.call @__nv_rsqrtf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.rsqrt %arg_f64 : f64
+    // CHECK: llvm.call @__nv_rsqrt(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return %result16, %result32, %result64 : f16, f32, f64
+  }
+}
+
+// -----
+
 // Test that we handled properly operation with SymbolTable other than module op
 gpu.module @test_module {
   "test.symbol_scope"() ({