[Mlir-commits] [mlir] [mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder (PR #164064)

[Matthias Springer]

================
@@ -2610,6 +2610,21 @@ static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
   return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
 }
 
+/// Calculate the constant value of the loop's induction variable for its last
+/// trip, construct an OpFoldResult using this value and return it.
+static OpFoldResult getConstantInductionVarForLastTrip(AffineForOp forOp) {
+  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
+  if (!tripCount.has_value())
+    return {};
+  if (tripCount.value() == 0)
+    return {};
+  int64_t lb = forOp.getConstantLowerBound();
+  int64_t step = forOp.getStepAsInt();
+  int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
+  return OpFoldResult(
+      IntegerAttr::get(IndexType::get(forOp.getContext()), lastTripIv));
+}
+
 /// Fold the empty loop.
 static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
----------------
matthias-springer wrote:

Is this folder actually correct? What if there's a side-effecting op in the body?

https://github.com/llvm/llvm-project/pull/164064