[Mlir-commits] [mlir] [mlir][affine] implement `promoteIfSingleIteration` for `AffineForOp` (PR #72547)

[Jakub Kuderski]

================
@@ -2440,6 +2442,53 @@ std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
   return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
 }
 
+/// Promotes the loop body of a forOp to its containing block if the forOp
+/// was known to have a single iteration.
+LogicalResult AffineForOp::promoteIfSingleIteration(RewriterBase &rewriter) {
+  auto forOp = cast<AffineForOp>(getOperation());
+  std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
+  if (!tripCount || *tripCount != 1)
+    return failure();
+
+  // TODO: extend this for arbitrary affine bounds.
+  if (forOp.getLowerBoundMap().getNumResults() != 1)
+    return failure();
+
+  // Replaces all IV uses to its single iteration value.
+  auto iv = forOp.getInductionVar();
+  auto *parentBlock = forOp->getBlock();
----------------
kuhar wrote:

Same nit about types

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