From edb7292a511171fb1fe75f85fc85464b91130a8f Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Mon, 3 Mar 2025 14:55:36 -0800
Subject: [PATCH] [mlir] Add use nameloc to OpPrintingFlags (#129584)

---
 mlir/include/mlir/IR/OperationSupport.h    |  3 +++
 mlir/lib/IR/AsmPrinter.cpp                 |  5 +++++
 mlir/unittests/IR/OperationSupportTest.cpp | 20 ++++++++++++++++++++
 3 files changed, 28 insertions(+)

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..bac2b72b68deb 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\n");
+
+  op->destroy();
+}
+
 TEST(NamedAttrListTest, TestAppendAssign) {
   MLIRContext ctx;
   NamedAttrList attrs;