[Mlir-commits] [mlir] [mlir][Linalg] Preserve discardable/user-defined attributes during generalization (PR #178599)

[Abhishek Varma]

https://github.com/Abhishek-Varma updated https://github.com/llvm/llvm-project/pull/178599

>From 5da56a3ab212de9b0da9e262bad603d34f7b49e0 Mon Sep 17 00:00:00 2001
From: Abhishek Varma <abhvarma at amd.com>
Date: Thu, 29 Jan 2026 05:50:15 +0000
Subject: [PATCH 1/2] [Linalg] Preserve discardable attributes during
 generalization

-- As observed in a [downstream project](https://github.com/iree-org/iree/pull/23294#discussion_r2734982998) : the named to generize
   linalg op conversion wasn't preserving discardable attributes.
-- This commit aims to fix the same.
-- Only a single test case is added as the change applies to any
   named linalg op's generalization.

Signed-off-by: Abhishek Varma <abhvarma at amd.com>
---
 .../Linalg/Transforms/Generalization.cpp       |  4 ++++
 .../Dialect/Linalg/generalize-named-ops.mlir   | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index 75bb1757a55f5..67a116827362f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -64,6 +64,10 @@ FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
                         outputs, indexingMaps, iterators);
   rewriter.inlineRegionBefore(linalgOp->getRegion(0), genericOp.getRegion(),
                               genericOp.getRegion().begin());
+
+  // Preserve discardable (user-defined) attributes from the original op.
+  genericOp->setDiscardableAttrs(linalgOp->getDiscardableAttrDictionary());
+
   rewriter.replaceOp(linalgOp, genericOp->getResults());
   return genericOp;
 }
diff --git a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
index dcdd6c8db4b21..e346bee901f1d 100644
--- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
@@ -1261,3 +1261,21 @@ func.func @contract_matmul_bcast_a_b(
       outs(%C: memref<3x7xf32>)
   return
 }
+
+// -----
+
+// Test that discardable (user-defined) attributes are preserved during
+// generalization.
+
+func.func @preserve_discardable_attrs(%A : tensor<16x8xf32>,
+                                            %B : tensor<8x32xf32>,
+                                            %C : tensor<16x32xf32>) -> tensor<16x32xf32> {
+  %0 = linalg.matmul {my_custom_attr = "preserved", another_attr = 42 : i64}
+                     ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
+                     outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+  return %0: tensor<16x32xf32>
+}
+
+// CHECK-LABEL: func @preserve_discardable_attrs
+// CHECK:         linalg.generic
+// CHECK-SAME:        attrs = {another_attr = 42 : i64, my_custom_attr = "preserved"}

>From 2d07703c8e9e73d8382c0017cfaf27c9e8f7f22f Mon Sep 17 00:00:00 2001
From: Abhishek Varma <abhvarma at amd.com>
Date: Thu, 29 Jan 2026 13:03:39 +0000
Subject: [PATCH 2/2] Nit comment

---
 mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index 67a116827362f..35abdc027f194 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -65,7 +65,11 @@ FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
   rewriter.inlineRegionBefore(linalgOp->getRegion(0), genericOp.getRegion(),
                               genericOp.getRegion().begin());
 
-  // Preserve discardable (user-defined) attributes from the original op.
+  // Discardable attributes carry user-defined metadata (e.g., annotations for
+  // downstream passes). Generalization is a semantics-preserving
+  // transformation, so dropping this metadata would be unexpected. This is safe
+  // because discardable attributes are by definition independent of op
+  // semantics.
   genericOp->setDiscardableAttrs(linalgOp->getDiscardableAttrDictionary());
 
   rewriter.replaceOp(linalgOp, genericOp->getResults());