[Mlir-commits] [mlir] [MLIR] Expose approximation patterns for tanh/erf. (PR #82750)

[Johannes Reifferscheid]

https://github.com/jreiffers created https://github.com/llvm/llvm-project/pull/82750

These patterns can already be used via
populateMathPolynomialApproximationPatterns, but that includes a number of other patterns that may not be needed.

There are already similar functions for expansion.

For now only adding tanh and erf since I have a concrete use case for these two.

>From bf857c7a63dcd1f1eb6755e2f02ca6c7b3a8531b Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers at google.com>
Date: Fri, 23 Feb 2024 11:49:10 +0100
Subject: [PATCH] [MLIR] Expose approximation patterns for tanh/erf.

These patterns can already be used via
populateMathPolynomialApproximationPatterns, but that includes a number
of other patterns that may not be needed.

There are already similar functions for expansion.

For now only adding tanh and erf since I have a concrete use case for
these two.
---
 mlir/include/mlir/Dialect/Math/Transforms/Passes.h     |  4 ++++
 .../Math/Transforms/PolynomialApproximation.cpp        | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
index 010dde5ea73847..c0b5c0080bdcef 100644
--- a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
@@ -45,6 +45,10 @@ struct MathPolynomialApproximationOptions {
   bool enableAvx2 = false;
 };
 
+
+void populatePolynomialApproximateTanhPattern(RewritePatternSet &patterns);
+void populatePolynomialApproximateErfPattern(RewritePatternSet &patterns);
+
 void populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns,
     const MathPolynomialApproximationOptions &options = {});
diff --git a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
index 71e4e13103f514..962cb28b7c2ab9 100644
--- a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
@@ -1471,6 +1471,16 @@ RsqrtApproximation::matchAndRewrite(math::RsqrtOp op,
 
 //----------------------------------------------------------------------------//
 
+void mlir::populatePolynomialApproximateTanhPattern(
+    RewritePatternSet &patterns) {
+  patterns.add<TanhApproximation>(patterns.getContext());
+}
+
+void mlir::populatePolynomialApproximateErfPattern(
+    RewritePatternSet &patterns) {
+  patterns.add<ErfPolynomialApproximation>(patterns.getContext());
+}
+
 void mlir::populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns,
     const MathPolynomialApproximationOptions &options) {