diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h --- a/mlir/include/mlir-c/IR.h +++ b/mlir/include/mlir-c/IR.h @@ -776,6 +776,13 @@ 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, + MlirOperation parentOp, + 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); diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -156,6 +156,14 @@ equivalent to printing the operation that produced it. )"; +static const char kValuePrintAsOperand[] = + R"(Returns the string form of value as an operand (i.e., the ValueID). + +Note, this isn't canonical; 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. +)"; + 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. @@ -3336,6 +3344,31 @@ return printAccum.join(); }, kValueDunderStrDocstring) + .def( + "print_as_operand", + [](PyValue &self) { + MlirValue value = self.get(); + MlirOperation parent; + if (mlirValueIsAOpResult(value)) { + parent = + mlirOperationGetParentOperation(mlirOpResultGetOwner(value)); + } else if (mlirValueIsABlockArgument(value)) { + parent = + mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value)); + } + if (mlirOperationIsNull(parent)) { + auto origRepr = py::str(py::cast(self)).cast(); + mlirValueDump(value); + throw SetPyError(PyExc_ValueError, + Twine("No parent operation for ") + origRepr); + } + PyPrintAccumulator printAccum; + mlirValuePrintAsOperand(self.get(), parent, + printAccum.getCallback(), + printAccum.getUserData()); + return printAccum.join(); + }, + kValuePrintAsOperand) .def_property_readonly("type", [](PyValue &self) { return PyType( diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp --- a/mlir/lib/CAPI/IR/IR.cpp +++ b/mlir/lib/CAPI/IR/IR.cpp @@ -767,6 +767,15 @@ unwrap(value).print(stream); } +void mlirValuePrintAsOperand(MlirValue value, MlirOperation parentOp, + MlirStringCallback callback, void *userData) { + detail::CallbackOstream stream(callback, userData); + Value cppValue = unwrap(value); + Operation *cppParentOp = unwrap(parentOp); + AsmState state(cppParentOp); + cppValue.printAsOperand(stream, state); +} + MlirOpOperand mlirValueGetFirstUse(MlirValue value) { Value cppValue = unwrap(value); if (cppValue.use_empty()) diff --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py --- a/mlir/test/python/ir/value.py +++ b/mlir/test/python/ir/value.py @@ -2,6 +2,7 @@ import gc from mlir.ir import * +from mlir.dialects import func def run(f): @@ -90,6 +91,7 @@ assert hash(block.arguments[0]) == hash(op.operands[0]) assert hash(op.result) == hash(ret.operands[0]) + # CHECK-LABEL: TEST: testValueUses @run def testValueUses(): @@ -112,6 +114,7 @@ print(f"Use owner: {use.owner}") print(f"Use operand_number: {use.operand_number}") + # CHECK-LABEL: TEST: testValueReplaceAllUsesWith @run def testValueReplaceAllUsesWith(): @@ -137,3 +140,41 @@ 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) + + # CHECK: %[[VAL1]] + print(value.print_as_operand()) + # CHECK: %[[VAL2]] + print(value2.print_as_operand()) + + value2.owner.detach_from_parent() + try: + print(value2.print_as_operand()) + except ValueError as e: + # CHECK: No parent operation for Value(%0 = "custom.op2"() : () -> i32 + print(e) + + f = func.FuncOp("test", ([i32], [])) + entry_block = Block.create_at_start(f.operation.regions[0], [i32, i32]) + + # CHECK: %arg0 + print(entry_block.arguments[0].print_as_operand()) + # CHECK: %arg1 + print(entry_block.arguments[1].print_as_operand())