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

[Jakub Kuderski]

================
@@ -2610,6 +2610,19 @@ 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 {};
+  int64_t lb = forOp.getConstantLowerBound();
+  int64_t step = forOp.getStepAsInt();
+  int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
----------------
kuhar wrote:

Can the trip count be 0? If not, maybe assert that for completeness?

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