[mlir] Add use nameloc to OpPrintingFlags (llvm#129584)

jwanggit86 · Mar 3, 2025 · edb7292 · edb7292
1 parent f6e8366
commit edb7292
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
@@ -1195,6 +1195,9 @@ class OpPrintingFlags {
   /// conflicts across all regions
   OpPrintingFlags &printUniqueSSAIDs(bool enable = true);
 
+  /// Print SSA IDs using their NameLoc, if provided, as prefix.
+  OpPrintingFlags &printNameLocAsPrefix(bool enable = true);
+
   /// Return if the given ElementsAttr should be elided.
   bool shouldElideElementsAttr(ElementsAttr attr) const;
 

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
@@ -327,6 +327,11 @@ bool OpPrintingFlags::shouldPrintElementsAttrWithHex(ElementsAttr attr) const {
          !llvm::isa<SplatElementsAttr>(attr);
 }
 
+OpPrintingFlags &OpPrintingFlags::printNameLocAsPrefix(bool enable) {
+  useNameLocAsPrefix = enable;
+  return *this;
+}
+
 /// Return the size limit for printing large ElementsAttr.
 std::optional<int64_t> OpPrintingFlags::getLargeElementsAttrLimit() const {
   return elementsAttrElementLimit;

diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -230,6 +230,26 @@ TEST(OperationFormatPrintTest, CanUseVariadicFormat) {
   op->destroy();
 }
 
+TEST(OperationFormatPrintTest, CanPrintNameAsPrefix) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  context.allowUnregisteredDialects();
+  Operation *op = Operation::create(
+      NameLoc::get(StringAttr::get(&context, "my_named_loc")),
+      OperationName("t.op", &context), builder.getIntegerType(16), std::nullopt,
+      std::nullopt, nullptr, std::nullopt, 0);
+
+  std::string str;
+  OpPrintingFlags flags;
+  flags.printNameLocAsPrefix(true);
+  llvm::raw_string_ostream os(str);
+  op->print(os, flags);
+  ASSERT_STREQ(str.c_str(), "%my_named_loc = \"t.op\"() : () -> i16\n");
+
+  op->destroy();
+}
+
 TEST(NamedAttrListTest, TestAppendAssign) {
   MLIRContext ctx;
   NamedAttrList attrs;