[mlir] Add support for expanding tanh op into exp form.

[hanchung]

Summary

Fixed build of D81618

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.

[hanchung]

It's fixed by adding a missing dep in mlir/lib/Conversion/LinalgToLLVM/CMakeLists.txt

[hanchung]

Oh wait me a sec, I think I should populate the pattern at populateStdToLLVMConversionPatterns, not in LinalgToLLVM.cpp file. Let me fix it.

[hanchung]

Remove the pattern from ConvertLinalgToLLVMPass.

[hanchung]

Remove unneeded include from mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

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

add_mlir_conversion_library(MLIRLinalgToLLVM
 
   LINK_LIBS PUBLIC
   MLIRAffineToStandard
   MLIREDSC
   MLIRIR
   MLIRLinalgOps
   MLIRLLVMIR
   MLIRSCFToStandard
+  MLIRStandardOpsTransforms
   MLIRStandardToLLVM
   MLIRTransforms
   MLIRVectorToLLVM
   MLIRVectorToSCF
   )

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();