[MLIR][Affine] Fix crash in affine.for empty loop folder

[bondhugula]

A yield value can be the loop IV itself.

Fixes: #74301

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Uday Bondhugula (bondhugula)

Changes

An yield value can be the loop IV itself.

Fixes: #74301

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+4)
• (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+9)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 06493de2b09d4..3d71c49ccfb17 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2304,6 +2304,10 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
     for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
       Value val = yieldOp.getOperand(i);
       auto *iterArgIt = llvm::find(iterArgs, val);
+      // 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 failure();
       if (iterArgIt == iterArgs.end()) {
         // `val` is defined outside of the loop.
         assert(forOp.isDefinedOutsideOfLoop(val) &&
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index a9ac13ad71624..d39c0c6e41df2 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2262,3 +2262,12 @@ func.func @cst_value_to_cst_attr_basis_linearize_index(%arg0 : index, %arg1 : in
   %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by  (%c2, 3, %c4) : index
   return %0 : index
 }
+
+// CHECK-LABEL: func @for_empty_body_folder_iv_yield
+func.func @for_empty_body_folder_iv_yield() -> index {
+  %c18 = arith.constant 18 : index
+  %10 = affine.for %arg3 = 0 to 114 iter_args(%arg4 = %c18) -> (index) {
+    affine.yield %arg3 : index
+  }
+  return %10 : index
+}

[ayzhuang]

LGTM, thanks!