diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -178,18 +178,19 @@ return } ``` - `affine.for` can also operate on loop-carried variables and return the final - values after loop termination. The initial values of the variables are - passed as additional SSA operands to the "affine.for" following the 2 loop - control values lower bound, upper bound. The operation region has equivalent - arguments for each variable representing the value of the variable at the - current iteration. + `affine.for` can also operate on loop-carried variables (`iter_args`) and + return the final values after loop termination. The initial values of the + variables are passed as additional SSA operands to the `affine.for` + following the operands for the loop's lower and upper bounds. The + operation's region has equivalent arguments for each variable representing + the value of the variable at the current iteration. The region must terminate with an `affine.yield` that passes all the current - iteration variables to the next iteration, or to the `affine.for` result, if - at the last iteration. + iteration variables to the next iteration, or to the `affine.for`'s results + if at the last iteration. For `affine.for`'s that execute zero iterations, the + initial values of the loop-carried variables (corresponding to the SSA + operands) will be the op's results. - `affine.for` results hold the final values after the last iteration. For example, to sum-reduce a memref: ```mlir @@ -208,6 +209,15 @@ return %sum : f32 } ``` + + ```mlir + %res:2 = affine.for %i = 0 to 128 iter_args(%arg0 = %init0, %arg1 = %init1) + -> (index, index) { + %y0 = addi %arg0, %c1 : index + %y1 = addi %arg1, %c2 : index + affine.yield %y0, %y1 : index, index + } + ``` If the `affine.for` defines any values, a yield terminator must be explicitly present. The number and types of the "affine.for" results must match the initial values in the `iter_args` binding and the yield operands. 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 @@ -1648,10 +1648,25 @@ results.add(context); } +/// Returns true if the affine.for has zero iterations in trivial cases. +static bool hasTrivialZeroTripCount(AffineForOp op) { + if (!op.hasConstantBounds()) + return false; + int64_t lb = op.getConstantLowerBound(); + int64_t ub = op.getConstantUpperBound(); + return ub - lb <= 0; +} + LogicalResult AffineForOp::fold(ArrayRef operands, SmallVectorImpl &results) { bool folded = succeeded(foldLoopBounds(*this)); folded |= succeeded(canonicalizeLoopBounds(*this)); + if (hasTrivialZeroTripCount(*this)) { + // The initial values of the loop-carried variables (iter_args) are the + // results of the op. + results.assign(getIterOperands().begin(), getIterOperands().end()); + folded = true; + } return success(folded); } 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 @@ -469,6 +469,23 @@ } // CHECK: return +// ----- + +// CHECK-LABEL: func @fold_zero_iter_loops +// CHECK-SAME: %[[ARG:.*]]: index +func @fold_zero_iter_loops(%in : index) -> index { + %c1 = 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 = addi %loop_arg, %c1 : index + affine.yield %yield : index + } + // CHECK-NEXT: return %[[ARG]] + return %res : index +} // -----