[Mlir-commits] [mlir] [mlir][affine] Make AffineForEmptyLoopFolder as folder function (PR #163929)

[lonely eagle]

================
@@ -2460,6 +2460,67 @@ static LogicalResult foldLoopBounds(AffineForOp forOp) {
   return success(folded);
 }
 
+/// Returns constant trip count in trivial cases.
+static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
+  int64_t step = forOp.getStepAsInt();
+  if (!forOp.hasConstantBounds() || step <= 0)
+    return std::nullopt;
+  int64_t lb = forOp.getConstantLowerBound();
+  int64_t ub = forOp.getConstantUpperBound();
+  return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
+}
+
+/// Fold the empty loop.
+static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) {
+  if (!llvm::hasSingleElement(*forOp.getBody()))
+    return failure();
+  if (forOp.getNumResults() == 0)
+    return success();
+  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
+  if (tripCount == 0) {
+    // The initial values of the iteration arguments would be the op's
+    // results.
+    forOp.getResults().replaceAllUsesWith(forOp.getInits());
----------------
linuxlonelyeagle wrote:

Thank you for reminding me; it's been so long since I last wrote the folder function that I'd rather forgotten how.Today I discovered a pattern for eliminating for loops in scf.for. Subsequently, I identified a TODO in affine.for and submitted the corresponding pull request.

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