[mlir] Add use nameloc to OpPrintingFlags

[jpienaar]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Jacques Pienaar (jpienaar)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/129584.diff

3 Files Affected:

• (modified) mlir/include/mlir/IR/OperationSupport.h (+3)
• (modified) mlir/lib/IR/AsmPrinter.cpp (+5)
• (modified) mlir/unittests/IR/OperationSupportTest.cpp (+20)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index fbe0ed29a4d11..2b0e50437afcc 100644
--- 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
index 1f22d4f37a813..16bc857ad3416 100644
--- 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
index f94dc78445807..c03b79aae3321 100644
--- 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)");
+
+  op->destroy();
+}
+
 TEST(NamedAttrListTest, TestAppendAssign) {
   MLIRContext ctx;
   NamedAttrList attrs;

[makslevental]

Can you plumb through to Python pretty please :)

[jpienaar]

Can you plumb through to Python pretty please :)

I'll look in follow up yes.