diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2482,9 +2482,12 @@ SmallVectorImpl &results) { bool folded = succeeded(foldLoopBounds(*this)); folded |= succeeded(canonicalizeLoopBounds(*this)); - if (hasTrivialZeroTripCount(*this)) { + if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) { // The initial values of the loop-carried variables (iter_args) are the - // results of the op. + // results of the op. But this must be avoided for an affine.for op that + // does not return any results. Since ops that do not return results cannot + // be folded away, we would enter an infinite loop of folds on the same + // affine.for op. results.assign(getIterOperands().begin(), getIterOperands().end()); folded = true; } diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -582,14 +582,23 @@ // ----- +// CHECK-LABEL: func @zero_iter_loop_not_folded +func.func @zero_iter_loop_not_folded() { + %A = memref.alloc() : memref<4xf32> + affine.for %i = 0 to 0 { + %load = affine.load %A[%i] : memref<4xf32> + affine.store %load, %A[%i] : memref<4xf32> + } + // CHECK: affine.for {{.*}} = 0 to 0 { + return +} + +// ----- + // CHECK-LABEL: func @fold_zero_iter_loops // CHECK-SAME: %[[ARG:.*]]: index func.func @fold_zero_iter_loops(%in : index) -> index { %c1 = arith.constant 1 : index - affine.for %i = 0 to 0 { - affine.for %j = 0 to -1 { - } - } %res = affine.for %i = 0 to 0 iter_args(%loop_arg = %in) -> index { %yield = arith.addi %loop_arg, %c1 : index affine.yield %yield : index