[Mlir-commits] [mlir] [mlir][py] ability to downcast AffineExpr after #172892 (PR #174808)

[Oleksandr Alex Zinenko]

https://github.com/ftynse updated https://github.com/llvm/llvm-project/pull/174808

>From 25118ce5932e3490bb8c8cd7a2c5b70450181c34 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <git at ozinenko.com>
Date: Wed, 7 Jan 2026 18:20:12 +0100
Subject: [PATCH] [mlir][py] ability to downcast AffineExpr after #172892

AffineExpr is a separate hierarchy of LLVM-style nested classes that doesn't
rely on TypeID and is not extensible. We need the ability to downcast the
Python equivalent of those to a specific subclass that was seemingly lost in
PR #172892. Bring it back by having an explicit cast. We don't really need
user-defined type casters here since AffineExpr is entirely closed and not
typed, unlike values.
---
 mlir/include/mlir/Bindings/Python/IRCore.h |  2 ++
 mlir/lib/Bindings/Python/IRAffine.cpp      | 30 +++++++++++++++++++---
 mlir/test/python/ir/affine_expr.py         | 11 ++++++++
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 729cbb6df3267..1f19683dfe80d 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1214,6 +1214,8 @@ class MLIR_PYTHON_API_EXPORTED PyAffineExpr : public BaseContextObject {
   PyAffineExpr ceilDiv(const PyAffineExpr &other) const;
   PyAffineExpr mod(const PyAffineExpr &other) const;
 
+  nanobind::typed<nanobind::object, PyAffineExpr> maybeDownCast();
+
 private:
   MlirAffineExpr affineExpr;
 };
diff --git a/mlir/lib/Bindings/Python/IRAffine.cpp b/mlir/lib/Bindings/Python/IRAffine.cpp
index b3d15ee59566b..3c2da03181e6a 100644
--- a/mlir/lib/Bindings/Python/IRAffine.cpp
+++ b/mlir/lib/Bindings/Python/IRAffine.cpp
@@ -193,14 +193,14 @@ class PyAffineBinaryExpr : public PyConcreteAffineExpr<PyAffineBinaryExpr> {
   static constexpr const char *pyClassName = "AffineBinaryExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
-  PyAffineExpr lhs() {
+  nb::typed<nb::object, PyAffineExpr> lhs() {
     MlirAffineExpr lhsExpr = mlirAffineBinaryOpExprGetLHS(get());
-    return PyAffineExpr(getContext(), lhsExpr);
+    return PyAffineExpr(getContext(), lhsExpr).maybeDownCast();
   }
 
-  PyAffineExpr rhs() {
+  nb::typed<nb::object, PyAffineExpr> rhs() {
     MlirAffineExpr rhsExpr = mlirAffineBinaryOpExprGetRHS(get());
-    return PyAffineExpr(getContext(), rhsExpr);
+    return PyAffineExpr(getContext(), rhsExpr).maybeDownCast();
   }
 
   static void bindDerived(ClassTy &c) {
@@ -375,6 +375,27 @@ PyAffineExpr PyAffineExpr::createFromCapsule(const nb::object &capsule) {
       rawAffineExpr);
 }
 
+nb::typed<nb::object, PyAffineExpr> PyAffineExpr::maybeDownCast() {
+  MlirAffineExpr expr = get();
+  if (mlirAffineExprIsAConstant(expr))
+    return nb::cast(PyAffineConstantExpr(getContext(), expr));
+  if (mlirAffineExprIsADim(expr))
+    return nb::cast(PyAffineDimExpr(getContext(), expr));
+  if (mlirAffineExprIsASymbol(expr))
+    return nb::cast(PyAffineSymbolExpr(getContext(), expr));
+  if (mlirAffineExprIsAAdd(expr))
+    return nb::cast(PyAffineAddExpr(getContext(), expr));
+  if (mlirAffineExprIsAMul(expr))
+    return nb::cast(PyAffineMulExpr(getContext(), expr));
+  if (mlirAffineExprIsAMod(expr))
+    return nb::cast(PyAffineModExpr(getContext(), expr));
+  if (mlirAffineExprIsAFloorDiv(expr))
+    return nb::cast(PyAffineFloorDivExpr(getContext(), expr));
+  if (mlirAffineExprIsACeilDiv(expr))
+    return nb::cast(PyAffineCeilDivExpr(getContext(), expr));
+  return nb::cast(*this);
+}
+
 //------------------------------------------------------------------------------
 // PyAffineMap and utilities.
 //------------------------------------------------------------------------------
@@ -593,6 +614,7 @@ void populateIRAffine(nb::module_ &m) {
              return PyAffineExpr(self.getContext(),
                                  mlirAffineExprCompose(self, other));
            })
+      .def("maybe_downcast", &PyAffineExpr::maybeDownCast)
       .def(
           "shift_dims",
           [](PyAffineExpr &self, uint32_t numDims, uint32_t shift,
diff --git a/mlir/test/python/ir/affine_expr.py b/mlir/test/python/ir/affine_expr.py
index 82c509efdf8fc..cf1d5258333d7 100644
--- a/mlir/test/python/ir/affine_expr.py
+++ b/mlir/test/python/ir/affine_expr.py
@@ -424,3 +424,14 @@ def testAffineExprSimplify():
     with Context() as ctx:
         expr = AffineExpr.get_dim(0) + AffineExpr.get_symbol(0)
         assert expr == AffineExpr.simplify_affine_expr(expr, 1, 1)
+
+
+# CHECK-LABEL: TEST: testAffineExprDowncast
+ at run
+def testAffineExprDowncast():
+    with Context() as ctx:
+        expr = AffineExpr.get_dim(0) + AffineExpr.get_symbol(0)
+        assert isinstance(expr.lhs, AffineDimExpr)
+        assert isinstance(expr.rhs, AffineSymbolExpr)
+        assert expr.lhs.position == 0
+        assert expr.rhs.position == 0