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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: lianjinfeng2003 (mygitljf)

<details>
<summary>Changes</summary>

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 

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


3 Files Affected:

- (modified) mlir/include/mlir/Interfaces/LoopLikeInterface.td (+3-1) 
- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+2-3) 
- (modified) mlir/test/Transforms/test-convert-func-op.mlir (+15) 


``````````diff
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

``````````

</details>


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