Add a phase that lowers std.tanh to a C runtime call

[llitchev]

Summary

The TanhOp is not properly lowered to LLVM. To address that a phase was added to lower the TanhOp to a C runtime call.

[llitchev]

Code formatting changes.

[rriddle]

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

19

This doesn't really seem like the right granularity for a pass. Seems better to just expose the patterns and use a test pass instead.

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

27

Can you add proper documentation here?

[ftynse]

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

1

Please make sure the header fits into 80 cols

54

This is wrong. You should use the pattern rewriter provided as function argument to _create_ ops as well.

56

rewriter.replaceWithNewOp<CallOp>(op, ...)

62

This function should take a PatternRewriter and use it to create the functions instead of defining local builders.

mlir/test/Dialect/Standard/tanh.mlir

2

Can we also check that the appropriate function references were inserted?

18

This file needs a better name

[llitchev]

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

19

Thanks!! Could you elaborate a bit more on this - using a test pass. Are there examples I can look at? I just modeled this after the ExpandAtomic pass (that was suggested to use as an example - https://reviews.llvm.org/D75509).

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

 
 namespace mlir {
 
 class Pass;
 
 /// Creates an instance of the ExpandAtomic pass.
 std::unique_ptr<Pass> createExpandAtomicPass();
 
+/// Creates an instance of the TanhLowering pass.
+std::unique_ptr<Pass> createTanhLoweringPass();
+
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_

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

 
 include "mlir/Pass/PassBase.td"
 
 def ExpandAtomic : FunctionPass<"expand-atomic"> {
   let summary = "Expands AtomicRMWOp into GenericAtomicRMWOp.";
   let constructor = "mlir::createExpandAtomicPass()";
 }
 
+def TanhLowering : FunctionPass<"tanh-lowering"> {
+  let summary = "Lowers TanhOp to a runtime call.";
+  let constructor = "mlir::createTanhLoweringPass()";
+}
+
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   ExpandAtomic.cpp
   FuncConversions.cpp
+  StdOpToCall.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/StdOpToCall.cpp

+//===- StdOpToCall.cpp - Code to perform tanh lowering to appropriate runtime
+// call ---------------------===//
+//
+// 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 lowering of TanhOp to a C runtime function call.
+//
+//===----------------------------------------------------------------------===//
+
+#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/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+
+template <typename OpType>
+struct StdOpToCallConverter : public OpRewritePattern<OpType> {
+public:
+  StdOpToCallConverter(MLIRContext *context, Type matchingType,
+                       StringRef funcName)
+      : OpRewritePattern<OpType>(context), matchingType(matchingType),
+        funcName(funcName) {}
+
+  LogicalResult matchAndRewrite(OpType op,
+                                PatternRewriter &rewriter) const final {
+    auto resultTypes = op.getOperation()->getResultTypes();
+    for (auto type : resultTypes) {
+      if (type != matchingType)
+        return failure();
+    }
+
+    auto operands = op.getOperation()->getOperands();
+    SmallVector<Type, 2> operandTypes;
+    for (auto operand : operands) {
+      auto type = operand.getType();
+      if (type != matchingType)
+        return failure();
+      operandTypes.push_back(type);
+    }
+
+    auto module = op.template getParentOfType<ModuleOp>();
+    auto callee = getOrInsertFunc(module, funcName, resultTypes, operandTypes,
+                                  op.getLoc());
+    OpBuilder builder(op);
+    auto callOp =
+        builder.create<CallOp>(op.getLoc(), callee, resultTypes, operands);
+    rewriter.replaceOp(op, callOp.getResults());
+    return success();
+  }
+
+private:
+  static FlatSymbolRefAttr getOrInsertFunc(ModuleOp module, StringRef symbol,
+                                           ArrayRef<Type> resultTypes,
+                                           ArrayRef<Type> operandTypes,
+                                           Location loc) {
+    auto context = module.getContext();
+    if (module.lookupSymbol(symbol)) {
+      return SymbolRefAttr::get(symbol, context);
+    }
+    OpBuilder builder(module.getBodyRegion());
+    auto funcType = builder.getFunctionType(resultTypes, operandTypes);
+    ArrayRef<NamedAttribute> attrs;
+    builder.create<FuncOp>(loc, symbol, funcType, attrs);
+    return SymbolRefAttr::get(symbol, context);
+  }
+
+  Type matchingType;
+  StringRef funcName;
+};
+
+struct TanhLowering : public TanhLoweringBase<TanhLowering> {
+  void runOnFunction() override {
+    auto &context = getContext();
+
+    OwningRewritePatternList patterns;
+    patterns.insert<StdOpToCallConverter<TanhOp>>(
+        &context, FloatType::getF32(&context), "tanhf");
+    patterns.insert<StdOpToCallConverter<TanhOp>>(
+        &context, FloatType::getF64(&context), "tanh");
+
+    ConversionTarget target(context);
+    target.addIllegalOp<TanhOp>();
+    if (failed(applyPartialConversion(getFunction(), target, patterns))) {
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::createTanhLoweringPass() {
+  return std::make_unique<TanhLowering>();
+}
 No newline at end of file

mlir/test/Dialect/Standard/tanh.mlir

+// RUN: mlir-opt -tanh-lowering %s | FileCheck %s
+
+// CHECK-LABEL: @tanh_f64
+func @tanh_f64() {
+  %cst = constant 1.0 : f64
+  // CHECK: call @tanh(%{{.*}}) : (f64) -> f64
+  %1 = tanh %cst : f64
+  return
+}
+
+// CHECK-LABEL: @tanh_f32
+func @tanh_f32() {
+  %cst = constant 1.0 : f32
+  // CHECK: call @tanhf(%{{.*}}) : (f32) -> f32
+  %1 = tanh %cst : f32
+  return
+}
 No newline at end of file