[mlir] [transform dialect] NamedSequenceOp build: honor arg_attrs when building

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,

mlir/unittests/Dialect/Transform/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_unittest(MLIRTransformDialectTests
+  TransformNamedSequenceCreate.cpp
   BuildOnlyExtensionTest.cpp
   Preload.cpp
 )

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);
+}