[mlir] Add math polynomial approximation pass

[ezhulenev]

Summary

This gives ~30x speedup compared to expanding Tanh into exp operations:

name                  old cpu/op  new cpu/op  delta
BM_mlir_Tanh_f32/10    253ns ± 3%    55ns ± 7%  -78.35%  (p=0.000 n=44+41)
BM_mlir_Tanh_f32/100  2.21µs ± 4%  0.14µs ± 8%  -93.85%  (p=0.000 n=48+49)
BM_mlir_Tanh_f32/1k   22.6µs ± 4%   0.7µs ± 5%  -96.68%  (p=0.000 n=32+42)
BM_mlir_Tanh_f32/10k   225µs ± 5%     7µs ± 6%  -96.88%  (p=0.000 n=49+55)

name                  old time/op             new time/op             delta
BM_mlir_Tanh_f32/10    259ns ± 1%               56ns ± 2%  -78.31%        (p=0.000 n=41+39)
BM_mlir_Tanh_f32/100  2.27µs ± 1%             0.14µs ± 5%  -93.89%        (p=0.000 n=46+49)
BM_mlir_Tanh_f32/1k   22.9µs ± 1%              0.8µs ± 4%  -96.67%        (p=0.000 n=30+42)
BM_mlir_Tanh_f32/10k   230µs ± 0%                7µs ± 3%  -96.88%        (p=0.000 n=37+55)

This approximations is based on Eigen::generic_fast_tanh function

[mehdi_amini]

I don't think a pass is the right way to model this, if we want to have support for fast-math we should model it through ops on each attributes, on the model of what LLVM is doing.

[herhut]

In D96739#2564409, @mehdi_amini wrote:

I don't think a pass is the right way to model this, if we want to have support for fast-math we should model it through ops on each attributes, on the model of what LLVM is doing.

I think this is different than usual fast-math flags. This is not about special properties that allow for optimizations to apply that normally would not. These are different approximations for trigonometric functions, so in a way a different set of intrinsics or library implementations for math functions. For those, I think it makes sense to model them as lowering pattern that compilers can mix in if this is the implementation they want (and not a libm based implementation via function calls).

mlir/lib/Dialect/Math/Transforms/FastMathExpansion.cpp

39 ↗

(On Diff #323851)

This adds yet another convenience builder infrastructure. What happened to EDSC, is that still the way to go? @ftynse do you know?

204 ↗

(On Diff #323851)

Can you expose these via a populate method?

mlir/test/mlir-cpu-runner/fast-math.mlir

8 ↗

(On Diff #323851)

Please drop the dump input.

[ezhulenev]

In D96739#2564794, @herhut wrote:

I think this is different than usual fast-math flags.

Fast-math is probably a bad name for this, this is mostly about function approximations using polynomials, and other techniques to avoid using library functions.

[ezhulenev]

Rename fast-math to math-polynomial-approximation pass

[mehdi_amini]

In D96739#2564794, @herhut wrote:

In D96739#2564409, @mehdi_amini wrote:

I don't think a pass is the right way to model this, if we want to have support for fast-math we should model it through ops on each attributes, on the model of what LLVM is doing.

I think this is different than usual fast-math flags. This is not about special properties that allow for optimizations to apply that normally would not. These are different approximations for trigonometric functions, so in a way a different set of intrinsics or library implementations for math functions. For those, I think it makes sense to model them as lowering pattern that compilers can mix in if this is the implementation they want (and not a libm based implementation via function calls).

You're right: the "fast" naming of the original revision threw me off here. Are these expansions equivalent to libm always right now?

Also the pass only makes sense as a "test" for the patterns, so I rather see the pass in the test folder. A "real" compiler should have a concept of target and legality and load the patterns it needs.

[ezhulenev]

Register polynomial expansion as test pass

[ezhulenev]

In D96739#2567381, @mehdi_amini wrote:

Are these expansions equivalent to libm always right now?

No, currently in Eigen they are usually different. And the plan for MLIR is to use https://www.sollya.org to build these math approximations going forward.

Also the pass only makes sense as a "test" for the patterns, so I rather see the pass in the test folder. A "real" compiler should have a concept of target and legality and load the patterns it needs.

Done.

[tpopp]

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

10

Maybe include a link to Sollya if that's the desired way to create new approximations.

148

I'm guessing Sollya generates the code in this way? It would still be nice to follow the LLVM case style.
s/alpha_/alpha/
s/beta_/beta/

[ezhulenev]

Fix xtyle warning

[mravishankar]

mlir/test/Dialect/Math/polynomial-approximation.mlir

7

Confused by the test. What is it supposed to generate instead?

[ezhulenev]

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

10

I've discussed this today with Rasmus, and the plan is to create a tiny-dsl with roughly the same operations as Eigen packet operations (abstract layer that hides SIMD details) and port important approximations from Eigen. Maybe Sollya will be used to build some of them, but that is not clear yet.

148

Fixed all of that. This particular approximation is based on Eigen code (although it is used pretty much everywhere).

mlir/test/Dialect/Math/polynomial-approximation.mlir

7

Approximation built from std dialect operations. Don't want to put a lot of checks in this test, and tested correctness in a separate mlir-cpu-runner test.

mlir/include/mlir/Dialect/Math/Transforms/Passes.h

 
 namespace mlir {
 
 class OwningRewritePatternList;
 
 void populateExpandTanhPattern(OwningRewritePatternList &patterns,
                                MLIRContext *ctx);
 
+void populateMathPolynomialApproximationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

mlir/lib/Dialect/Math/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMathTransforms
   ExpandTanh.cpp
+  PolynomialApproximation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Math/Transforms
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRMath
   MLIRPass
   MLIRStandard
   MLIRTransforms
   )

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

+//===- PolynomialApproximation.cpp - Approximate math operations ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements expansion of math operations to fast approximations
+// that do not rely on any of the library functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+
+static bool isValidFloatType(Type type) {
+  if (auto vectorType = type.dyn_cast<VectorType>())
+    return vectorType.getElementType().isa<FloatType>();
+  return type.isa<FloatType>();
+}
+
+//----------------------------------------------------------------------------//
+// A PatternRewriter wrapper that provides concise API for building expansions
+// for operations on float scalars or vectors.
+//----------------------------------------------------------------------------//
+
+namespace {
+class FloatApproximationBuilder {
+public:
+  FloatApproximationBuilder(Location loc, Type type, PatternRewriter &rewriter);
+
+  Value constant(double value) const;
+
+  Value abs(Value a) const;                    // abs(a)
+  Value min(Value a, Value b) const;           // min(a, b)
+  Value max(Value a, Value b) const;           // max(a, b)
+  Value mul(Value a, Value b) const;           // a * b
+  Value div(Value a, Value b) const;           // a / b
+  Value madd(Value a, Value b, Value c) const; // a * b + c
+
+  // Compares values `a` and `b` with the given `predicate`.
+  Value cmp(CmpFPredicate predicate, Value a, Value b) const;
+
+  // Selects values from `a` or `b` based on the `predicate`.
+  Value select(Value predicate, Value a, Value b) const;
+
+private:
+  Location loc;
+  PatternRewriter &rewriter;
+  VectorType vectorType; // can be null for scalar type
+  FloatType elementType;
+};
+} // namespace
+
+FloatApproximationBuilder::FloatApproximationBuilder(Location loc, Type type,
+                                                     PatternRewriter &rewriter)
+    : loc(loc), rewriter(rewriter) {
+  vectorType = type.dyn_cast<VectorType>();
+
+  if (vectorType)
+    elementType = vectorType.getElementType().cast<FloatType>();
+  else
+    elementType = type.cast<FloatType>();
+}
+
+Value FloatApproximationBuilder::constant(double value) const {
+  auto attr = rewriter.getFloatAttr(elementType, value);
+  Value scalar = rewriter.create<ConstantOp>(loc, attr);
+
+  if (vectorType)
+    return rewriter.create<BroadcastOp>(loc, vectorType, scalar);
+  return scalar;
+}
+
+Value FloatApproximationBuilder::abs(Value a) const {
+  return rewriter.create<AbsFOp>(loc, a);
+}
+
+Value FloatApproximationBuilder::min(Value a, Value b) const {
+  return select(cmp(CmpFPredicate::OLT, a, b), a, b);
+}
+Value FloatApproximationBuilder::max(Value a, Value b) const {
+  return select(cmp(CmpFPredicate::OGT, a, b), a, b);
+}
+Value FloatApproximationBuilder::mul(Value a, Value b) const {
+  return rewriter.create<MulFOp>(loc, a, b);
+}
+
+Value FloatApproximationBuilder::div(Value a, Value b) const {
+  return rewriter.create<DivFOp>(loc, a, b);
+}
+
+Value FloatApproximationBuilder::madd(Value a, Value b, Value c) const {
+  if (vectorType)
+    return rewriter.create<vector::FMAOp>(loc, a, b, c);
+
+  Value mul = rewriter.create<MulFOp>(loc, a, b);
+  return rewriter.create<AddFOp>(loc, mul, c);
+}
+
+Value FloatApproximationBuilder::cmp(CmpFPredicate predicate, Value a,
+                                     Value b) const {
+  return rewriter.create<CmpFOp>(loc, predicate, a, b);
+}
+
+Value FloatApproximationBuilder::select(Value predicate, Value a,
+                                        Value b) const {
+  return rewriter.create<SelectOp>(loc, predicate, a, b);
+}
+
+//----------------------------------------------------------------------------//
+// TanhOp expansion.
+//----------------------------------------------------------------------------//
+
+namespace {
+struct TanhApproximation : public OpRewritePattern<math::TanhOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::TanhOp op,
+                                PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+LogicalResult
+TanhApproximation::matchAndRewrite(math::TanhOp op,
+                                   PatternRewriter &rewriter) const {
+  if (!isValidFloatType(op.operand().getType()))
+    return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+  Value operand = op.operand();
+  FloatApproximationBuilder builder(op->getLoc(), operand.getType(), rewriter);
+
+  // Clamp operand into [plusClamp, minusClamp] range.
+  Value plusClamp = builder.constant(7.90531110763549805);
+  Value minusClamp = builder.constant(-7.9053111076354980);
+  Value x = builder.max(builder.min(operand, plusClamp), minusClamp);
+
+  // The monomial coefficients of the numerator polynomial (odd).
+  Value alpha_1 = builder.constant(4.89352455891786e-03);
+  Value alpha_3 = builder.constant(6.37261928875436e-04);
+  Value alpha_5 = builder.constant(1.48572235717979e-05);
+  Value alpha_7 = builder.constant(5.12229709037114e-08);
+  Value alpha_9 = builder.constant(-8.60467152213735e-11);
+  Value alpha_11 = builder.constant(2.00018790482477e-13);
+  Value alpha_13 = builder.constant(-2.76076847742355e-16);
+
+  // The monomial coefficients of the denominator polynomial (even).
+  Value beta_0 = builder.constant(4.89352518554385e-03);
+  Value beta_2 = builder.constant(2.26843463243900e-03);
+  Value beta_4 = builder.constant(1.18534705686654e-04);
+  Value beta_6 = builder.constant(1.19825839466702e-06);
+
+  // Since the polynomials are odd/even, we need x^2.
+  Value x2 = builder.mul(x, x);
+
+  // Evaluate the numerator polynomial p.
+  Value p = builder.madd(x2, alpha_13, alpha_11);
+  p = builder.madd(x2, p, alpha_9);
+  p = builder.madd(x2, p, alpha_7);
+  p = builder.madd(x2, p, alpha_5);
+  p = builder.madd(x2, p, alpha_3);
+  p = builder.madd(x2, p, alpha_1);
+  p = builder.mul(x, p);
+
+  // Evaluate the denominator polynomial q.
+  Value q = builder.madd(x2, beta_6, beta_4);
+  q = builder.madd(x2, q, beta_2);
+  q = builder.madd(x2, q, beta_0);
+
+  // Divide the numerator by the denominator.
+  Value tiny = builder.constant(0.0004f);
+  Value tinyMask = builder.cmp(CmpFPredicate::OLT, builder.abs(operand), tiny);
+  Value res = builder.select(tinyMask, x, builder.div(p, q));
+
+  rewriter.replaceOp(op, res);
+
+  return success();
+}
+
+//----------------------------------------------------------------------------//
+
+void mlir::populateMathPolynomialApproximationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<TanhApproximation>(ctx);
+}

mlir/test/Dialect/Math/polynomial-approximation.mlir

+// RUN: mlir-opt %s -test-math-polynomial-approximation | FileCheck %s
+
+// CHECK-LABEL: @tanh_scalar
+func @tanh_scalar(%arg0: f32) -> f32 {
+  // CHECK-NOT: tanh
+  %0 = math.tanh %arg0 : f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @tanh_vector
+func @tanh_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
+  // CHECK-NOT: tanh
+  %0 = math.tanh %arg0 : vector<8xf32>
+  return %0 : vector<8xf32>
+}

mlir/test/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTestTransforms
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
   TestNumberOfExecutions.cpp
   TestOpaqueLoc.cpp
+  TestPolynomialApproximation.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSCFUtils.cpp
   TestSparsification.cpp
   TestVectorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR

mlir/test/lib/Transforms/TestPolynomialApproximation.cpp

+//===- TestPolynomialApproximation.cpp - Test math ops approximations -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for expanding math operations into
+// polynomial approximations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestMathPolynomialApproximationPass
+    : public PassWrapper<TestMathPolynomialApproximationPass, FunctionPass> {
+  void runOnFunction() override;
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<vector::VectorDialect, math::MathDialect>();
+  }
+};
+} // end anonymous namespace
+
+void TestMathPolynomialApproximationPass::runOnFunction() {
+  OwningRewritePatternList patterns;
+  populateMathPolynomialApproximationPatterns(patterns, &getContext());
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+namespace mlir {
+namespace test {
+void registerTestMathPolynomialApproximationPass() {
+  PassRegistration<TestMathPolynomialApproximationPass> pass(
+      "test-math-polynomial-approximation",
+      "Test math polynomial approximations");
+}
+} // namespace test
+} // namespace mlir

mlir/test/mlir-cpu-runner/math_polynomial_approx.mlir

+// RUN:   mlir-opt %s -test-math-polynomial-approximation                      \
+// RUN:               -convert-vector-to-llvm                                  \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:     -e main -entry-point-result=void -O0                               \
+// RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext  \
+// RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext    \
+// RUN: | FileCheck %s
+
+
+func @main() {
+  // ------------------------------------------------------------------------ //
+  // Tanh.
+  // ------------------------------------------------------------------------ //
+
+  // CHECK: 0.848284
+  %0 = constant 1.25 : f32
+  %1 = math.tanh %0 : f32
+  vector.print %1 : f32
+
+  // CHECK: 0.244919, 0.635149, 0.761594, 0.848284
+  %2 = constant dense<[0.25, 0.75, 1.0, 1.25]> : vector<4xf32>
+  %3 = math.tanh %2 : vector<4xf32>
+  vector.print %3 : vector<4xf32>
+
+  // CHECK: 0.099668, 0.197375, 0.291313, 0.379949, 0.462117, 0.53705, 0.604368, 0.664037
+  %4 = constant dense<[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]> : vector<8xf32>
+  %5 = math.tanh %4 : vector<8xf32>
+  vector.print %5 : vector<8xf32>
+
+  return
+}

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
+void registerTestMathPolynomialApproximationPass();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
#endif
   test::registerTestLinalgGreedyFusion();
   test::registerTestLinalgHoisting();
   test::registerTestLinalgTileAndFuseSequencePass();
   test::registerTestLinalgTransforms();
   test::registerTestLivenessPass();
   test::registerTestLoopFusion();
   test::registerTestLoopMappingPass();
   test::registerTestLoopUnrollingPass();
+  test::registerTestMathPolynomialApproximationPass();
   test::registerTestMemRefDependenceCheck();
   test::registerTestMemRefStrideCalculation();
   test::registerTestNumberOfBlockExecutionsPass();
   test::registerTestNumberOfOperationExecutionsPass();
   test::registerTestOpaqueLoc();
   test::registerTestPDLByteCodePass();
   test::registerTestRecursiveTypesPass();
   test::registerTestSCFUtilsPass();