Add Py API for composing an affine expression with a map. Also allows extracting constant values for const expressions.

PiperOrigin-RevId: 284809623

mlir/bindings/python/pybind.cpp

@@ -1122,6 +1122,20 @@ PYBIND11_MODULE(pybind, m) {
            [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
              return PythonAffineExpr(lhs.get() % rhs.get());
            })
+      .def("compose",
+           [](PythonAffineExpr self, PythonAffineMap map) -> PythonAffineExpr {
+             return PythonAffineExpr(self.get().compose(map));
+           })
+      .def(
+          "get_constant_value",
+          [](PythonAffineExpr self) -> py::object {
+            auto const_expr = self.get().dyn_cast<AffineConstantExpr>();
+            if (const_expr)
+              return py::cast(const_expr.getValue());
+            return py::none();
+          },
+          "Returns the constant value for the affine expression if any, or "
+          "returns None.")
       .def("__str__", &PythonAffineExpr::str);
 
   py::class_<PythonAffineMap>(m, "AffineMap",

mlir/bindings/python/test/test_py2and3.py

@@ -195,6 +195,22 @@ class EdscTest:
     # CHECK-LABEL: testCondBr
     #       CHECK:   cond_br %{{.*}}, ^bb{{.*}}, ^bb{{.*}}(%{{.*}} : index)
 
+  def testConstantAffineExpr(self):
+    self.setUp()
+    with self.module.function_context("constant_affine", [], []) as fun:
+      a1 = self.module.affine_dim_expr(0)
+      a2 = self.module.affine_dim_expr(1)
+      a3 = a1 + a2 + 3
+      composedExpr = a3.compose(
+          self.module.affine_map(2, 0, [
+              self.module.affine_constant_expr(4),
+              self.module.affine_constant_expr(7)
+          ]))
+      printWithCurrentFunctionName(str(fun))
+      print("constant value : %d" % composedExpr.get_constant_value())
+    # CHECK-LABEL: testConstantAffineExpr
+    #       CHECK: constant value : 14
+
   def testConstants(self):
     self.setUp()
     with self.module.function_context("constants", [], []) as fun: