[Mlir-commits] [mlir] [MLIR][NVVM] Add `nvvm.tanh` OP (PR #199525)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Varad Rahul Kamthe (varadk27)

<details>
<summary>Changes</summary>

Implement `nvvm.tanh` lowering to` llvm.tanh` with `afn `fast-math flag. NVPTX backend pattern-matches this into `tanh.approx.f32`. PTX does not expose an `ftz` modifier for `tanh.approx`, so the op has no `ftz` attribute.

---
Full diff: https://github.com/llvm/llvm-project/pull/199525.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+24) 
- (modified) mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir (+7) 
- (modified) mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir (+7) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 030c33526b16a..5cc90e585e542 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -501,6 +501,30 @@ def NVVM_Ex2Op : NVVM_F32UnaryApproxOp<"ex2"> {
   }];
 }
 
+def NVVM_TanhOp : NVVM_Op<"tanh", [Pure, SameOperandsAndResultType]> {
+  let summary = "Hyperbolic tangent (fast approximation)";
+  let description = [{
+    Computes a fast approximation of the hyperbolic tangent of the input
+    value. Lowers to PTX `tanh.approx.f32` (sm_75+, PTX 7.0+). PTX does not
+    expose an `ftz` modifier for `tanh.approx`, so this op has no `ftz`
+    attribute.
+
+    For more information, see PTX ISA:
+    [tanh](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-tanh)
+  }];
+  let arguments = (ins F32:$src);
+  let results = (outs F32:$res);
+  let assemblyFormat = "$src attr-dict `:` type($src)";
+  string llvmBuilder = [{
+    llvm::CallInst *call = createIntrinsicCall(
+        builder, llvm::Intrinsic::tanh, {$src}, {$src->getType()});
+    llvm::FastMathFlags fmf;
+    fmf.setApproxFunc();
+    call->setFastMathFlags(fmf);
+    $res = call;
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM rsqrt op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
index eaaedd335d3d2..3ad03cebbae1c 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
@@ -55,3 +55,10 @@ func.func @nvvm_ex2_ftz_f32(%arg0: f32) -> f32 {
   %0 = nvvm.ex2 %arg0 {ftz = true} : f32
   return %0 : f32
 }
+
+// CHECK-LABEL: @nvvm_tanh_f32
+func.func @nvvm_tanh_f32(%arg0: f32) -> f32 {
+  // CHECK: nvvm.tanh {{.*}} : f32
+  %0 = nvvm.tanh %arg0 : f32
+  return %0 : f32
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 5b7ecc9e1c87c..319db3bdd778e 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -55,3 +55,10 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> f32 {
   %0 = nvvm.ex2 %arg0 {ftz = true} : f32
   llvm.return %0 : f32
 }
+
+// CHECK-LABEL: @nvvm_tanh
+llvm.func @nvvm_tanh(%arg0: f32) -> f32 {
+  // CHECK: call afn float @llvm.tanh.f32(float %{{.*}})
+  %0 = nvvm.tanh %arg0 : f32
+  llvm.return %0 : f32
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/199525