[mlir] Add support for lowering tanh to LLVMIR.

[hanchung]

Summary

Add a pattern for expanding tanh op into exp form.
A tanh is expanded into:

1. 1-exp^{-2x} / 1+exp^{-2x}, if x => 0
2. exp^{2x}-1 / exp^{2x}+1 , if x < 0.

[asaadaldien]

Didn't go through the code yet. high level comments:

• Is it possible to get this lowered to a tanh function call instead of the exp rewrite ?

• If we have to rewrite tanh(x) then lets express it in a numerically stable way because tanh(x) = exp^{2x}-1 / exp^{2x}+1 will blow up for relatively not very big positive values of x, so we can do:

tanh(x) = 1-exp^{-2x} / 1+exp^{-2x}, x => 0
tanh(x) = exp^{2x}-1 / exp^{2x}+1, x < 0

[ftynse]

It looks like this is better done as an intra-dialect transform, everything that you need is available in the standard dialect so the code can just expand std.tanh to a set of std. operations and then use the regular conversion. We already do so for atomics https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/StandardOps/Transforms/ExpandAtomic.cpp, we can put ExpandTanh.cpp next to it.

[ftynse]

Is it possible to get this lowered to a tanh function call instead of the exp rewrite ?

Unless we assume libc is linked, there's no tanh function.

[asaadaldien]

In D81618#2086974, @ftynse wrote:

Is it possible to get this lowered to a tanh function call instead of the exp rewrite ?

Unless we assume libc is linked, there's no tanh function.

other functions e.g llvm.inter.exp aren't libc mathlib calls ?

[ftynse]

In D81618#2088020, @asaadaldien wrote:

In D81618#2086974, @ftynse wrote:

Is it possible to get this lowered to a tanh function call instead of the exp rewrite ?

Unless we assume libc is linked, there's no tanh function.

other functions e.g llvm.inter.exp aren't libc mathlib calls ?

It depends on the LLVM backend. LLVM models _all_ intrinsics as function calls, but it does not necessarily mean they actually become function calls in the generated assembly (e.g., llvm.intr.fma won't).

[hanchung]

Add expanding tanh pattern to StandardOps/Transforms/

[hanchung]

Add a test

[hanchung]

Add missing files...

[hanchung]

Thanks for the advice, I think it's better now. :)

[ftynse]

Thanks!

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

30

Does something actually use it as a pass? I only see the pattern to be used inside Linalg conversion. If there is no direct user of this pass, let's put it under test/lib/Transforms.

mlir/lib/Dialect/StandardOps/Transforms/ExpandTanh.cpp

2

It does not perform loop fusion ;)

36

Please don't commit commented-out code

44

Could we have a better name? Like doubledX ?

[asaadaldien]

mlir/lib/Dialect/StandardOps/Transforms/ExpandTanh.cpp

38

return failure() if operand isn't float ?

[hanchung]

Address comments.

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

30

Done, good to know this, thank you!

mlir/lib/Dialect/StandardOps/Transforms/ExpandTanh.cpp

2

Good catch, thanks! I copied it from ExpandAtomic, let me send out another patch to fix that part.

36

Forgot to remove debug code, done.

38

I think we don't need this check because it's guaranteed that the type of the operand must be float.

class FloatUnaryOp<string mnemonic, list<OpTrait> traits = []> :
    UnaryOpSameOperandAndResultType<mnemonic, traits>,
    Arguments<(ins FloatLike:$operand)>;
...
def TanhOp : FloatUnaryOp<"tanh">
 ...

44

Done, also fixed the name in the test.

[ftynse]

Thanks for iterating!

[mehdi_amini]

Reverted in a9a21bb4b682474248dc85f9e7db4b260d249ab9 ; the standalone test was broken

[hanchung]

In D81618#2093524, @mehdi_amini wrote:

Reverted in a9a21bb4b682474248dc85f9e7db4b260d249ab9 ; the standalone test was broken

It passed with bazel on my local machine, I'll take a look at what happened, sorry about that.

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

 #ifndef MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 
 #include <memory>
 
 namespace mlir {
 
 class Pass;
+class MLIRContext;
+class OwningRewritePatternList;
 
 /// Creates an instance of the ExpandAtomic pass.
 std::unique_ptr<Pass> createExpandAtomicPass();
 
+void populateExpandTanhPattern(OwningRewritePatternList &patterns,
+                               MLIRContext *ctx);
+
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
void ConvertLinalgToLLVMPass::runOnOperation() {
   auto module = getOperation();
 
   // Convert to the LLVM IR dialect using the converter defined above.
   OwningRewritePatternList patterns;
   LLVMTypeConverter converter(&getContext());
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());
   populateStdToLLVMConversionPatterns(converter, patterns);
+  populateExpandTanhPattern(patterns, &getContext());
   populateVectorToSCFConversionPatterns(patterns, &getContext());
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext());
 
   LLVMConversionTarget target(getContext());
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
   if (failed(applyFullConversion(module, target, patterns, &converter)))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertLinalgToLLVMPass() {
   return std::make_unique<ConvertLinalgToLLVMPass>();
 }

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

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   ExpandAtomic.cpp
+  ExpandTanh.cpp
   FuncConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRStandardOps
   MLIRTransforms
   )

mlir/lib/Dialect/StandardOps/Transforms/ExpandTanh.cpp

+//===- ExpandTanh.cpp - Code to perform expanding tanh op -----------------===//
+//
+// 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 tanh op.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+
+/// Expands tanh op into
+///   1) 1-exp^{-2x} / 1+exp^{-2x}, if x => 0
+///   2) exp^{2x}-1 / exp^{2x}+1  , if x < 0
+struct TanhOpConverter : public OpRewritePattern<TanhOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TanhOp op,
+                                PatternRewriter &rewriter) const final {
+    auto floatType = op.operand().getType();
+    Location loc = op.getLoc();
+    auto floatOne = rewriter.getFloatAttr(floatType, 1.0);
+    auto floatTwo = rewriter.getFloatAttr(floatType, 2.0);
+    Value one = rewriter.create<ConstantOp>(loc, floatOne);
+    Value two = rewriter.create<ConstantOp>(loc, floatTwo);
+    Value doubledX = rewriter.create<MulFOp>(loc, op.operand(), two);
+
+    // Case 1: tanh(x) = 1-exp^{-2x} / 1+exp^{-2x}
+    Value negDoubledX = rewriter.create<NegFOp>(loc, doubledX);
+    Value exp2x = rewriter.create<ExpOp>(loc, negDoubledX);
+    Value dividend = rewriter.create<SubFOp>(loc, one, exp2x);
+    Value divisor = rewriter.create<AddFOp>(loc, one, exp2x);
+    Value positiveRes = rewriter.create<DivFOp>(loc, dividend, divisor);
+
+    // Case 2: tanh(x) = exp^{2x}-1 / exp^{2x}+1
+    exp2x = rewriter.create<ExpOp>(loc, doubledX);
+    dividend = rewriter.create<SubFOp>(loc, exp2x, one);
+    divisor = rewriter.create<AddFOp>(loc, exp2x, one);
+    Value negativeRes = rewriter.create<DivFOp>(loc, dividend, divisor);
+
+    // tanh(x) = x >= 0 ? positiveRes : negativeRes
+    auto floatZero = rewriter.getFloatAttr(floatType, 0.0);
+    Value zero = rewriter.create<ConstantOp>(loc, floatZero);
+    Value cmpRes =
+        rewriter.create<CmpFOp>(loc, CmpFPredicate::OGE, op.operand(), zero);
+    rewriter.replaceOpWithNewOp<SelectOp>(op, cmpRes, positiveRes, negativeRes);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateExpandTanhPattern(OwningRewritePatternList &patterns,
+                                     MLIRContext *ctx) {
+  patterns.insert<TanhOpConverter>(ctx);
+}

mlir/test/Dialect/Standard/expand-tanh.mlir

+// RUN: mlir-opt %s -test-expand-tanh | FileCheck %s
+
+// CHECK-LABEL: func @tanh
+func @tanh(%arg: f32) -> f32 {
+  %res = tanh %arg : f32
+  return %res : f32
+}
+// CHECK-DAG: %[[ZERO:.+]] = constant 0.000000e+00 : f32
+// CHECK-DAG: %[[ONE:.+]] = constant 1.000000e+00 : f32
+// CHECK-DAG: %[[TWO:.+]] = constant 2.000000e+00 : f32
+// CHECK: %[[DOUBLEDX:.+]] = mulf %arg0, %[[TWO]] : f32
+// CHECK: %[[NEGDOUBLEDX:.+]] = negf %[[DOUBLEDX]] : f32
+// CHECK: %[[EXP1:.+]] = exp %[[NEGDOUBLEDX]] : f32
+// CHECK: %[[DIVIDEND1:.+]] = subf %[[ONE]], %[[EXP1]] : f32
+// CHECK: %[[DIVISOR1:.+]] = addf %[[ONE]], %[[EXP1]] : f32
+// CHECK: %[[RES1:.+]] = divf %[[DIVIDEND1]], %[[DIVISOR1]] : f32
+// CHECK: %[[EXP2:.+]] = exp %[[DOUBLEDX]] : f32
+// CHECK: %[[DIVIDEND2:.+]] = subf %[[EXP2]], %[[ONE]] : f32
+// CHECK: %[[DIVISOR2:.+]] = addf %[[EXP2]], %[[ONE]] : f32
+// CHECK: %[[RES2:.+]] = divf %[[DIVIDEND2]], %[[DIVISOR2]] : f32
+// CHECK: %[[COND:.+]] = cmpf "oge", %arg0, %[[ZERO]] : f32
+// CHECK: %[[RESULT:.+]] = select %[[COND]], %[[RES1]], %[[RES2]] : f32
+// CHECK: return %[[RESULT]]

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
   TestBufferPlacement.cpp
+  TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp

mlir/test/lib/Transforms/TestExpandTanh.cpp

+//===- TestExpandTanh.cpp - Test expand tanh op into exp form ------===//
+//
+// 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 tanh.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestExpandTanhPass
+    : public PassWrapper<TestExpandTanhPass, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+void TestExpandTanhPass::runOnFunction() {
+  OwningRewritePatternList patterns;
+  populateExpandTanhPattern(patterns, &getContext());
+  applyPatternsAndFoldGreedily(getOperation(), patterns);
+}
+
+namespace mlir {
+void registerTestExpandTanhPass() {
+  PassRegistration<TestExpandTanhPass> pass("test-expand-tanh",
+                                            "Test expanding tanh");
+}
+} // namespace mlir

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

 void registerTestAffineLoopUnswitchingPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestLoopPermutationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
+void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLinalgHoisting();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
#if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
   registerTestFunc();
+  registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
   registerTestMatchers();