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

[Takayuki Todokoro]

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

## 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`.


>From c04f3e6990e446ca867a1594b3e9140a28ad8e51 Mon Sep 17 00:00:00 2001
From: takatodo <takatodo1227 at gmail.com>
Date: Sun, 19 Jul 2026 12:04:15 +0900
Subject: [PATCH] [MLIR][Affine] Preserve discardable attrs when adding loop
 yields

AffineForOp::replaceWithAdditionalYields rebuilds the loop but previously dropped discardable attributes attached by scheduling and transformation pipelines.

Copy the original discardable attribute dictionary to the replacement loop and add an unroll-and-jam regression that exercises the LoopLikeOpInterface path.
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 1 +
 mlir/test/Dialect/Affine/unroll-jam.mlir | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

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)