[Mlir-commits] [mlir] [mlir][Python] downcast ir.Value to BlockArgument or OpResult (PR #175264)
Maksim Levental
llvmlistbot at llvm.org
Fri Jan 9 19:52:53 PST 2026
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/175264
>From 7c2c7f0ac6e721e42445d45f3ec3fff1e1368aa3 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Fri, 9 Jan 2026 15:02:44 -0800
Subject: [PATCH] [mlir][Python] downcast ir.Value to BlockArgument or OpResult
---
mlir/include/mlir/Bindings/Python/IRCore.h | 13 ++--
mlir/include/mlir/Bindings/Python/Nanobind.h | 1 +
.../mlir/Bindings/Python/NanobindUtils.h | 2 +-
mlir/lib/Bindings/Python/IRAffine.cpp | 2 +-
mlir/lib/Bindings/Python/IRCore.cpp | 61 +++++++++----------
mlir/test/python/dialects/python_test.py | 2 +-
mlir/test/python/ir/value.py | 22 +++----
7 files changed, 50 insertions(+), 53 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 330318683c15e..bfc35ca5b9d50 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -36,20 +36,21 @@
namespace mlir {
namespace python {
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
-
+class DefaultingPyLocation;
+class DefaultingPyMlirContext;
class PyBlock;
+class PyBlockArgument;
class PyDiagnostic;
class PyDiagnosticHandler;
class PyInsertionPoint;
class PyLocation;
-class DefaultingPyLocation;
class PyMlirContext;
-class DefaultingPyMlirContext;
class PyModule;
+class PyOpResult;
class PyOperation;
class PyOperationBase;
-class PyType;
class PySymbolTable;
+class PyType;
class PyValue;
/// Wrapper for the global LLVM debugging flag.
@@ -1188,7 +1189,9 @@ class MLIR_PYTHON_API_EXPORTED PyValue {
/// Gets a capsule wrapping the void* within the MlirValue.
nanobind::object getCapsule();
- nanobind::typed<nanobind::object, PyValue> maybeDownCast();
+ nanobind::typed<nanobind::object,
+ std::variant<PyBlockArgument, PyOpResult, PyValue>>
+ maybeDownCast();
/// Creates a PyValue from the MlirValue wrapped by a capsule. Ownership of
/// the underlying MlirValue is still tied to the owning operation.
diff --git a/mlir/include/mlir/Bindings/Python/Nanobind.h b/mlir/include/mlir/Bindings/Python/Nanobind.h
index 8dc8a0d063d70..f7cf8fd38981d 100644
--- a/mlir/include/mlir/Bindings/Python/Nanobind.h
+++ b/mlir/include/mlir/Bindings/Python/Nanobind.h
@@ -29,6 +29,7 @@
#include <nanobind/stl/string.h>
#include <nanobind/stl/string_view.h>
#include <nanobind/stl/tuple.h>
+#include <nanobind/stl/variant.h>
#include <nanobind/stl/vector.h>
#include <nanobind/typing.h>
#if defined(__clang__) || defined(__GNUC__)
diff --git a/mlir/include/mlir/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
index aea195fecae82..215daf245b902 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindUtils.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
@@ -277,7 +277,7 @@ class Sliceable {
/// Returns the element at the given slice index. Supports negative indices
/// by taking elements in inverse order. Returns a nullptr object if out
/// of bounds.
- nanobind::object getItem(intptr_t index) {
+ nanobind::typed<nanobind::object, ElementTy> getItem(intptr_t index) {
// Negative indices mean we count from the end.
index = wrapIndex(index);
if (index < 0) {
diff --git a/mlir/lib/Bindings/Python/IRAffine.cpp b/mlir/lib/Bindings/Python/IRAffine.cpp
index 3c2da03181e6a..2e760e6e6f830 100644
--- a/mlir/lib/Bindings/Python/IRAffine.cpp
+++ b/mlir/lib/Bindings/Python/IRAffine.cpp
@@ -614,7 +614,7 @@ void populateIRAffine(nb::module_ &m) {
return PyAffineExpr(self.getContext(),
mlirAffineExprCompose(self, other));
})
- .def("maybe_downcast", &PyAffineExpr::maybeDownCast)
+ .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, &PyAffineExpr::maybeDownCast)
.def(
"shift_dims",
[](PyAffineExpr &self, uint32_t numDims, uint32_t shift,
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 19db41fae4fe2..646490bba366b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -10,6 +10,7 @@
#include "mlir/Bindings/Python/Globals.h"
#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bindings/Python/NanobindUtils.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
// clang-format on
#include "mlir-c/BuiltinAttributes.h"
@@ -17,10 +18,6 @@
#include "mlir-c/Diagnostics.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
-#include "mlir/Bindings/Python/Nanobind.h"
-#include "mlir/Bindings/Python/NanobindAdaptors.h"
-#include "nanobind/nanobind.h"
-#include "nanobind/typing.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
@@ -2003,7 +2000,22 @@ nb::object PyValue::getCapsule() {
return nb::steal<nb::object>(mlirPythonValueToCapsule(get()));
}
-nb::typed<nb::object, PyValue> PyValue::maybeDownCast() {
+static PyOperationRef getValueOwnerRef(MlirValue value) {
+ MlirOperation owner;
+ if (mlirValueIsAOpResult(value))
+ owner = mlirOpResultGetOwner(value);
+ else if (mlirValueIsABlockArgument(value))
+ owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));
+ else
+ assert(false && "Value must be an block arg or op result.");
+ if (mlirOperationIsNull(owner))
+ throw nb::python_error();
+ MlirContext ctx = mlirOperationGetContext(owner);
+ return PyOperation::forOperation(PyMlirContext::forContext(ctx), owner);
+}
+
+nb::typed<nb::object, std::variant<PyBlockArgument, PyOpResult, PyValue>>
+PyValue::maybeDownCast() {
MlirType type = mlirValueGetType(get());
MlirTypeID mlirTypeID = mlirTypeGetTypeID(type);
assert(!mlirTypeIDIsNull(mlirTypeID) &&
@@ -2012,26 +2024,23 @@ nb::typed<nb::object, PyValue> PyValue::maybeDownCast() {
PyGlobals::get().lookupValueCaster(mlirTypeID, mlirTypeGetDialect(type));
// nb::rv_policy::move means use std::move to move the return value
// contents into a new instance that will be owned by Python.
- nb::object thisObj = nb::cast(this, nb::rv_policy::move);
- if (!valueCaster)
- return thisObj;
- return valueCaster.value()(thisObj);
+ nb::object thisObj;
+ if (mlirValueIsAOpResult(value))
+ thisObj = nb::cast<PyOpResult>(*this, nb::rv_policy::move);
+ else if (mlirValueIsABlockArgument(value))
+ thisObj = nb::cast<PyBlockArgument>(*this, nb::rv_policy::move);
+ else
+ assert(false && "Value must be an block arg or op result.");
+ if (valueCaster)
+ return valueCaster.value()(thisObj);
+ return thisObj;
}
PyValue PyValue::createFromCapsule(nb::object capsule) {
MlirValue value = mlirPythonCapsuleToValue(capsule.ptr());
if (mlirValueIsNull(value))
throw nb::python_error();
- MlirOperation owner;
- if (mlirValueIsAOpResult(value))
- owner = mlirOpResultGetOwner(value);
- if (mlirValueIsABlockArgument(value))
- owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));
- if (mlirOperationIsNull(owner))
- throw nb::python_error();
- MlirContext ctx = mlirOperationGetContext(owner);
- PyOperationRef ownerRef =
- PyOperation::forOperation(PyMlirContext::forContext(ctx), owner);
+ PyOperationRef ownerRef = getValueOwnerRef(value);
return PyValue(ownerRef, value);
}
@@ -2279,15 +2288,7 @@ intptr_t PyOpOperandList::getRawNumElements() {
PyValue PyOpOperandList::getRawElement(intptr_t pos) {
MlirValue operand = mlirOperationGetOperand(operation->get(), pos);
- MlirOperation owner;
- if (mlirValueIsAOpResult(operand))
- owner = mlirOpResultGetOwner(operand);
- else if (mlirValueIsABlockArgument(operand))
- owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(operand));
- else
- assert(false && "Value must be an block arg or op result.");
- PyOperationRef pyOwner =
- PyOperation::forOperation(operation->getContext(), owner);
+ PyOperationRef pyOwner = getValueOwnerRef(operand);
return PyValue(pyOwner, operand);
}
@@ -4679,9 +4680,7 @@ void populateIRCore(nb::module_ &m) {
"with_"_a, "exceptions"_a, kValueReplaceAllUsesExceptDocstring)
.def(
MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,
- [](PyValue &self) -> nb::typed<nb::object, PyValue> {
- return self.maybeDownCast();
- },
+ [](PyValue &self) { return self.maybeDownCast(); },
"Downcasts the `Value` to a more specific kind if possible.")
.def_prop_ro(
"location",
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 9c0966d2d8798..8814f4d66fdf9 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -862,7 +862,7 @@ def values(lst):
]
is Value
)
- assert type(variadic_operands.non_variadic) is Value
+ assert type(variadic_operands.non_variadic) is OpResult
# CHECK: ['Value(%{{.*}} = arith.constant 0 : i32)', 'Value(%{{.*}} = arith.constant 1 : i32)']
print(values(variadic_operands.variadic1))
diff --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py
index 45efb880bab44..4cd1c2fa96c27 100644
--- a/mlir/test/python/ir/value.py
+++ b/mlir/test/python/ir/value.py
@@ -347,13 +347,6 @@ def __init__(self, v):
def __str__(self):
return super().__str__().replace(Value.__name__, NOPResult.__name__)
- class NOPValue(Value):
- def __init__(self, v):
- super().__init__(v)
-
- def __str__(self):
- return super().__str__().replace(Value.__name__, NOPValue.__name__)
-
class NOPBlockArg(BlockArgument):
def __init__(self, v):
super().__init__(v)
@@ -362,14 +355,12 @@ def __str__(self):
return super().__str__().replace(Value.__name__, NOPBlockArg.__name__)
@register_value_caster(IntegerType.static_typeid)
- def cast_int(v) -> Value:
+ def cast_int(v) -> NOPResult | NOPBlockArg:
print("in caster", v.__class__.__name__)
if isinstance(v, OpResult):
return NOPResult(v)
if isinstance(v, BlockArgument):
return NOPBlockArg(v)
- elif isinstance(v, Value):
- return NOPValue(v)
ctx = Context()
ctx.allow_unregistered_dialects = True
@@ -400,12 +391,15 @@ def cast_int(v) -> Value:
# CHECK: "custom.op2"(%0#0, %0#1) : (i32, i32) -> ()
print(op1)
- # CHECK: in caster Value
- # CHECK: operand 0 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))
+ # CHECK: in caster OpResult
+ # CHECK: operand 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
print("operand 0", op1.operands[0])
- # CHECK: in caster Value
- # CHECK: operand 1 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))
+ assert isinstance(op1.operands[0], Value)
+ assert isinstance(op1.operands[0], OpResult)
+ # CHECK: in caster OpResult
+ # CHECK: operand 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))
print("operand 1", op1.operands[1])
+ assert isinstance(op1.operands[1], OpResult)
# CHECK: in caster BlockArgument
# CHECK: in caster BlockArgument
More information about the Mlir-commits
mailing list