Skip to content

Commit 27597c6

Browse files
committed
[MLIR][Affine] Fix crash in affine.for empty loop folder
An yield value can be the loop IV itself. Fixes: #74301
1 parent ec54ec6 commit 27597c6

File tree

2 files changed

+13
-0
lines changed

2 files changed

+13
-0
lines changed

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2304,6 +2304,10 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
23042304
for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
23052305
Value val = yieldOp.getOperand(i);
23062306
auto *iterArgIt = llvm::find(iterArgs, val);
2307+
// TODO: It should be possible to perform a replacement by computing the
2308+
// last value of the IV based on the bounds and the step.
2309+
if (val == forOp.getInductionVar())
2310+
return failure();
23072311
if (iterArgIt == iterArgs.end()) {
23082312
// `val` is defined outside of the loop.
23092313
assert(forOp.isDefinedOutsideOfLoop(val) &&

mlir/test/Dialect/Affine/canonicalize.mlir

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2262,3 +2262,12 @@ func.func @cst_value_to_cst_attr_basis_linearize_index(%arg0 : index, %arg1 : in
22622262
%0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (%c2, 3, %c4) : index
22632263
return %0 : index
22642264
}
2265+
2266+
// CHECK-LABEL: func @for_empty_body_folder_iv_yield
2267+
func.func @for_empty_body_folder_iv_yield() -> index {
2268+
%c18 = arith.constant 18 : index
2269+
%10 = affine.for %arg3 = 0 to 114 iter_args(%arg4 = %c18) -> (index) {
2270+
affine.yield %arg3 : index
2271+
}
2272+
return %10 : index
2273+
}

0 commit comments

Comments
 (0)