[mlir][python] add use_name_loc_as_prefix to value.get_name()

[makslevental]

Add use_name_loc_as_prefix to value.get_name().

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

Add use_name_loc_as_prefix to value.get_name().

---

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

3 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+5-2)
• (modified) mlir/python/mlir/_mlir_libs/_mlir/ir.pyi (+1-1)
• (modified) mlir/test/python/ir/value.py (+8-2)

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

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.