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

[Maksim Levental]

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

Add `use_name_loc_as_prefix` to `value.get_name()`.

>From 02dc3b837c292ac95500d50663ea13f801c7de7b 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 |  2 +-
 mlir/test/python/ir/value.py             | 10 ++++++++--
 3 files changed, 14 insertions(+), 5 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..af463cfe45d0c 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:
         """
diff --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py
index 9a8146bd9350b..57a5bec7e9659 100644
--- a/mlir/test/python/ir/value.py
+++ b/mlir/test/python/ir/value.py
@@ -248,6 +248,10 @@ def testValuePrintAsOperand():
                 print(value4)
                 func.ReturnOp([])
 
+            named_value = Operation.create(
+                "custom.op3", results=[i32], location=Location.name("apple")
+            ).results[0]
+
         # CHECK: %[[VAL1]]
         print(value.get_name())
         # CHECK: %[[VAL2]]
@@ -292,6 +296,9 @@ def testValuePrintAsOperand():
         # CHECK: %0
         print(value2.get_name())
 
+        # CHECK: %apple
+        print(named_value.get_name(use_name_loc_as_prefix=True))
+
 
 # CHECK-LABEL: TEST: testValueSetType
 @run
@@ -404,8 +411,7 @@ def reduction(arg0, arg1):
     try:
 
         @register_value_caster(IntegerType.static_typeid)
-        def dont_cast_int_shouldnt_register(v):
-            ...
+        def dont_cast_int_shouldnt_register(v): ...
 
     except RuntimeError as e:
         # CHECK: Value caster is already registered: {{.*}}cast_int