[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

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

204

Can you expose these via a populate method?

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

8

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 ↗

(On Diff #324236)

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

148 ↗

(On Diff #324236)

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

6 ↗

(On Diff #324807)

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

[ezhulenev]

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

10 ↗

(On Diff #324236)

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 ↗

(On Diff #324236)

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

6 ↗

(On Diff #324807)

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/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

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

+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Math)
+add_public_tablegen_target(MLIRMathPassIncGen)
+
+add_mlir_doc(Passes -gen-pass-doc MathPasses ./)

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

 
 namespace mlir {
 
 class OwningRewritePatternList;
 
 void populateExpandTanhPattern(OwningRewritePatternList &patterns,
                                MLIRContext *ctx);
 
+std::unique_ptr<OperationPass<FuncOp>> createFastMathExpansionPass();
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

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

+//===-- Passes.td - Math pass definition file --------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MATH_TRANSORMS_PASSES
+#define MLIR_DIALECT_MATH_TRANSORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def FastMathExpansion : FunctionPass<"fast-math-expansion"> {
+  let summary = "Expands math operations to fast approximations that do not "
+                "rely on any library functions.";
+  let constructor = "mlir::createFastMathExpansionPass()";
+  let dependentDialects = ["math::MathDialect", "vector::VectorDialect"];
+}
+
+#endif // MLIR_DIALECT_MATH_TRANSORMS_PASSES

mlir/include/mlir/InitAllPasses.h

 #define MLIR_INITALLPASSES_H_
 
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/Quant/Passes.h"
 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
 #include "mlir/Transforms/Passes.h"
inline void registerAllPasses() {
 
   // Dialect passes
   registerAffinePasses();
   registerAsyncPasses();
   registerGPUPasses();
   registerLinalgPasses();
   LLVM::registerLLVMPasses();
   quant::registerQuantPasses();
+  registerMathPasses();
   registerSCFPasses();
   registerShapePasses();
   spirv::registerSPIRVPasses();
   registerStandardPasses();
   tensor::registerTensorPasses();
   tosa::registerTosaOptPasses();
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

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

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

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

+//===- FastMathExpansion.cpp - Expand math ops to fast 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 implements expansion of math operations to fast approximations
+// that do not rely on any of the library functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#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 FloatExpansionBuilder {
+public:
+  FloatExpansionBuilder(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
+
+FloatExpansionBuilder::FloatExpansionBuilder(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 FloatExpansionBuilder::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 FloatExpansionBuilder::abs(Value a) const {
+  return rewriter.create<AbsFOp>(loc, a);
+}
+
+Value FloatExpansionBuilder::min(Value a, Value b) const {
+  return select(cmp(CmpFPredicate::OLT, a, b), a, b);
+}
+Value FloatExpansionBuilder::max(Value a, Value b) const {
+  return select(cmp(CmpFPredicate::OGT, a, b), a, b);
+}
+Value FloatExpansionBuilder::mul(Value a, Value b) const {
+  return rewriter.create<MulFOp>(loc, a, b);
+}
+
+Value FloatExpansionBuilder::div(Value a, Value b) const {
+  return rewriter.create<DivFOp>(loc, a, b);
+}
+
+Value FloatExpansionBuilder::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 FloatExpansionBuilder::cmp(CmpFPredicate predicate, Value a,
+                                 Value b) const {
+  return rewriter.create<CmpFOp>(loc, predicate, a, b);
+}
+
+Value FloatExpansionBuilder::select(Value predicate, Value a, Value b) const {
+  return rewriter.create<SelectOp>(loc, predicate, a, b);
+}
+
+//----------------------------------------------------------------------------//
+// TanhOp expansion.
+//----------------------------------------------------------------------------//
+
+namespace {
+struct TanhExpansion : public OpRewritePattern<math::TanhOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::TanhOp op,
+                                PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+LogicalResult TanhExpansion::matchAndRewrite(math::TanhOp op,
+                                             PatternRewriter &rewriter) const {
+  if (!isValidFloatType(op.operand().getType()))
+    return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+  Value operand = op.operand();
+  FloatExpansionBuilder 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();
+}
+
+//----------------------------------------------------------------------------//
+// Function pass that expands math operations to fast approximations.
+//----------------------------------------------------------------------------//
+
+namespace {
+class FastMathExpansionPass
+    : public FastMathExpansionBase<FastMathExpansionPass> {
+public:
+  FastMathExpansionPass() = default;
+  void runOnFunction() override;
+};
+} // namespace
+
+void FastMathExpansionPass::runOnFunction() {
+  mlir::MLIRContext *ctx = &getContext();
+  mlir::OwningRewritePatternList patterns;
+  patterns.insert<TanhExpansion>(ctx);
+
+  ConversionTarget target(*ctx);
+  target.addLegalDialect<StandardOpsDialect>();
+  target.addLegalDialect<vector::VectorDialect>();
+  target.addIllegalOp<math::TanhOp>();
+
+  if (failed(applyPartialConversion(getOperation(), target,
+                                    std::move(patterns)))) {
+    signalPassFailure();
+    return;
+  }
+}
+
+std::unique_ptr<OperationPass<FuncOp>> mlir::createFastMathExpansionPass() {
+  return std::make_unique<FastMathExpansionPass>();
+}

mlir/lib/Dialect/Math/Transforms/PassDetail.h

+//===- PassDetail.h - Math Pass class details -------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DIALECT_MATH_TRANSFORMS_PASSDETAIL_H_
+#define DIALECT_MATH_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+namespace math {
+class MathDialect;
+} // namespace math
+
+namespace vector {
+class VectorDialect;
+} // namespace vector
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // DIALECT_MATH_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/Math/fast-math.mlir

+// RUN: mlir-opt %s -fast-math-expansion | 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/mlir-cpu-runner/fast-math.mlir

+// RUN:   mlir-opt %s -fast-math-expansion                                     \
+// 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 --dump-input=always
+
+
+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
+}