[MLIR][python bindings] Add `PyValue.get_name`

[makslevental]

Summary

Useful for easier debugging (no need to regex out all of the stuff around the id).

[makslevental]

mlir/lib/Bindings/Python/IRCore.cpp

3344

Probably this should be named something other than print but the conventional name escapes me (only thing that comes to mind is show vs print in Haskell).

[mehdi_amini]

mlir/include/mlir-c/IR.h

780

Can you document the arguments? In particular the parentOp here and the contract with "value"

[makslevental]

move parent into C function

[makslevental]

mlir/include/mlir-c/IR.h

780

Prior I had the "get parent" stuff inside the C function and moved it out because it seemed like most of the functions in IR.cpp are "small"; your comment made me realize it's much safer/better to have them there in order to enforce the relationship.

mlir/lib/CAPI/IR/IR.cpp

781

Block::printAsOperand similarly prints <<UNLINKED BLOCK>> in such a case (see here).

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3344

imo this should be the behaviour of __str__, and the current __str__ should be __repr__

mlir/lib/CAPI/IR/IR.cpp

775–779

To get the expected result for values in nested regions I think you'll need to continue traversing up to find the top-level op, similar to here: https://github.com/llvm/llvm-project/blob/2e3cabe172c6f2eaf1f097ffeff1664b3768223a/mlir/lib/IR/AsmPrinter.cpp#L3680-L3691

780–783

In this case I think printing with AsmState(ctx) should work fine?

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3344

on second thought, if the current __str__ prints the full operation, using it for __repr__ is probably not a good idea

[rkayaith]

mlir/lib/CAPI/IR/IR.cpp

780–783

nevermind, after looking at the implementation of that, it would just end up printing something similar to what you did already

[makslevental]

handle vals in nested regions correctly

mlir/lib/Bindings/Python/IRCore.cpp

3344

I agree! I can make the change and then patch the failing lit tests if there's a consensus.

[makslevental]

mlir/lib/Bindings/Python/IRCore.cpp

3344

f the interpretations of __repr__ and __str__ are

• __repr__ is to be unambiguous
• __str__ is to be readable

then I think printing only the valueid for __str__ and printing the whole op for __repr__ would be the right thing to do?

I mean it doesn't matter to me as long as there's someway to print the valueid (i.e., this dunder).

mlir/lib/CAPI/IR/IR.cpp

780–783

With AsmState(ctx) you get %0 always I believe, which either is or isn't counterintutive.

[makslevental]

add one more test case

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

162–164

I'm not sure this should be documented here, it seems more like a bug in the C++ API I think?

3344

If you keep this as a method, how about get_name or get_id?

mlir/lib/CAPI/IR/IR.cpp

776–777

This traversal should happen for block args as well

795

This is never freed, I'd just go back to printing the << ... >> string if that simplifies things

[makslevental]

traverse for block args, free asmstate, handle corner case

[makslevental]

clang-format

[rkayaith]

mlir/include/mlir-c/IR.h

779–780

The "this isn't canonical" comment doesn't seem like it should be here (documenting a bug?)

mlir/lib/Bindings/Python/IRCore.cpp

3344

imo get_name by itself is unambiguous and the _as_operand doesn't add anything here

3352

to be consistent with ir.Operation.print

mlir/lib/CAPI/IR/IR.cpp

770

There's additional printing options that can affect the value id, like printGenericOpForm. If you want to support different options it seems better to just take a MlirOpPrintingFlags here (probably not worth exposing all the options python-side though, that seems easy to extend later).

778–779

780–785

Values can only be OpResults or BlockArguments, so this else can't be reached. I think the check should look something like this instead

788

I would add a Value::printAsOperand(raw_ostream &os, const OpPrintingFlags &flags={}) overload that does this. That would also let you share the logic here with Operation::print.

800–802

[makslevental]

refactor

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3347

You'll need to call mlirOpPrintingFlagsDestroy at the end to avoid this leaking

3348–3350

llvm style is to not use braces here

mlir/lib/IR/AsmPrinter.cpp

3676–3677 ↗

(On Diff #520219)

the llvm style is to use static for functions instead of anonymous namespaces

3699–3700 ↗

(On Diff #520219)

cast() will do this assert for you

[makslevental]

fix mistakes

[rkayaith]

LGTM, can you also update the revision title

mlir/include/mlir-c/IR.h

 MLIR_CAPI_EXPORTED void mlirValueDump(MlirValue value);
 
 /// Prints a value by sending chunks of the string representation and
 /// forwarding `userData to `callback`. Note that the callback may be called
 /// several times with consecutive chunks of the string.
 MLIR_CAPI_EXPORTED void
 mlirValuePrint(MlirValue value, MlirStringCallback callback, void *userData);
 
+/// Prints a value as an operand (i.e., the ValueID). Note, this isn't
+/// canonical (see comment in AsmPrinter::Value::printAsOperand.
+MLIR_CAPI_EXPORTED void mlirValuePrintAsOperand(MlirValue value,
+                                                bool shouldUseLocalScope,
+                                                MlirStringCallback callback,
+                                                void *userData);
+
 /// Returns an op operand representing the first use of the value, or a null op
 /// operand if there are no uses.
 MLIR_CAPI_EXPORTED MlirOpOperand mlirValueGetFirstUse(MlirValue value);
 
 /// Replace all uses of 'of' value with the 'with' value, updating anything in
 /// the IR that uses 'of' to use the other value instead.  When this returns
 /// there are zero uses of 'of'.
 MLIR_CAPI_EXPORTED void mlirValueReplaceAllUsesOfWith(MlirValue of,

mlir/lib/Bindings/Python/IRCore.cpp

 static const char kValueDunderStrDocstring[] =
     R"(Returns the string form of the value.
 
 If the value is a block argument, this is the assembly form of its type and the
 position in the argument list. If the value is an operation result, this is
 equivalent to printing the operation that produced it.
 )";
 
+static const char kGetNameAsOperand[] =
+    R"(Returns the string form of value as an operand (i.e., the ValueID).
+)";
+
 static const char kValueReplaceAllUsesWithDocstring[] =
     R"(Replace all uses of value with the new value, updating anything in
 the IR that uses 'self' to use the other value instead.
 )";
 
 //------------------------------------------------------------------------------
 // Utilities.
 //------------------------------------------------------------------------------
 
 /// Helper for creating an @classmethod.
py::class_<PyValue>(m, "Value", py::module_local())
             PyPrintAccumulator printAccum;
             printAccum.parts.append("Value(");
             mlirValuePrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             printAccum.parts.append(")");
             return printAccum.join();
           },
           kValueDunderStrDocstring)
+      .def(
+          "get_name_as_operand",
+          [](PyValue &self, bool shouldUseLocalScope) {
+            PyPrintAccumulator printAccum;
+            mlirValuePrintAsOperand(self.get(), shouldUseLocalScope,
+                                    printAccum.getCallback(),
+                                    printAccum.getUserData());
+            return printAccum.join();
+          },
+          py::arg("should_use_local_scope") = false, kGetNameAsOperand)
       .def_property_readonly("type",
                              [](PyValue &self) {
                                return PyType(
                                    self.getParentOperation()->getContext(),
                                    mlirValueGetType(self.get()));
                              })
       .def(
           "replace_all_uses_with",

mlir/lib/CAPI/IR/IR.cpp

 void mlirValueDump(MlirValue value) { unwrap(value).dump(); }
 
 void mlirValuePrint(MlirValue value, MlirStringCallback callback,
                     void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(value).print(stream);
 }
 
+void mlirValuePrintAsOperand(MlirValue value, bool shouldUseLocalScope,
+                             MlirStringCallback callback, void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+
+  Value cppValue = unwrap(value);
+  AsmState *state;
+  Operation *op;
+
+  if (cppValue.isa<OpResult>()) {
+    op = cppValue.cast<OpResult>().getOwner();
+  } else if (cppValue.isa<BlockArgument>()) {
+    op = cppValue.cast<BlockArgument>().getOwner()->getParentOp();
+  } else {
+    stream << "<<UNKNOWN SSA VALUE>>";
+    return;
+  }
+
+  do {
+    // If we are printing local scope, stop at the first operation that is
+    // isolated from above.
+    if (shouldUseLocalScope && op->hasTrait<OpTrait::IsIsolatedFromAbove>())
+      break;
+
+    // Otherwise, traverse up to the next parent.
+    Operation *parentOp = op->getParentOp();
+    if (!parentOp)
+      break;
+    op = parentOp;
+  } while (true);
+
+  state = new AsmState(op);
+  cppValue.printAsOperand(stream, *state);
+  delete state;
+}
+
 MlirOpOperand mlirValueGetFirstUse(MlirValue value) {
   Value cppValue = unwrap(value);
   if (cppValue.use_empty())
     return {};
 
   OpOperand *opOperand = cppValue.use_begin().getOperand();
 
   return wrap(opOperand);

mlir/test/python/ir/value.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
+from mlir.dialects import func
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f
module = Module.parse(
     }""", ctx)
 
   [func] = module.body.operations
   block = func.entry_block
   op, ret = block.operations
   assert hash(block.arguments[0]) == hash(op.operands[0])
   assert hash(op.result) == hash(ret.operands[0])
 
+
 # CHECK-LABEL: TEST: testValueUses
 @run
 def testValueUses():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signless(32)
     module = Module.create()
     with InsertionPoint(module.body):
       value = Operation.create("custom.op1", results=[i32]).results[0]
       op1 = Operation.create("custom.op2", operands=[value])
       op2 = Operation.create("custom.op2", operands=[value])
 
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   for use in value.uses:
     assert use.owner in [op1, op2]
     print(f"Use owner: {use.owner}")
     print(f"Use operand_number: {use.operand_number}")
 
+
 # CHECK-LABEL: TEST: testValueReplaceAllUsesWith
 @run
 def testValueReplaceAllUsesWith():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signless(32)
     module = Module.create()
def testValueReplaceAllUsesWith():
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   for use in value2.uses:
     assert use.owner in [op1, op2]
     print(f"Use owner: {use.owner}")
     print(f"Use operand_number: {use.operand_number}")
+
+
+# CHECK-LABEL: TEST: testValuePrintAsOperand
+@run
+def testValuePrintAsOperand():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    i32 = IntegerType.get_signless(32)
+    module = Module.create()
+    with InsertionPoint(module.body):
+      value = Operation.create("custom.op1", results=[i32]).results[0]
+      # CHECK: Value(%[[VAL1:.*]] = "custom.op1"() : () -> i32)
+      print(value)
+
+      value2 = Operation.create("custom.op2", results=[i32]).results[0]
+      # CHECK: Value(%[[VAL2:.*]] = "custom.op2"() : () -> i32)
+      print(value2)
+
+      f = func.FuncOp("test", ([i32, i32], []))
+      entry_block1 = Block.create_at_start(f.operation.regions[0], [i32, i32])
+
+      with InsertionPoint(entry_block1):
+        value3 = Operation.create("custom.op3", results=[i32]).results[0]
+        # CHECK: Value(%[[VAL3:.*]] = "custom.op3"() : () -> i32)
+        print(value3)
+        value4 = Operation.create("custom.op4", results=[i32]).results[0]
+        # CHECK: Value(%[[VAL4:.*]] = "custom.op4"() : () -> i32)
+        print(value4)
+
+        f = func.FuncOp("test", ([i32, i32], []))
+        entry_block2 = Block.create_at_start(f.operation.regions[0], [i32, i32])
+        with InsertionPoint(entry_block2):
+          value5 = Operation.create("custom.op5", results=[i32]).results[0]
+          # CHECK: Value(%[[VAL5:.*]] = "custom.op5"() : () -> i32)
+          print(value5)
+          value6 = Operation.create("custom.op6", results=[i32]).results[0]
+          # CHECK: Value(%[[VAL6:.*]] = "custom.op6"() : () -> i32)
+          print(value6)
+
+          func.ReturnOp([])
+
+        func.ReturnOp([])
+
+    # CHECK: %[[VAL1]]
+    print(value.get_name_as_operand())
+    # CHECK: %[[VAL2]]
+    print(value2.get_name_as_operand())
+    # CHECK: %[[VAL3]]
+    print(value3.get_name_as_operand())
+    # CHECK: %[[VAL4]]
+    print(value4.get_name_as_operand())
+
+    # CHECK: %0
+    print(value3.get_name_as_operand(should_use_local_scope=True))
+    # CHECK: %1
+    print(value4.get_name_as_operand(should_use_local_scope=True))
+
+    # CHECK: %[[VAL5]]
+    print(value5.get_name_as_operand())
+    # CHECK: %[[VAL6]]
+    print(value6.get_name_as_operand())
+
+    # CHECK: %[[ARG0:.*]]
+    print(entry_block1.arguments[0].get_name_as_operand())
+    # CHECK: %[[ARG1:.*]]
+    print(entry_block1.arguments[1].get_name_as_operand())
+
+    # CHECK: %[[ARG2:.*]]
+    print(entry_block2.arguments[0].get_name_as_operand())
+    # CHECK: %[[ARG3:.*]]
+    print(entry_block2.arguments[1].get_name_as_operand())
+
+    # CHECK: module {
+    # CHECK:   %[[VAL1]] = "custom.op1"() : () -> i32
+    # CHECK:   %[[VAL2]] = "custom.op2"() : () -> i32
+    # CHECK:   func.func @test(%[[ARG0]]: i32, %[[ARG1]]: i32) {
+    # CHECK:     %[[VAL3]] = "custom.op3"() : () -> i32
+    # CHECK:     %[[VAL4]] = "custom.op4"() : () -> i32
+    # CHECK:     func @test(%[[ARG2]]: i32, %[[ARG3]]: i32) {
+    # CHECK:       %[[VAL5]] = "custom.op5"() : () -> i32
+    # CHECK:       %[[VAL6]] = "custom.op6"() : () -> i32
+    # CHECK:       return
+    # CHECK:     }
+    # CHECK:     return
+    # CHECK:   }
+    # CHECK: }
+    print(module)
+
+    value2.owner.detach_from_parent()
+    # CHECK: %0
+    print(value2.get_name_as_operand())