Skip to content

Support lhs constant folding for commutative ops #156

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
tancheng merged 4 commits into from
Oct 17, 2025
Merged

Support lhs constant folding for commutative ops #156

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
56a5b4a
support lhs folding for cummutative ops
ShangkunLi Oct 15, 2025
b484956
[clean] rm some local files
ShangkunLi Oct 15, 2025
e4ad69a
fix bug and add test
ShangkunLi Oct 16, 2025
4352ece
release constant folding constraints
ShangkunLi Oct 16, 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
51 changes: 37 additions & 14 deletions lib/NeuraDialect/Transforms/Optimizations/HwAgnosticOpt/FoldConstantPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,11 @@ template <typename OpType>
struct FuseRhsConstantPattern : public OpRewritePattern<OpType> {
using OpRewritePattern<OpType>::OpRewritePattern;

// By default, we assume the operation is not commutative.
// If the operation is commutative, we can extend this pattern to support
// constant folding on the left-hand side operand as well.
virtual bool isCommutative() const { return false; }

virtual Operation *
createOpWithFusedRhsConstant(OpType op, Attribute rhs_const_value,
PatternRewriter &rewriter) const = 0;
Expand All @@ -110,32 +115,46 @@ struct FuseRhsConstantPattern : public OpRewritePattern<OpType> {
Value lhs = op.getLhs();
Value rhs = op.getRhs();

if (isOriginConstantOp(lhs)) {
llvm::errs() << "LHS constant folding not supported yet.\n";
return failure();
}
bool lhs_is_const = isOriginConstantOp(lhs);
bool rhs_is_const = rhs && isOriginConstantOp(rhs);

if (rhs_is_const && !lhs_is_const) {
auto constant_op = dyn_cast<neura::ConstantOp>(rhs.getDefiningOp());

if (!rhs || !isOriginConstantOp(rhs)) {
return failure();
Attribute rhs_const_value = getOriginConstantValue(rhs);
Operation *fused_op =
createOpWithFusedRhsConstant(op, rhs_const_value, rewriter);

rewriter.replaceOp(op, fused_op->getResults());
if (constant_op->use_empty()) {
rewriter.eraseOp(constant_op);
}
return success();
}

auto constant_op = dyn_cast<neura::ConstantOp>(rhs.getDefiningOp());
if (lhs_is_const && !rhs_is_const && isCommutative()) {
auto constant_op = dyn_cast<neura::ConstantOp>(lhs.getDefiningOp());

Attribute rhs_const_value = getOriginConstantValue(rhs);
Operation *fused_op =
createOpWithFusedRhsConstant(op, rhs_const_value, rewriter);
Attribute lhs_const_value = getOriginConstantValue(lhs);
Operation *fused_op =
createOpWithFusedRhsConstant(op, lhs_const_value, rewriter);

rewriter.replaceOp(op, fused_op->getResults());
if (constant_op->use_empty()) {
rewriter.eraseOp(constant_op);
rewriter.replaceOp(op, fused_op->getResults());
if (constant_op->use_empty()) {
rewriter.eraseOp(constant_op);
}
return success();
}
return success();

return failure();
}
};

struct FuseAddRhsConstantPattern : public FuseRhsConstantPattern<neura::AddOp> {
using FuseRhsConstantPattern<neura::AddOp>::FuseRhsConstantPattern;

bool isCommutative() const override { return true; }

Operation *
createOpWithFusedRhsConstant(neura::AddOp op, Attribute rhs_const_value,
PatternRewriter &rewriter) const override {
Expand Down Expand Up @@ -164,6 +183,8 @@ struct FuseSubRhsConstantPattern : public FuseRhsConstantPattern<neura::SubOp> {
struct FuseMulRhsConstantPattern : public FuseRhsConstantPattern<neura::MulOp> {
using FuseRhsConstantPattern<neura::MulOp>::FuseRhsConstantPattern;

bool isCommutative() const override { return true; }

Operation *
createOpWithFusedRhsConstant(neura::MulOp op, Attribute rhs_const_value,
PatternRewriter &rewriter) const override {
Expand Down Expand Up @@ -194,6 +215,8 @@ struct FuseFAddRhsConstantPattern
: public FuseRhsConstantPattern<neura::FAddOp> {
using FuseRhsConstantPattern<neura::FAddOp>::FuseRhsConstantPattern;

bool isCommutative() const override { return true; }

Operation *
createOpWithFusedRhsConstant(neura::FAddOp op, Attribute rhs_const_value,
PatternRewriter &rewriter) const override {
Expand Down