[Mlir-commits] [mlir] [MLIR][Affine] Fix crash in affine.for empty loop folder (PR #130372)

[Uday Bondhugula]

https://github.com/bondhugula updated https://github.com/llvm/llvm-project/pull/130372

>From c216b12e36ff1ee343cf5b885655ac4db86046b1 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Sat, 8 Mar 2025 05:51:56 +0530
Subject: [PATCH] [MLIR][Affine] Fix crash in affine.for empty loop folder

A yield value can be the loop IV itself.

Fixes: https://github.com/llvm/llvm-project/issues/74301
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   | 4 ++++
 mlir/test/Dialect/Affine/canonicalize.mlir | 9 +++++++++
 2 files changed, 13 insertions(+)

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