Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marbre merged 6 commits into from
Jan 20, 2025
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d5bd00c
[mlir][EmitC] Add MathToEmitC pass for math function lowering to EmitC
recursion-man Oct 22, 2024
ad9af42
[MLIR][MathToEmitC] Ensure scalar type handling and refactor
recursion-man Oct 29, 2024
f6c2406
[MLIR][MathToEmitC] Refactor code, add tests for unsupported types, a…
recursion-man Oct 29, 2024
23ada46
[MLIR][MathToEmitC] Add support for C and C++ targets with Lit tests
recursion-man Oct 31, 2024
8992fd5
[MLIR][MathToEmitC] Add language standard option, create LanguageTarg…
recursion-man Nov 20, 2024
52c35a6
[mlir][EmitC] Apply formatting fixes
recursion-man Jan 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions mlir/include/mlir/Conversion/MathToEmitC/MathToEmitC.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
//===- MathToEmitC.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_MATHTOEMITC_H
#define MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITC_H

#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include <memory>

namespace mlir {

#define GEN_PASS_DECL_CONVERTMATHTOEMITC
#include "mlir/Conversion/Passes.h.inc"

std::unique_ptr<OperationPass<mlir::ModuleOp>> createConvertMathToEmitCPass();

} // namespace mlir

#endif // MLIR_CONVERSION_MATHTOEMITC_MATHTOEMITC_H
1 change: 1 addition & 0 deletions mlir/include/mlir/Conversion/Passes.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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/MathToEmitC.h"
#include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
#include "mlir/Conversion/MathToLibm/MathToLibm.h"
Expand Down
19 changes: 19 additions & 0 deletions mlir/include/mlir/Conversion/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -780,6 +780,25 @@ def ConvertMathToSPIRV : Pass<"convert-math-to-spirv"> {
let dependentDialects = ["spirv::SPIRVDialect"];
}

//===----------------------------------------------------------------------===//
// MathToEmitC
//===----------------------------------------------------------------------===//

def ConvertMathToEmitC : Pass<"convert-math-to-emitc", "ModuleOp"> {
let summary = "Convert some Math operations to EmitC Call_opaque";
let description = [{
This pass converts supported Math ops to call_opaque calls to compiler generated
functions implementing these operations in software.
Copy link
Contributor

@simon-camp simon-camp Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Maybe something like:
This pass converts supported Math ops to opaque_call ops targeting libc/libm functions.

Copy link
Contributor Author

@recursion-man recursion-man Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed to what you suggested.
We might add new targets as suggested in this thread

Unlike convert-math-to-funcs pass, this pass uses call_opaque,
therefore enables us to overload the same funtion with different argument types
}];

let constructor = "mlir::createConvertMathToEmitCPass()";
let dependentDialects = ["emitc::EmitCDialect",
"math::MathDialect"
];
}

//===----------------------------------------------------------------------===//
// MathToFuncs
//===----------------------------------------------------------------------===//
Expand Down
1 change: 1 addition & 0 deletions mlir/lib/Conversion/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
19 changes: 19 additions & 0 deletions mlir/lib/Conversion/MathToEmitC/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
add_mlir_conversion_library(MLIRMathToEmitC
MathToEmitC.cpp

ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/MathToEmitC

DEPENDS
MLIRConversionPassIncGen

LINK_COMPONENTS
Core

LINK_LIBS PUBLIC
MLIRLLVMCommonConversion
MLIREmitCDialect
MLIRMathDialect
MLIRPass
MLIRTransforms
)
99 changes: 99 additions & 0 deletions mlir/lib/Conversion/MathToEmitC/MathToEmitC.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@

//===- MathToEmitC.cpp - Math to EmitC Pass Implementation ----------===//
//
// 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/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 ConvertMathToEmitCPass
: public impl::ConvertMathToEmitCBase<ConvertMathToEmitCPass> {
public:
void runOnOperation() final;
};

} // end anonymous namespace

template <typename OpType>
class LowerToEmitCCallOpaque : public mlir::OpRewritePattern<OpType> {
Copy link
Contributor

@simon-camp simon-camp Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Drop all the mlir:: namespaces.

Copy link
Contributor Author

@recursion-man recursion-man Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, Deleted

std::string calleeStr;

public:
LowerToEmitCCallOpaque(MLIRContext *context, std::string calleeStr)
: OpRewritePattern<OpType>(context), calleeStr(calleeStr) {}

LogicalResult matchAndRewrite(OpType op,
PatternRewriter &rewriter) const override;
};

// Populates patterns to replace `math` operations with `emitc.call_opaque`,
// using function names consistent with those in <math.h>.
static void populateConvertMathToEmitCPatterns(RewritePatternSet &patterns) {
auto *context = patterns.getContext();
patterns.insert<LowerToEmitCCallOpaque<math::FloorOp>>(context, "floor");
patterns.insert<LowerToEmitCCallOpaque<math::RoundEvenOp>>(context, "rint");
Copy link
Contributor

@simon-camp simon-camp Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This may be unsafe, not sure. rint uses the current rounding mode, RoundToEven ignores the rounding mode.

Copy link
Contributor Author

@recursion-man recursion-man Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right- for now I support only round

patterns.insert<LowerToEmitCCallOpaque<math::ExpOp>>(context, "exp");
patterns.insert<LowerToEmitCCallOpaque<math::CosOp>>(context, "cos");
patterns.insert<LowerToEmitCCallOpaque<math::SinOp>>(context, "sin");
patterns.insert<LowerToEmitCCallOpaque<math::AcosOp>>(context, "acos");
patterns.insert<LowerToEmitCCallOpaque<math::AsinOp>>(context, "asin");
patterns.insert<LowerToEmitCCallOpaque<math::Atan2Op>>(context, "atan2");
patterns.insert<LowerToEmitCCallOpaque<math::CeilOp>>(context, "ceil");
patterns.insert<LowerToEmitCCallOpaque<math::AbsFOp>>(context, "fabs");
patterns.insert<LowerToEmitCCallOpaque<math::FPowIOp>>(context, "powf");
patterns.insert<LowerToEmitCCallOpaque<math::IPowIOp>>(context, "pow");
Copy link
Contributor

@cferry-AMD cferry-AMD Oct 28, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you have a look at the page for fabs for instance, you'll see that the C implementation (without overloading) bears a different name for each type used. How about a more generic pattern with a predicate:

patterns.insert<LowerToEmitCCallOpaque<math::AbsFOp>>(context, "fabsf", [](Type t){ return t == IntegerType::get(t.getContext(), 32); });
patterns.insert<LowerToEmitCCallOpaque<math::AbsFOp>>(context, "fabs", [](Type t){ return t == IntegerType::get(t.getContext(), 64); });

Copy link
Member

@marbre marbre Oct 28, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for addressing this @cferry-AMD. This is something that we definitely should consider. However, I think it depends on what language standard one targets (see https://en.cppreference.com/w/c/numeric/math/fabs and https://en.cppreference.com/w/cpp/numeric/math/fabs). What about an optional switch to specify which standard to target as part of the conversion?

Copy link
Contributor

@simon-camp simon-camp Oct 28, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also pass a pair of strings for f32 and f64 types. Or as an alternative a map<Type, string>.

Copy link
Contributor Author

@recursion-man recursion-man Oct 29, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for raising this! I propose two options for handling function names:

Option 1: Hardcoded names targeting C's math.h or C++'s
C target would use type-specific suffixes (since C lacks overloading) for example :

if (operandType.isF32()) return (baseName + "f").str();  // "fabsf"
else if (operandType.isF64()) return baseName.str();      // "fabs"

C++ target would use base names with std:: for example: std::fabs

Option 2: Custom library support via user-defined mapping

std::map<TypeID, std::string> customMap = {
    {TypeID::get<math::AbsFOp>(), "native_abs"},
    {TypeID::get<math::ExpOp>(), "native_exp"}
};
patterns.insert<LowerToEmitCCallOpaque<math::AbsFOp>>(context, customMap);

This option allows users to pass a custom mapping, supporting other libraries (e.g., OpenCL). If no mapping is provided, it defaults to math.h.

Thanks again for your input! I’d really like to hear your opinions on these options

Copy link
Contributor

@cferry-AMD cferry-AMD Nov 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How specific to Math functions do you see this pass? The second choice makes the lowering much more generic than I thought: it may apply to any op that could be lowered to opaque library calls, and not just Math dialect ops to math.h functions. I like this ability to lower anything to opaque calls with language filters, but this may not be what you intended.

Copy link
Contributor Author

@recursion-man recursion-man Nov 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the feedback!
I intended this pass to focus on lowering math functions specifically. I accepted the suggestion to use libraries like and am also considering the option to add custom math libraries, not just the standard ones, which motivated me to use a map. I can enforce that only math dialect ops will be converted to ensure it remains targeted.

}

template <typename OpType>
LogicalResult LowerToEmitCCallOpaque<OpType>::matchAndRewrite(
OpType op, PatternRewriter &rewriter) const {
mlir::StringAttr callee = rewriter.getStringAttr(calleeStr);
auto actualOp = mlir::cast<OpType>(op);
rewriter.replaceOpWithNewOp<mlir::emitc::CallOpaqueOp>(
actualOp, actualOp.getType(), callee, actualOp->getOperands());
return mlir::success();
}

void ConvertMathToEmitCPass::runOnOperation() {
auto moduleOp = getOperation();
// Insert #include <math.h> at the beginning of the module
OpBuilder builder(moduleOp.getBodyRegion());
builder.setInsertionPointToStart(&moduleOp.getBodyRegion().front());
builder.create<emitc::IncludeOp>(moduleOp.getLoc(),
builder.getStringAttr("math.h"));

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,
math::FPowIOp, math::IPowIOp>();

RewritePatternSet patterns(&getContext());
populateConvertMathToEmitCPatterns(patterns);

if (failed(applyPartialConversion(moduleOp, target, std::move(patterns))))
signalPassFailure();
}

std::unique_ptr<OperationPass<mlir::ModuleOp>>
mlir::createConvertMathToEmitCPass() {
return std::make_unique<ConvertMathToEmitCPass>();
}
140 changes: 140 additions & 0 deletions mlir/test/Conversion/MathToEmitC/math-to-emitc.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
// RUN: mlir-opt --split-input-file -convert-math-to-emitc %s | FileCheck %s

// CHECK-LABEL: emitc.include "math.h"

// CHECK-LABEL: func.func @absf_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "fabs"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @absf_to_call_opaque(%arg0: f32) {
%1 = math.absf %arg0 : f32
Copy link
Contributor

@simon-camp simon-camp Oct 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add the f64 variants for the tests please

Copy link
Contributor Author

@recursion-man recursion-man Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I decided to support only float32 until I decide how to proceed with this thread. I would appreciate you opinion about it.

Copy link
Contributor Author

@recursion-man recursion-man Oct 31, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a commit that support both f64 and f32, and updated the lit tests

Copy link
Contributor

@simon-camp simon-camp Nov 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you make the tests exhaustive, with something like this:

func.func @floor(%arg0: f32, %arg1: f64) {
    // C: emitc.call_opaque "floorf" (%arg0)
    // C-NEXT: emitc.call_opaque "floor" (%arg1)
    // CPP: emitc.call_opaque "std::floor" (%arg0)
    // CPP-NEXT: emitc.call_opaque "std::floor" (%arg1)
    %0 = math.floor %arg0 : f32
    %1 = math.floor %arg1 : f64
    return
}

Addiotionally a test for math.round is missing currently.

Copy link
Contributor Author

@recursion-man recursion-man Nov 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. I changed the test accordingly and added one for math.round.

return
}

// -----

// CHECK-LABEL: func.func @floor_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "floor"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @floor_to_call_opaque(%arg0: f32) {
%1 = math.floor %arg0 : f32
return
}

// -----

// CHECK-LABEL: func.func @sin_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "sin"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @sin_to_call_opaque(%arg0: f32) {
%1 = math.sin %arg0 : f32
return
}

// -----

// CHECK-LABEL: func.func @cos_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "cos"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @cos_to_call_opaque(%arg0: f32) {
%1 = math.cos %arg0 : f32
return
}


// -----

// CHECK-LABEL: func.func @asin_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "asin"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @asin_to_call_opaque(%arg0: f32) {
%1 = math.asin %arg0 : f32
return
}

// -----

// CHECK-LABEL: func.func @acos_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "acos"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @acos_to_call_opaque(%arg0: f32) {
%1 = math.acos %arg0 : f32
return
}

// -----

// CHECK-LABEL: func.func @atan2_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32,
// CHECK-SAME: %[[VAL_1:.*]]: f32) {
// CHECK: %[[VAL_2:.*]] = emitc.call_opaque "atan2"(%[[VAL_0]], %[[VAL_1]]) : (f32, f32) -> f32
// CHECK: return
// CHECK: }
func.func @atan2_to_call_opaque(%arg0: f32, %arg1: f32) {
%1 = math.atan2 %arg0, %arg1 : f32
return
}

// -----

// CHECK-LABEL: func.func @ceil_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "ceil"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @ceil_to_call_opaque(%arg0: f32) {
%1 = math.ceil %arg0 : f32
return
}

// -----

// CHECK-LABEL: func.func @exp_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32) {
// CHECK: %[[VAL_1:.*]] = emitc.call_opaque "exp"(%[[VAL_0]]) : (f32) -> f32
// CHECK: return
// CHECK: }
func.func @exp_to_call_opaque(%arg0: f32) {
%1 = math.exp %arg0 : f32
return
}


// -----

// CHECK-LABEL: func.func @fpowi_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: f32,
// CHECK-SAME: %[[VAL_1:.*]]: i32) {
// CHECK: %[[VAL_2:.*]] = emitc.call_opaque "powf"(%[[VAL_0]], %[[VAL_1]]) : (f32, i32) -> f32
// CHECK: return
// CHECK: }
func.func @fpowi_to_call_opaque(%arg0: f32, %arg1: i32) {
%1 = math.fpowi %arg0, %arg1 : f32, i32
return
}

// -----

// CHECK-LABEL: func.func @ipowi_to_call_opaque(
// CHECK-SAME: %[[VAL_0:.*]]: i32,
// CHECK-SAME: %[[VAL_1:.*]]: i32) {
// CHECK: %[[VAL_2:.*]] = emitc.call_opaque "pow"(%[[VAL_0]], %[[VAL_1]]) : (i32, i32) -> i32
// CHECK: return
// CHECK: }
func.func @ipowi_to_call_opaque(%arg0: i32, %arg1: i32) {
%1 = math.ipowi %arg0, %arg1 : i32
return
}


22 changes: 22 additions & 0 deletions utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4201,6 +4201,7 @@ cc_library(
":IndexToLLVM",
":IndexToSPIRV",
":LinalgToStandard",
":MathToEmitC",
":MathToFuncs",
":MathToLLVM",
":MathToLibm",
Expand Down Expand Up @@ -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"]),
Expand Down
Loading