[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

782

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

776–780

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

781–784

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

781–784

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

781–784

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

777–778

This traversal should happen for block args as well

796

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

771

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).

779–780

781–786

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

789

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.

801–803

[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

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

3699–3700

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).
+MLIR_CAPI_EXPORTED void mlirValuePrintAsOperand(MlirValue value,
+                                                MlirOpPrintingFlags flags,
+                                                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/include/mlir/IR/Value.h

public:
 
   void print(raw_ostream &os);
   void print(raw_ostream &os, const OpPrintingFlags &flags);
   void print(raw_ostream &os, AsmState &state);
   void dump();
 
   /// Print this value as if it were an operand.
   void printAsOperand(raw_ostream &os, AsmState &state);
+  void printAsOperand(raw_ostream &os, const OpPrintingFlags &flags);
 
   /// Methods for supporting PointerLikeTypeTraits.
   void *getAsOpaquePointer() const { return impl; }
   static Value getFromOpaquePointer(const void *pointer) {
     return reinterpret_cast<detail::ValueImpl *>(const_cast<void *>(pointer));
   }
   detail::ValueImpl *getImpl() const { return impl; }
 

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",
+          [](PyValue &self, bool useLocalScope) {
+            PyPrintAccumulator printAccum;
+            MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
+            if (useLocalScope) {
+              mlirOpPrintingFlagsUseLocalScope(flags);
+            }
+            mlirValuePrintAsOperand(self.get(), flags, printAccum.getCallback(),
+                                    printAccum.getUserData());
+            return printAccum.join();
+          },
+          py::arg("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

 #include "mlir/CAPI/Support.h"
 #include "mlir/CAPI/Utils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser/Parser.h"
 
 #include <cstddef>
 #include <optional>
 
 using namespace mlir;

 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, MlirOpPrintingFlags flags,
+                             MlirStringCallback callback, void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+  Value cppValue = unwrap(value);
+  cppValue.printAsOperand(stream, *unwrap(flags));
+}
+
 MlirOpOperand mlirValueGetFirstUse(MlirValue value) {
   Value cppValue = unwrap(value);
   if (cppValue.use_empty())
     return {};
 
   OpOperand *opOperand = cppValue.use_begin().getOperand();
 
   return wrap(opOperand);
 }
 
 void mlirValueReplaceAllUsesOfWith(MlirValue oldValue, MlirValue newValue) {
   unwrap(oldValue).replaceAllUsesWith(unwrap(newValue));
 }
 
 //===----------------------------------------------------------------------===//
 // OpOperand API.
 //===----------------------------------------------------------------------===//
 
 bool mlirOpOperandIsNull(MlirOpOperand opOperand) { return !opOperand.ptr; }
 
 MlirOperation mlirOpOperandGetOwner(MlirOpOperand opOperand) {
   return wrap(unwrap(opOperand)->getOwner());
 }
 
 unsigned mlirOpOperandGetOperandNumber(MlirOpOperand opOperand) {
   return unwrap(opOperand)->getOperandNumber();
 }
 
 MlirOpOperand mlirOpOperandGetNextUse(MlirOpOperand opOperand) {
   if (mlirOpOperandIsNull(opOperand))
     return {};
 
   OpOperand *nextOpOperand = static_cast<OpOperand *>(
       unwrap(opOperand)->getNextOperandUsingThisValue());

mlir/lib/IR/AsmPrinter.cpp

 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/Threading.h"
 
-#include <tuple>
 #include <optional>
+#include <tuple>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 #define DEBUG_TYPE "mlir-asm-printer"
 
 void OperationName::print(raw_ostream &os) const { os << getStringRef(); }
 
void Value::printAsOperand(raw_ostream &os, AsmState &state) {
   // TODO: This doesn't necessarily capture all potential cases.
   // Currently, region arguments can be shadowed when printing the main
   // operation. If the IR hasn't been printed, this will produce the old SSA
   // name and not the shadowed name.
   state.getImpl().getSSANameState().printValueID(*this, /*printResultNo=*/true,
                                                  os);
 }
 
-void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
-  // Find the operation to number from based upon the provided flags.
-  Operation *op = this;
-  bool shouldUseLocalScope = printerFlags.shouldUseLocalScope();
+namespace {
+Operation *findParent(Operation *op, bool shouldUseLocalScope) {
   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);
+  return op;
+}
+} // namespace
+
+void Value::printAsOperand(raw_ostream &os, const OpPrintingFlags &flags) {
+  Operation *op;
+  if (auto result = dyn_cast<OpResult>()) {
+    op = result.getOwner();
+  } else {
+    assert(isa<BlockArgument>() && "expected result to be BlockArgument");
+    op = dyn_cast<BlockArgument>().getOwner()->getParentOp();
+    if (!op) {
+      os << "<<UNKNOWN SSA VALUE>>";
+      return;
+    }
+  }
+  op = findParent(op, flags.shouldUseLocalScope());
+  AsmState state(op, flags);
+  printAsOperand(os, state);
+}
 
+void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
+  // Find the operation to number from based upon the provided flags.
+  Operation *op = findParent(this, printerFlags.shouldUseLocalScope());
   AsmState state(op, printerFlags);
   print(os, state);
 }
 void Operation::print(raw_ostream &os, AsmState &state) {
   OperationPrinter printer(os, state.getImpl());
   if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope()) {
     state.getImpl().initializeAliases(this);
     printer.printTopLevelOperation(this);

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())
+    # CHECK: %[[VAL2]]
+    print(value2.get_name())
+    # CHECK: %[[VAL3]]
+    print(value3.get_name())
+    # CHECK: %[[VAL4]]
+    print(value4.get_name())
+
+    # CHECK: %0
+    print(value3.get_name(use_local_scope=True))
+    # CHECK: %1
+    print(value4.get_name(use_local_scope=True))
+
+    # CHECK: %[[VAL5]]
+    print(value5.get_name())
+    # CHECK: %[[VAL6]]
+    print(value6.get_name())
+
+    # CHECK: %[[ARG0:.*]]
+    print(entry_block1.arguments[0].get_name())
+    # CHECK: %[[ARG1:.*]]
+    print(entry_block1.arguments[1].get_name())
+
+    # CHECK: %[[ARG2:.*]]
+    print(entry_block2.arguments[0].get_name())
+    # CHECK: %[[ARG3:.*]]
+    print(entry_block2.arguments[1].get_name())
+
+    # 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())