[Mlir-commits] [mlir] [mlir] [transform dialect] NamedSequenceOp build: honor arg_attrs when building (PR #168101)

[llvmlistbot at llvm.org]

https://github.com/colawithsauce updated https://github.com/llvm/llvm-project/pull/168101

>From 5cebc4877dcb757514e82f85460cc09cd8e6666a Mon Sep 17 00:00:00 2001
From: colawithsauce <51731521+colawithsauce at users.noreply.github.com>
Date: Sat, 15 Nov 2025 02:24:22 +0800
Subject: [PATCH 1/2] [mlir] NamedSequence builder honor arg_attrs when
 building

Previously the builder did not attach argument attributes when calling NamedSequenceOp::build with `arg_attrs`, which makes we can not create a legal named_sequence operation by this builder.

This patch ensures `arg_attrs` are stored in the op state
so verification and printing show them.
---
 mlir/lib/Dialect/Transform/IR/TransformOps.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 062606e7e10b6..415e2af491c07 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -2571,6 +2571,11 @@ void transform::NamedSequenceOp::build(OpBuilder &builder,
                      TypeAttr::get(FunctionType::get(builder.getContext(),
                                                      rootType, resultTypes)));
   state.attributes.append(attrs.begin(), attrs.end());
+  if (!argAttrs.empty()) {
+    SmallVector<Attribute> argAttrsVec(argAttrs.begin(), argAttrs.end());
+    state.getOrAddProperties<Properties>().arg_attrs =
+        ArrayAttr::get(builder.getContext(), argAttrsVec);
+  }
   state.addRegion();
 
   buildSequenceBody(builder, state, rootType,

>From a44361c3e4af1e3161b6bad3cbf609947c7ed3df Mon Sep 17 00:00:00 2001
From: colawithsauce <51731521+colawithsauce at users.noreply.github.com>
Date: Tue, 18 Nov 2025 03:27:33 +0000
Subject: [PATCH 2/2] [mlir][transform-dialect] add unittest of named_sequence
 build.

---
 .../Dialect/Transform/CMakeLists.txt          |  1 +
 .../TransformNamedSequenceCreate.cpp          | 45 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 mlir/unittests/Dialect/Transform/TransformNamedSequenceCreate.cpp

diff --git a/mlir/unittests/Dialect/Transform/CMakeLists.txt b/mlir/unittests/Dialect/Transform/CMakeLists.txt
index 20cdc63966ec0..6bffd71d82fc2 100644
--- a/mlir/unittests/Dialect/Transform/CMakeLists.txt
+++ b/mlir/unittests/Dialect/Transform/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_mlir_unittest(MLIRTransformDialectTests
+  TransformNamedSequenceCreate.cpp
   BuildOnlyExtensionTest.cpp
   Preload.cpp
 )
diff --git a/mlir/unittests/Dialect/Transform/TransformNamedSequenceCreate.cpp b/mlir/unittests/Dialect/Transform/TransformNamedSequenceCreate.cpp
new file mode 100644
index 0000000000000..849b2b56021a8
--- /dev/null
+++ b/mlir/unittests/Dialect/Transform/TransformNamedSequenceCreate.cpp
@@ -0,0 +1,45 @@
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::transform;
+
+TEST(NamedSequenceOpTest, ArgAttrsAreHonoredByBuilder) {
+  MLIRContext ctx;
+  ctx.loadDialect<TransformDialect>();
+
+  OpBuilder builder(&ctx);
+  auto module = ModuleOp::create(UnknownLoc::get(&ctx));
+  builder.setInsertionPointToEnd(module.getBody());
+
+  Location loc = UnknownLoc::get(&ctx);
+
+  static constexpr StringLiteral kMainSequenceName = "__transform_main";
+
+  NamedSequenceOp seqOp = builder.create<NamedSequenceOp>(
+      loc,
+      /*sym_name=*/kMainSequenceName,
+      /*rootType=*/builder.getType<AnyOpType>(),
+      /*resultType=*/TypeRange{},
+      [](OpBuilder &b, Location nested, Value rootH) {
+        b.create<YieldOp>(nested, ValueRange());
+      },
+      /*args=*/ArrayRef<NamedAttribute>{},
+      /*attrArgs=*/
+      ArrayRef<DictionaryAttr>{
+          builder.getDictionaryAttr(ArrayRef<NamedAttribute>{
+              builder.getNamedAttr(TransformDialect::kArgConsumedAttrName,
+                                   builder.getUnitAttr())})});
+
+  // 检查 body argument 上有没有 transform.consumed
+  Block &body = seqOp.getBody().front();
+  ASSERT_EQ(body.getNumArguments(), 1u);
+
+  StringAttr arg0Name = seqOp.getArgAttrsAttrName();
+  EXPECT_TRUE(arg0Name);
+}
\ No newline at end of file