Skip to content

[MLIR][Math] Add optional benefit arg to populate math lowering patterns #127291

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
merged 7 commits into from
Feb 15, 2025
Merged

[MLIR][Math] Add optional benefit arg to populate math lowering patterns #127291

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
282d90a
[MLIR][Math] Add optional benefit arg to populate math lowering patterns
wsmoses Feb 15, 2025
2902627
Update mlir/lib/Conversion/MathToLibm/MathToLibm.cpp
wsmoses Feb 15, 2025
79e569a
Update MathToLLVM.h
wsmoses Feb 15, 2025
c5cef7d
Update MathToLibm.cpp
wsmoses Feb 15, 2025
0aa729a
fixup arg order
wsmoses Feb 15, 2025
304c33d
fmt
wsmoses Feb 15, 2025
83ac782
Update MathToLibm.cpp
wsmoses Feb 15, 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
4 changes: 3 additions & 1 deletion mlir/include/mlir/Conversion/MathToLLVM/MathToLLVM.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#ifndef MLIR_CONVERSION_MATHTOLLVM_MATHTOLLVM_H
#define MLIR_CONVERSION_MATHTOLLVM_MATHTOLLVM_H

#include "mlir/IR/PatternMatch.h"
#include <memory>

namespace mlir {
Expand All @@ -23,7 +24,8 @@ class Pass;

void populateMathToLLVMConversionPatterns(const LLVMTypeConverter &converter,
RewritePatternSet &patterns,
bool approximateLog1p = true);
bool approximateLog1p = true,
PatternBenefit benefit = 1);

void registerConvertMathToLLVMInterface(DialectRegistry &registry);

Expand Down
3 changes: 2 additions & 1 deletion mlir/include/mlir/Conversion/MathToLibm/MathToLibm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ class OperationPass;

/// Populate the given list with patterns that convert from Math to Libm calls.
/// If log1pBenefit is present, use it instead of benefit for the Log1p op.
void populateMathToLibmConversionPatterns(RewritePatternSet &patterns);
void populateMathToLibmConversionPatterns(RewritePatternSet &patterns,
PatternBenefit benefit = 1);

/// Create a pass to convert Math operations to libm calls.
std::unique_ptr<OperationPass<ModuleOp>> createConvertMathToLibmPass();
Expand Down
6 changes: 3 additions & 3 deletions mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,9 +304,9 @@ struct ConvertMathToLLVMPass

void mlir::populateMathToLLVMConversionPatterns(
const LLVMTypeConverter &converter, RewritePatternSet &patterns,
bool approximateLog1p) {
bool approximateLog1p, PatternBenefit benefit) {
if (approximateLog1p)
patterns.add<Log1pOpLowering>(converter);
patterns.add<Log1pOpLowering>(converter, benefit);
// clang-format off
patterns.add<
AbsFOpLowering,
Expand Down Expand Up @@ -337,7 +337,7 @@ void mlir::populateMathToLLVMConversionPatterns(
FTruncOpLowering,
TanOpLowering,
TanhOpLowering
>(converter);
>(converter, benefit);
// clang-format on
}

Expand Down
95 changes: 54 additions & 41 deletions mlir/lib/Conversion/MathToLibm/MathToLibm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,10 @@ template <typename Op>
struct ScalarOpToLibmCall : public OpRewritePattern<Op> {
public:
using OpRewritePattern<Op>::OpRewritePattern;
ScalarOpToLibmCall(MLIRContext *context, StringRef floatFunc,
StringRef doubleFunc)
: OpRewritePattern<Op>(context), floatFunc(floatFunc),
doubleFunc(doubleFunc){};
ScalarOpToLibmCall(MLIRContext *context, PatternBenefit benefit,
StringRef floatFunc, StringRef doubleFunc)
: OpRewritePattern<Op>(context, benefit), floatFunc(floatFunc),
doubleFunc(doubleFunc) {};

LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;

Expand All @@ -62,10 +62,11 @@ struct ScalarOpToLibmCall : public OpRewritePattern<Op> {
};

template <typename OpTy>
void populatePatternsForOp(RewritePatternSet &patterns, MLIRContext *ctx,
StringRef floatFunc, StringRef doubleFunc) {
patterns.add<VecOpToScalarOp<OpTy>, PromoteOpToF32<OpTy>>(ctx);
patterns.add<ScalarOpToLibmCall<OpTy>>(ctx, floatFunc, doubleFunc);
void populatePatternsForOp(RewritePatternSet &patterns, PatternBenefit benefit,
MLIRContext *ctx, StringRef floatFunc,
StringRef doubleFunc) {
patterns.add<VecOpToScalarOp<OpTy>, PromoteOpToF32<OpTy>>(ctx, benefit);
patterns.add<ScalarOpToLibmCall<OpTy>>(ctx, benefit, floatFunc, doubleFunc);
}

} // namespace
Expand Down Expand Up @@ -159,42 +160,54 @@ ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,
return success();
}

void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns) {
void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns,
PatternBenefit benefit) {
MLIRContext *ctx = patterns.getContext();

populatePatternsForOp<math::AbsFOp>(patterns, ctx, "fabsf", "fabs");
populatePatternsForOp<math::AcosOp>(patterns, ctx, "acosf", "acos");
populatePatternsForOp<math::AcoshOp>(patterns, ctx, "acoshf", "acosh");
populatePatternsForOp<math::AsinOp>(patterns, ctx, "asinf", "asin");
populatePatternsForOp<math::AsinhOp>(patterns, ctx, "asinhf", "asinh");
populatePatternsForOp<math::Atan2Op>(patterns, ctx, "atan2f", "atan2");
populatePatternsForOp<math::AtanOp>(patterns, ctx, "atanf", "atan");
populatePatternsForOp<math::AtanhOp>(patterns, ctx, "atanhf", "atanh");
populatePatternsForOp<math::CbrtOp>(patterns, ctx, "cbrtf", "cbrt");
populatePatternsForOp<math::CeilOp>(patterns, ctx, "ceilf", "ceil");
populatePatternsForOp<math::CosOp>(patterns, ctx, "cosf", "cos");
populatePatternsForOp<math::CoshOp>(patterns, ctx, "coshf", "cosh");
populatePatternsForOp<math::ErfOp>(patterns, ctx, "erff", "erf");
populatePatternsForOp<math::ExpOp>(patterns, ctx, "expf", "exp");
populatePatternsForOp<math::Exp2Op>(patterns, ctx, "exp2f", "exp2");
populatePatternsForOp<math::ExpM1Op>(patterns, ctx, "expm1f", "expm1");
populatePatternsForOp<math::FloorOp>(patterns, ctx, "floorf", "floor");
populatePatternsForOp<math::FmaOp>(patterns, ctx, "fmaf", "fma");
populatePatternsForOp<math::LogOp>(patterns, ctx, "logf", "log");
populatePatternsForOp<math::Log2Op>(patterns, ctx, "log2f", "log2");
populatePatternsForOp<math::Log10Op>(patterns, ctx, "log10f", "log10");
populatePatternsForOp<math::Log1pOp>(patterns, ctx, "log1pf", "log1p");
populatePatternsForOp<math::PowFOp>(patterns, ctx, "powf", "pow");
populatePatternsForOp<math::RoundEvenOp>(patterns, ctx, "roundevenf",
populatePatternsForOp<math::AbsFOp>(patterns, benefit, ctx, "fabsf", "fabs");
populatePatternsForOp<math::AcosOp>(patterns, benefit, ctx, "acosf", "acos");
populatePatternsForOp<math::AcoshOp>(patterns, benefit, ctx, "acoshf",
"acosh");
populatePatternsForOp<math::AsinOp>(patterns, benefit, ctx, "asinf", "asin");
populatePatternsForOp<math::AsinhOp>(patterns, benefit, ctx, "asinhf",
"asinh");
populatePatternsForOp<math::Atan2Op>(patterns, benefit, ctx, "atan2f",
"atan2");
populatePatternsForOp<math::AtanOp>(patterns, benefit, ctx, "atanf", "atan");
populatePatternsForOp<math::AtanhOp>(patterns, benefit, ctx, "atanhf",
"atanh");
populatePatternsForOp<math::CbrtOp>(patterns, benefit, ctx, "cbrtf", "cbrt");
populatePatternsForOp<math::CeilOp>(patterns, benefit, ctx, "ceilf", "ceil");
populatePatternsForOp<math::CosOp>(patterns, benefit, ctx, "cosf", "cos");
populatePatternsForOp<math::CoshOp>(patterns, benefit, ctx, "coshf", "cosh");
populatePatternsForOp<math::ErfOp>(patterns, benefit, ctx, "erff", "erf");
populatePatternsForOp<math::ExpOp>(patterns, benefit, ctx, "expf", "exp");
populatePatternsForOp<math::Exp2Op>(patterns, benefit, ctx, "exp2f", "exp2");
populatePatternsForOp<math::ExpM1Op>(patterns, benefit, ctx, "expm1f",
"expm1");
populatePatternsForOp<math::FloorOp>(patterns, benefit, ctx, "floorf",
"floor");
populatePatternsForOp<math::FmaOp>(patterns, benefit, ctx, "fmaf", "fma");
populatePatternsForOp<math::LogOp>(patterns, benefit, ctx, "logf", "log");
populatePatternsForOp<math::Log2Op>(patterns, benefit, ctx, "log2f", "log2");
populatePatternsForOp<math::Log10Op>(patterns, benefit, ctx, "log10f",
"log10");
populatePatternsForOp<math::Log1pOp>(patterns, benefit, ctx, "log1pf",
"log1p");
populatePatternsForOp<math::PowFOp>(patterns, benefit, ctx, "powf", "pow");
populatePatternsForOp<math::RoundEvenOp>(patterns, benefit, ctx, "roundevenf",
"roundeven");
populatePatternsForOp<math::RoundOp>(patterns, ctx, "roundf", "round");
populatePatternsForOp<math::SinOp>(patterns, ctx, "sinf", "sin");
populatePatternsForOp<math::SinhOp>(patterns, ctx, "sinhf", "sinh");
populatePatternsForOp<math::SqrtOp>(patterns, ctx, "sqrtf", "sqrt");
populatePatternsForOp<math::RsqrtOp>(patterns, ctx, "rsqrtf", "rsqrt");
populatePatternsForOp<math::TanOp>(patterns, ctx, "tanf", "tan");
populatePatternsForOp<math::TanhOp>(patterns, ctx, "tanhf", "tanh");
populatePatternsForOp<math::TruncOp>(patterns, ctx, "truncf", "trunc");
populatePatternsForOp<math::RoundOp>(patterns, benefit, ctx, "roundf",
"round");
populatePatternsForOp<math::SinOp>(patterns, benefit, ctx, "sinf", "sin");
populatePatternsForOp<math::SinhOp>(patterns, benefit, ctx, "sinhf", "sinh");
populatePatternsForOp<math::SqrtOp>(patterns, benefit, ctx, "sqrtf", "sqrt");
populatePatternsForOp<math::RsqrtOp>(patterns, benefit, ctx, "rsqrtf",
"rsqrt");
populatePatternsForOp<math::TanOp>(patterns, benefit, ctx, "tanf", "tan");
populatePatternsForOp<math::TanhOp>(patterns, benefit, ctx, "tanhf", "tanh");
populatePatternsForOp<math::TruncOp>(patterns, benefit, ctx, "truncf",
"trunc");
}

namespace {
Expand Down