[Mlir-commits] [mlir] [mlir][Python] downcast ir.Value to BlockArgument or OpResult (PR #175264)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/175264

None

>From c7a3e5dbfb669542a967dff201b87bf0bcad8f7d 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/lib/Bindings/Python/IRCore.cpp        | 36 ++++++++++++++--------
 2 files changed, 31 insertions(+), 18 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/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 19db41fae4fe2..0926229da404e 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -20,6 +20,7 @@
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 #include "nanobind/nanobind.h"
+#include "nanobind/stl/variant.h"
 #include "nanobind/typing.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
@@ -2003,7 +2004,20 @@ 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);
+  if (mlirValueIsABlockArgument(value))
+    owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));
+  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) &&
@@ -2013,8 +2027,13 @@ nb::typed<nb::object, PyValue> PyValue::maybeDownCast() {
   // 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;
+  if (!valueCaster) {
+    PyOperationRef ownerRef = getValueOwnerRef(value);
+    if (mlirValueIsAOpResult(value))
+      return nb::cast(PyOpResult(ownerRef, value));
+    if (mlirValueIsABlockArgument(value))
+      return nb::cast(PyBlockArgument(ownerRef, value));
+  }
   return valueCaster.value()(thisObj);
 }
 
@@ -2022,16 +2041,7 @@ 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);
 }