Skip to content
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

Fixbug #433

Merged
merged 5 commits into from
Mar 5, 2025
Merged

Fixbug #433

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
31ef3d1
Fix torem bug
wsmoses Mar 4, 2025
2b6c3e0
fix
wsmoses Mar 4, 2025
c34ad70
more fixes
wsmoses Mar 5, 2025
8629a92
fmt
wsmoses Mar 5, 2025
bf7ec98
Merge branch 'main' into fixbug
wsmoses Mar 5, 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
Prev Previous commit
Next Next commit
fmt
  • Loading branch information
@wsmoses
wsmoses committed Mar 5, 2025
commit 8629a92536cdeb792a2745b47dc2a135b38505b5
43 changes: 24 additions & 19 deletions src/enzyme_ad/jax/Passes/AffineCFG.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2874,21 +2874,24 @@ static AffineMap optimizeMap(AffineMap map,
map.getContext());
}

struct OptimizeRem
: public OpRewritePattern<arith::RemUIOp> {
struct OptimizeRem : public OpRewritePattern<arith::RemUIOp> {
using OpRewritePattern<arith::RemUIOp>::OpRewritePattern;

LogicalResult matchAndRewrite(arith::RemUIOp op,
PatternRewriter &rewriter) const override {
AddIOp sum = op.getLhs().getDefiningOp<arith::AddIOp>();
if (!sum) return failure();
for (int i=0; i<2; i++) {
auto val = sum->getOperand(i).getDefiningOp<arith::MulIOp>();
if (!val) continue;
if (val.getRhs() != op.getRhs()) continue;
rewriter.replaceOpWithNewOp<arith::RemUIOp>(op, sum->getOperand(1-i), op.getRhs());
}
return failure();
AddIOp sum = op.getLhs().getDefiningOp<arith::AddIOp>();
if (!sum)
return failure();
for (int i = 0; i < 2; i++) {
auto val = sum->getOperand(i).getDefiningOp<arith::MulIOp>();
if (!val)
continue;
if (val.getRhs() != op.getRhs())
continue;
rewriter.replaceOpWithNewOp<arith::RemUIOp>(op, sum->getOperand(1 - i),
op.getRhs());
}
return failure();
}
};

Expand Down Expand Up @@ -3046,10 +3049,10 @@ struct SplitParallelInductions
break;
}
} else {
if (pval == iv)
continue;
legal = false;
break;
if (pval == iv)
continue;
legal = false;
break;
}
auto findBasePattern = [](Value iv, AffineExpr root,
ValueRange operands, ValueOrInt &base,
Expand Down Expand Up @@ -3345,8 +3348,11 @@ struct SplitParallelInductions
rewriter.create<arith::ConstantIndexOp>(U->getLoc(),
base.i_val));
replacement.setOverflowFlags(IntegerOverflowFlags::nuw);
auto replacement2 = rewriter.create<arith::AddIOp>(U->getLoc(), replacement, newIv);
rewriter.replaceUsesWithIf(iv, replacement2->getResult(0), [&](OpOperand &op) { return op.getOwner() == U; });
auto replacement2 =
rewriter.create<arith::AddIOp>(U->getLoc(), replacement, newIv);
rewriter.replaceUsesWithIf(
iv, replacement2->getResult(0),
[&](OpOperand &op) { return op.getOwner() == U; });
}
}

Expand Down Expand Up @@ -3752,8 +3758,7 @@ void AffineCFGPass::runOnOperation() {
MoveStoreToAffine, MoveIfToAffine, MoveLoadToAffine,
AffineIfSimplification, CombineAffineIfs,
MergeNestedAffineParallelLoops, PrepMergeNestedAffineParallelLoops,
MergeNestedAffineParallelIf, MergeParallelInductions,
OptimizeRem,
MergeNestedAffineParallelIf, MergeParallelInductions, OptimizeRem,
CanonicalieForBounds, AddAddCstEnd>(getOperation()->getContext(), 2);
rpl.add<SplitParallelInductions>(getOperation()->getContext(), 1);
GreedyRewriteConfig config;
Expand Down
16 changes: 8 additions & 8 deletions src/enzyme_ad/jax/Passes/CanonicalizeLoops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,7 @@ std::optional<int64_t> maxSize(mlir::Value v) {
return lhs;

if (constValue.isNegative())
return {};
return {};

return (*lhs) * constValue.getZExtValue();
}
Expand Down Expand Up @@ -557,9 +557,9 @@ class DivMul final : public OpRewritePattern<arith::DivUIOp> {
if (!maxSizeOpt)
return failure();
if (!operand.getType().isIndex())
if (APInt::getMaxValue(operand.getType().getIntOrFloatBitWidth())
.ult(*maxSizeOpt))
return failure();
if (APInt::getMaxValue(operand.getType().getIntOrFloatBitWidth())
.ult(*maxSizeOpt))
return failure();
if (operand.getRhs() != ext.getRhs())
return failure();
rewriter.replaceOp(ext, operand.getLhs());
Expand Down Expand Up @@ -739,7 +739,7 @@ class ToRem final : public OpRewritePattern<arith::AddIOp> {
if (factor != -divisor)
continue;
if (div.getLhs() != val)
continue;
continue;
rewriter.replaceOpWithNewOp<arith::RemUIOp>(
ext, val, factor.isNegative() ? div.getRhs() : mul.getRhs());
return success();
Expand Down Expand Up @@ -1183,7 +1183,7 @@ struct CanonicalizeLoopsPass

void mlir::enzyme::addSingleIter(RewritePatternSet &patterns,
MLIRContext *ctx) {
patterns.add<RemoveAffineParallelSingleIter,ExtUIOfIndexUI, TruncIOfIndexUI, ShrUIOfIndexUI,
DivUIOfIndexUI, DivMul, AddIOfIndexUI, MulIOfIndexUI, ShLIOfIndexUI,
AddIOfDoubleIndex, ToRem>(ctx);
patterns.add<RemoveAffineParallelSingleIter, ExtUIOfIndexUI, TruncIOfIndexUI,
ShrUIOfIndexUI, DivUIOfIndexUI, DivMul, AddIOfIndexUI,
MulIOfIndexUI, ShLIOfIndexUI, AddIOfDoubleIndex, ToRem>(ctx);
}