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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>

Reverts llvm/llvm-project#<!-- -->164064

---
Full diff: https://github.com/llvm/llvm-project/pull/165607.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+5-24) 
- (modified) mlir/test/Dialect/Affine/canonicalize.mlir (-13) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 002f1f60bb58e..0c3592124cdec 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2610,21 +2610,6 @@ 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.
-static std::optional<int64_t>
-getConstantInductionVarForLastTrip(AffineForOp forOp) {
-  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
-  if (!tripCount.has_value())
-    return std::nullopt;
-  if (tripCount.value() == 0)
-    return std::nullopt;
-  int64_t lb = forOp.getConstantLowerBound();
-  int64_t step = forOp.getStepAsInt();
-  int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
-  return lastTripIv;
-}
-
 /// Fold the empty loop.
 static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   if (!llvm::hasSingleElement(*forOp.getBody()))
@@ -2637,7 +2622,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
     // results.
     return forOp.getInits();
   }
-  SmallVector<OpFoldResult, 4> replacements;
+  SmallVector<Value, 4> replacements;
   auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
   auto iterArgs = forOp.getRegionIterArgs();
   bool hasValDefinedOutsideLoop = false;
@@ -2645,14 +2630,10 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
     Value val = yieldOp.getOperand(i);
     BlockArgument *iterArgIt = llvm::find(iterArgs, val);
-    if (val == forOp.getInductionVar()) {
-      if (auto lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
-        replacements.push_back(IntegerAttr::get(
-            IndexType::get(forOp.getContext()), lastTripIv.value()));
-        continue;
-      }
+    // TODO: It should be possible to perform a replacement by computing the
+    // last value of the IV based on the bounds and the step.
+    if (val == forOp.getInductionVar())
       return {};
-    }
     if (iterArgIt == iterArgs.end()) {
       // `val` is defined outside of the loop.
       assert(forOp.isDefinedOutsideOfLoop(val) &&
@@ -2675,7 +2656,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   // out of order.
   if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
     return {};
-  return replacements;
+  return llvm::to_vector_of<OpFoldResult>(replacements);
 }
 
 /// Canonicalize the bounds of the given loop.
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 997f23b4bd669..1169cd1c29d74 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -609,19 +609,6 @@ func.func @fold_zero_iter_loops(%in : index) -> index {
 
 // -----
 
-// CHECK-LABEL: func @fold_empty_loop_iv
-//  CHECK-SAME:   %[[INIT:.*]]: index
-func.func @fold_empty_loop_iv(%init: index) -> (index, index) {
-  %res:2 = affine.for %i = 0 to 10 step 1 iter_args(%arg0 = %init, %arg1 = %init) -> (index, index) {
-    affine.yield %i, %arg1 : index, index
-  }
-  // CHECK: %[[C9:.*]] = arith.constant 9 : index
-  // CHECK: return %[[C9]], %[[INIT]] : index, index
-  return %res#0, %res#1 : index, index
-}
-
-// -----
-
 // CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
 
 // CHECK-LABEL: func @canonicalize_affine_if

``````````

</details>


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