[mlir][EmitC] Add MathToEmitC pass for math function lowering to EmitC

mlir/include/mlir/Conversion/MathToEmitC/MathToEmitC.h

@@ -0,0 +1,21 @@
+//===- MathToEmitC.h - Math to EmitCPatterns -------------------*- 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 MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITC_H
+#define MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITC_H
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+
+namespace mlir {
+class RewritePatternSet;
+
+void populateConvertMathToEmitCPatterns(
+    RewritePatternSet &patterns,
+    emitc::MathToEmitCLanguageTarget languageTarget);
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITC_H

mlir/include/mlir/Conversion/MathToEmitC/MathToEmitCPass.h

@@ -0,0 +1,21 @@
+//===- MathToEmitCPass.h - Math to EmitC Pass -------------------*- 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 MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITCPASS_H
+#define MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITCPASS_H
+
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include <memory>
+namespace mlir {
+class Pass;
+
+#define GEN_PASS_DECL_CONVERTMATHTOEMITC
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITCPASS_H

mlir/include/mlir/Conversion/Passes.h

@@ -43,6 +43,7 @@
 #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"
 #include "mlir/Conversion/IndexToSPIRV/IndexToSPIRV.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
+#include "mlir/Conversion/MathToEmitC/MathToEmitCPass.h"
 #include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MathToLibm/MathToLibm.h"

mlir/include/mlir/Conversion/Passes.td

@@ -780,6 +780,28 @@ def ConvertMathToSPIRV : Pass<"convert-math-to-spirv"> {
   let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
+//===----------------------------------------------------------------------===//
+// MathToEmitC
+//===----------------------------------------------------------------------===//
+
+def ConvertMathToEmitC  : Pass<"convert-math-to-emitc"> {
+  let summary = "Convert some Math operations to EmitC Call_opaque";
+  let description = [{
+    This pass converts supported Math ops to `opaque_call` ops targeting libc/libm
+    functions. Unlike convert-math-to-funcs pass, converting to `call_opaque` ops
+    allows to overload the same function with different argument types.
+  }];
+  let dependentDialects = ["emitc::EmitCDialect"];
+  let options = [
+    Option<"languageTarget", "language-target", "::mlir::emitc::MathToEmitCLanguageTarget",
+      /*default=*/"::mlir::emitc::MathToEmitCLanguageTarget::CPP", "Select the language target for callees (C or CPP).",
+      [{::llvm::cl::values(
+            clEnumValN(::mlir::emitc::MathToEmitCLanguageTarget::C, "C", "C"),
+            clEnumValN(::mlir::emitc::MathToEmitCLanguageTarget::CPP, "CPP", "CPP")
+          )}]>
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // MathToFuncs
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/EmitC/IR/EmitCAttributes.td

@@ -62,4 +62,13 @@ def EmitC_OpaqueAttr : EmitC_Attr<"Opaque", "opaque"> {
 
 def EmitC_OpaqueOrTypedAttr : AnyAttrOf<[EmitC_OpaqueAttr, TypedAttrInterface]>;
 
+def MathToEmitCLanguageTarget : I32EnumAttr<"MathToEmitCLanguageTarget",
+  "Specifies the language target for generating callees.", [
+    I32EnumAttrCase<"C", 0, "Use C-style function names">,
+    I32EnumAttrCase<"CPP", 1, "Use C++-style function names">
+  ]> {
+  let cppNamespace = "::mlir::emitc";
+}
+
+
 #endif // MLIR_DIALECT_EMITC_IR_EMITCATTRIBUTES

mlir/lib/Conversion/CMakeLists.txt

@@ -33,6 +33,7 @@ add_subdirectory(IndexToLLVM)
 add_subdirectory(IndexToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(LLVMCommon)
+add_subdirectory(MathToEmitC)
 add_subdirectory(MathToFuncs)
 add_subdirectory(MathToLibm)
 add_subdirectory(MathToLLVM)

mlir/lib/Conversion/MathToEmitC/CMakeLists.txt

@@ -0,0 +1,19 @@
+add_mlir_conversion_library(MLIRMathToEmitC
+  MathToEmitC.cpp
+  MathToEmitCPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/MathToEmitC
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIREmitCDialect
+  MLIRMathDialect
+  MLIRPass
+  MLIRTransforms
+)

mlir/lib/Conversion/MathToEmitC/MathToEmitC.cpp

@@ -0,0 +1,83 @@
+//===- MathToEmitC.cpp - Math to EmitC  Patterns ----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/MathToEmitC/MathToEmitC.h"
+
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+template <typename OpType>
+class LowerToEmitCCallOpaque : public OpRewritePattern<OpType> {
+  std::string calleeStr;
+  emitc::MathToEmitCLanguageTarget languageTarget;
+
+public:
+  LowerToEmitCCallOpaque(MLIRContext *context, std::string calleeStr,
+                         emitc::MathToEmitCLanguageTarget languageTarget)
+      : OpRewritePattern<OpType>(context), calleeStr(std::move(calleeStr)),
+        languageTarget(languageTarget) {}
+
+  LogicalResult matchAndRewrite(OpType op,
+                                PatternRewriter &rewriter) const override;
+};
+
+template <typename OpType>
+LogicalResult LowerToEmitCCallOpaque<OpType>::matchAndRewrite(
+    OpType op, PatternRewriter &rewriter) const {
+  if (!llvm::all_of(op->getOperandTypes(), llvm::IsaPred<Float32Type, Float64Type>)||
+      !llvm::all_of(op->getResultTypes(),llvm::IsaPred<Float32Type, Float64Type>))
+    return rewriter.notifyMatchFailure(
+        op.getLoc(), "expected all operands and results to be of type f32");
+  std::string modifiedCalleeStr = calleeStr;
+  if (languageTarget == emitc::MathToEmitCLanguageTarget::CPP) {
+    modifiedCalleeStr = "std::" + calleeStr;
+  } else if (languageTarget == emitc::MathToEmitCLanguageTarget::C) {
+    auto operandType = op->getOperandTypes()[0];
+    if (operandType.isF32())
+      modifiedCalleeStr = calleeStr + "f";
+  }
+  rewriter.replaceOpWithNewOp<emitc::CallOpaqueOp>(
+      op, op.getType(), modifiedCalleeStr, op->getOperands());
+  return success();
+}
+
+} // namespace
+
+// Populates patterns to replace `math` operations with `emitc.call_opaque`,
+// using function names consistent with those in <math.h>.
+void mlir::populateConvertMathToEmitCPatterns(
+    RewritePatternSet &patterns,
+    emitc::MathToEmitCLanguageTarget languageTarget) {
+  auto *context = patterns.getContext();
+  patterns.insert<LowerToEmitCCallOpaque<math::FloorOp>>(context, "floor",
+                                                         languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::RoundOp>>(context, "round",
+                                                         languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::ExpOp>>(context, "exp",
+                                                       languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::CosOp>>(context, "cos",
+                                                       languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::SinOp>>(context, "sin",
+                                                       languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::AcosOp>>(context, "acos",
+                                                        languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::AsinOp>>(context, "asin",
+                                                        languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::Atan2Op>>(context, "atan2",
+                                                         languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::CeilOp>>(context, "ceil",
+                                                        languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::AbsFOp>>(context, "fabs",
+                                                        languageTarget);
+  patterns.insert<LowerToEmitCCallOpaque<math::PowFOp>>(context, "pow",
+                                                        languageTarget);
+}

mlir/lib/Conversion/MathToEmitC/MathToEmitCPass.cpp

@@ -0,0 +1,52 @@
+//===- MathToEmitCPass.cpp - Math to EmitC Pass -----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to convert the Math dialect to the EmitC dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/MathToEmitC/MathToEmitCPass.h"
+#include "mlir/Conversion/MathToEmitC/MathToEmitC.h"
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTMATHTOEMITC
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+namespace {
+
+//  Replaces Math operations with `emitc.call_opaque` operations.
+struct ConvertMathToEmitC
+    : public impl::ConvertMathToEmitCBase<ConvertMathToEmitC> {
+  using ConvertMathToEmitCBase::ConvertMathToEmitCBase;
+
+public:
+  void runOnOperation() final;
+};
+
+} // namespace
+
+void ConvertMathToEmitC::runOnOperation() {
+  ConversionTarget target(getContext());
+  target.addLegalOp<emitc::CallOpaqueOp>();
+
+  target.addIllegalOp<math::FloorOp, math::ExpOp, math::RoundEvenOp,
+                      math::CosOp, math::SinOp, math::Atan2Op, math::CeilOp,
+                      math::AcosOp, math::AsinOp, math::AbsFOp, math::PowFOp>();
+
+  RewritePatternSet patterns(&getContext());
+  populateConvertMathToEmitCPatterns(patterns, languageTarget);
+
+  if (failed(applyPartialConversion(getOperation(), target, std::move(patterns))))
+    signalPassFailure();
+}

mlir/test/Conversion/MathToEmitC/math-to-emitc-failed.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt -split-input-file -convert-math-to-emitc -verify-diagnostics %s 
+
+func.func @unsupported_tensor_type(%arg0 : tensor<4xf32>) -> tensor<4xf32> {
+// expected-error @+1 {{failed to legalize operation 'math.absf' that was explicitly marked illegal}}
+  %0 = math.absf %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+func.func @unsupported_f16_type(%arg0 : f16) -> f16 {
+// expected-error @+1 {{failed to legalize operation 'math.absf' that was explicitly marked illegal}}
+  %0 = math.absf %arg0 : f16
+  return %0 : f16
+}
+
+// -----
+
+func.func @unsupported_f128_type(%arg0 : f128) -> f128 {
+// expected-error @+1 {{failed to legalize operation 'math.absf' that was explicitly marked illegal}}
+  %0 = math.absf %arg0 : f128
+  return %0 : f128
+}

mlir/test/Conversion/MathToEmitC/math-to-emitc.mlir

@@ -0,0 +1,65 @@
+// RUN: mlir-opt -convert-math-to-emitc=language-target=C %s | FileCheck %s --check-prefix=C
+// RUN: mlir-opt -convert-math-to-emitc=language-target=CPP %s | FileCheck %s --check-prefix=CPP
+
+func.func @absf_to_call_opaque(%arg0: f32) {
+    // C: emitc.call_opaque "fabsf"
+    // CPP: emitc.call_opaque "std::fabs"
+    %1 = math.absf %arg0 : f32
+    return
+  }
+func.func @floor_to_call_opaque(%arg0: f32) {
+    // C: emitc.call_opaque "floorf"
+    // CPP: emitc.call_opaque "std::floor"
+    %1 = math.floor %arg0 : f32
+    return
+  }
+func.func @sin_to_call_opaque(%arg0: f32) {
+    // C: emitc.call_opaque "sinf"
+    // CPP: emitc.call_opaque "std::sin"  
+    %1 = math.sin %arg0 : f32
+    return
+  }
+func.func @cos_to_call_opaque(%arg0: f32) {
+    // C: emitc.call_opaque "cosf"
+    // CPP: emitc.call_opaque "std::cos"    
+    %1 = math.cos %arg0 : f32
+    return
+  }
+func.func @asin_to_call_opaque(%arg0: f32) {
+    // C: emitc.call_opaque "asinf"
+    // CPP: emitc.call_opaque "std::asin"     
+    %1 = math.asin %arg0 : f32
+    return
+  }
+func.func @acos_to_call_opaque(%arg0: f64) {
+    // C: emitc.call_opaque "acos"
+    // CPP: emitc.call_opaque "std::acos"      
+    %1 = math.acos %arg0 : f64
+    return
+  }
+func.func @atan2_to_call_opaque(%arg0: f64, %arg1: f64) {
+    // C: emitc.call_opaque "atan2"
+    // CPP: emitc.call_opaque "std::atan2"       
+    %1 = math.atan2 %arg0, %arg1 : f64
+    return
+  }
+func.func @ceil_to_call_opaque(%arg0: f64) {
+    // C: emitc.call_opaque "ceil"
+    // CPP: emitc.call_opaque "std::ceil"    
+    %1 = math.ceil %arg0 : f64
+    return
+  }
+func.func @exp_to_call_opaque(%arg0: f64) {
+    // C: emitc.call_opaque "exp"
+    // CPP: emitc.call_opaque "std::exp"    
+    %1 = math.exp %arg0 : f64
+    return
+  }
+func.func @powf_to_call_opaque(%arg0: f64, %arg1: f64) {
+    // C: emitc.call_opaque "pow"
+    // CPP: emitc.call_opaque "std::pow"    
+    %1 = math.powf %arg0, %arg1 : f64
+    return
+  }
+
+

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -4201,6 +4201,7 @@ cc_library(
         ":IndexToLLVM",
         ":IndexToSPIRV",
         ":LinalgToStandard",
+        ":MathToEmitC",
         ":MathToFuncs",
         ":MathToLLVM",
         ":MathToLibm",
@@ -8721,6 +8722,27 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "MathToEmitC",
+    srcs = glob([
+        "lib/Conversion/MathToEmitC/*.cpp",
+    ]),
+    hdrs = glob([
+        "include/mlir/Conversion/MathToEmitC/*.h",
+    ]),
+    includes = [
+        "include",
+        "lib/Conversion/MathToEmitC",
+    ],
+    deps = [
+        ":ConversionPassIncGen",
+        ":EmitCDialect",
+        ":MathDialect",
+        ":Pass",
+        ":TransformUtils",
+    ],
+)
+
 cc_library(
     name = "MathToFuncs",
     srcs = glob(["lib/Conversion/MathToFuncs/*.cpp"]),