[Mlir-commits] [mlir] [mlir] Harden affine loop folding (PR #208736)

[llvmlistbot at llvm.org]

https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/208736

I made affine loop folding fail closed when conversion is looking at temporary IR whose value ownership cannot be proven, and I made the loop interface default treat unknown parentage conservatively. This keeps normal stable-IR folding behavior intact while avoiding unsafe conversion-time folding.
Fixes #207352 

>From 0582d9740cb10dc99afe852978e42d04138cabbe Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 10 Jul 2026 22:33:17 +0000
Subject: [PATCH] [mlir] Harden affine loop folding

---
 mlir/include/mlir/Interfaces/LoopLikeInterface.td |  4 +++-
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp          |  5 ++---
 mlir/test/Transforms/test-convert-func-op.mlir    | 15 +++++++++++++++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 5fb897339ffde..d64287b6c1f0b 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -62,7 +62,9 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*args=*/(ins "::mlir::Value ":$value),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return !$_op->isAncestor(value.getParentRegion()->getParentOp());
+        ::mlir::Region *region = value.getParentRegion();
+        ::mlir::Operation *parentOp = region ? region->getParentOp() : nullptr;
+        return parentOp && !$_op->isAncestor(parentOp);
       }]
     >,
     InterfaceMethod<[{
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f095500495f18..9caa5841b0e4d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2595,9 +2595,8 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
     if (val == forOp.getInductionVar())
       return {};
     if (iterArgIt == iterArgs.end()) {
-      // `val` is defined outside of the loop.
-      assert(forOp.isDefinedOutsideOfLoop(val) &&
-             "must be defined outside of the loop");
+      if (!forOp.isDefinedOutsideOfLoop(val))
+        return {};
       hasValDefinedOutsideLoop = true;
       replacements.push_back(val);
     } else {
diff --git a/mlir/test/Transforms/test-convert-func-op.mlir b/mlir/test/Transforms/test-convert-func-op.mlir
index 180f16a32991b..8587677e7bc30 100644
--- a/mlir/test/Transforms/test-convert-func-op.mlir
+++ b/mlir/test/Transforms/test-convert-func-op.mlir
@@ -38,3 +38,18 @@ func.func @byref(%arg0: !test.smpla {llvm.byref = !test.smpla}) -> !test.smpla {
 // CHECK-SAME: (%[[ARG0:.*]]: !llvm.ptr {llvm.byref = !llvm.struct<(i8, i8)>}) -> !llvm.struct<(i8, i8)>
 //      CHECK: %[[LD:.*]] = llvm.load %[[ARG0]] : !llvm.ptr -> !llvm.struct<(i8, i8)>
 //      CHECK: llvm.return %[[LD]] : !llvm.struct<(i8, i8)>
+
+// -----
+
+// CHECK-LABEL: llvm.func @affine_for_iter_args
+func.func @affine_for_iter_args(%arg0: memref<10xf32>) {
+  %0 = affine.for %arg1 = 0 to 10 iter_args(%arg2 = %arg0) -> (memref<10xf32>) {
+    affine.yield %arg0 : memref<10xf32>
+  }
+  return
+}
+
+// CHECK: builtin.unrealized_conversion_cast
+// CHECK: affine.for
+// CHECK: affine.yield
+// CHECK: llvm.return