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 +@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():