[Mlir-commits] [mlir] [mlir][python] add use_name_loc_as_prefix to value.get_name() (PR #135052)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/135052

>From fe433a5374470229e410c793dcc4e9771f5265a0 Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 9 Apr 2025 13:07:56 -0400
Subject: [PATCH] [mlir][python] add use_name_loc_as_prefix to value.get_name()

---
 mlir/lib/Bindings/Python/IRCore.cpp      |  7 +++++--
 mlir/python/mlir/_mlir_libs/_mlir/ir.pyi |  4 ++--
 mlir/test/python/ir/value.py             | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 5ffcf671741bd..b5720b7ad8b21 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3977,11 +3977,13 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           kValueDunderStrDocstring)
       .def(
           "get_name",
-          [](PyValue &self, bool useLocalScope) {
+          [](PyValue &self, bool useLocalScope, bool useNameLocAsPrefix) {
             PyPrintAccumulator printAccum;
             MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
             if (useLocalScope)
               mlirOpPrintingFlagsUseLocalScope(flags);
+            if (useNameLocAsPrefix)
+              mlirOpPrintingFlagsPrintNameLocAsPrefix(flags);
             MlirAsmState valueState =
                 mlirAsmStateCreateForValue(self.get(), flags);
             mlirValuePrintAsOperand(self.get(), valueState,
@@ -3991,7 +3993,8 @@ void mlir::python::populateIRCore(nb::module_ &m) {
             mlirAsmStateDestroy(valueState);
             return printAccum.join();
           },
-          nb::arg("use_local_scope") = false)
+          nb::arg("use_local_scope") = false,
+          nb::arg("use_name_loc_as_prefix") = false)
       .def(
           "get_name",
           [](PyValue &self, PyAsmState &state) {
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index c60ff72ff9fd4..1c8080c5d6d2e 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -577,7 +577,7 @@ class Value:
         Dumps a debug representation of the object to stderr.
         """
     @overload
-    def get_name(self, use_local_scope: bool = False) -> str: ...
+    def get_name(self, use_local_scope: bool = False, use_name_loc_as_prefix: bool = True) -> str: ...
     @overload
     def get_name(self, state: AsmState) -> str:
         """
@@ -2382,7 +2382,7 @@ class Operation(_OperationBase):
           attributes: Dict of str:Attribute.
           successors: List of Block for the operation's successors.
           regions: Number of regions to create.
-          location: A Location object (defaults to resolve from context manager).
+          loc: A Location object (defaults to resolve from context manager).
           ip: An InsertionPoint (defaults to resolve from context manager or set to
             False to disable insertion, even with an insertion point set in the
             context manager).
diff --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py
index 9a8146bd9350b..4a241afb8e89d 100644
--- a/mlir/test/python/ir/value.py
+++ b/mlir/test/python/ir/value.py
@@ -293,6 +293,28 @@ def testValuePrintAsOperand():
         print(value2.get_name())
 
 
+# CHECK-LABEL: TEST: testValuePrintAsOperandNamedLocPrefix
+ at run
+def testValuePrintAsOperandNamedLocPrefix():
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    with Location.unknown(ctx):
+        i32 = IntegerType.get_signless(32)
+
+        module = Module.create()
+        with InsertionPoint(module.body):
+            named_value = Operation.create(
+                "custom.op5", results=[i32], loc=Location.name("apple")
+            ).results[0]
+            # CHECK: Value(%[[VAL5:.*]] = "custom.op5"() : () -> i32)
+            print(named_value)
+
+        # CHECK: With use_name_loc_as_prefix
+        # CHECK: %apple
+        print("With use_name_loc_as_prefix")
+        print(named_value.get_name(use_name_loc_as_prefix=True))
+
+
 # CHECK-LABEL: TEST: testValueSetType
 @run
 def testValueSetType():