[Mlir-commits] [mlir] [MLIR][Affine] Preserve discardable attrs when adding loop yields (PR #210571)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Takayuki Todokoro (takatodo)

<details>
<summary>Changes</summary>

## Summary

- Preserve discardable attributes when
  `AffineForOp::replaceWithAdditionalYields` rebuilds an `affine.for`.
- Extend the existing unroll-and-jam regression to verify that an inner-loop
  attribute survives the `LoopLikeOpInterface` yield-extension path.

## Motivation

`replaceWithAdditionalYields` creates a replacement loop with additional
iteration arguments, but did not copy the original loop's discardable
attribute dictionary. Consequently, scheduling or transformation annotations
attached to the original loop were silently lost when a transform added loop
yields.

The replacement loop now receives the original discardable attributes. This
does not add an API, pass, or option; it only preserves existing IR metadata
across the replacement.

## Testing

- Built `mlir-opt` from an assertion-enabled Release configuration based on
  `origin/main`.
- Ran all three `RUN` variants in
  `mlir/test/Dialect/Affine/unroll-jam.mlir`.

Assisted-by: OpenAI Codex gpt-5.6-sol max

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+1) 
- (modified) mlir/test/Dialect/Affine/unroll-jam.mlir (+2-2) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f095500495f18..d4c4dd6b48424 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2893,6 +2893,7 @@ FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
   AffineForOp newLoop = AffineForOp::create(
       rewriter, getLoc(), getLowerBoundOperands(), getLowerBoundMap(),
       getUpperBoundOperands(), getUpperBoundMap(), getStepAsInt(), inits);
+  newLoop->setDiscardableAttrs(getOperation()->getDiscardableAttrDictionary());
 
   // Generate the new yield values and append them to the scf.yield operation.
   auto yieldOp = cast<AffineYieldOp>(getBody()->getTerminator());
diff --git a/mlir/test/Dialect/Affine/unroll-jam.mlir b/mlir/test/Dialect/Affine/unroll-jam.mlir
index 8ed7fccf7d251..bba24844c3484 100644
--- a/mlir/test/Dialect/Affine/unroll-jam.mlir
+++ b/mlir/test/Dialect/Affine/unroll-jam.mlir
@@ -185,7 +185,7 @@ func.func @unroll_jam_one_iter_arg() {
     %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) {
       %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32
       affine.yield %y : i32
-    }
+    } {test.keep = "inner"}
     %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32
   }
   return
@@ -201,7 +201,7 @@ func.func @unroll_jam_one_iter_arg() {
 // CHECK-NEXT:     [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])
 // CHECK-NEXT:     [[RES5:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]])
 // CHECK-NEXT:     affine.yield [[RES4]], [[RES5]]
-// CHECK-NEXT:   }
+// CHECK-NEXT:   } {test.keep = "inner"}
 // CHECK:        "foo"([[IV0]], [[RES1]], [[RES3]]#0)
 // CHECK-NEXT:   affine.apply [[$MAP_PLUS_1]]([[IV0]])
 // CHECK-NEXT:   "foo"({{.*}}, [[RES2]], [[RES3]]#1)

``````````

</details>


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