[MLIR][Affine] Fix crash in affine.for empty loop folder (#130372)

A yield value can be the loop IV itself.

Fixes: #74301

Author: bondhugula

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -2314,6 +2314,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) &&

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
+}