bind intmonomial (in order to later bind *polynomial)

Author: makslevental

mlir/include/mlir-c/Dialect/Polynomial.h

@@ -10,14 +10,36 @@
 #ifndef MLIR_C_DIALECT_POLYNOMIAL_H
 #define MLIR_C_DIALECT_POLYNOMIAL_H
 
-#include "mlir-c/IR.h"
+#include "mlir/CAPI/Wrap.h"
+#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Polynomial, polynomial);
 
+#define DEFINE_C_API_STRUCT(name, storage)                                     \
+  struct name {                                                                \
+    storage *ptr;                                                              \
+  };                                                                           \
+  typedef struct name name
+
+DEFINE_C_API_STRUCT(MlirIntMonomial, void);
+
+#undef DEFINE_C_API_STRUCT
+
+DEFINE_C_API_PTR_METHODS(MlirIntMonomial, mlir::polynomial::IntMonomial);
+
+MLIR_CAPI_EXPORTED MlirIntMonomial mlirPolynomialGetIntMonomial(int64_t coeff,
+                                                                uint64_t expo);
+
+MLIR_CAPI_EXPORTED int64_t
+mlirPolynomialIntMonomialGetCoefficient(MlirIntMonomial intMonomial);
+
+MLIR_CAPI_EXPORTED uint64_t
+mlirPolynomialIntMonomialGetExponent(MlirIntMonomial intMonomial);
+
 #ifdef __cplusplus
 }
 #endif

mlir/lib/Bindings/Python/DialectPolynomial.cpp

@@ -18,6 +18,85 @@ using namespace llvm;
 using namespace mlir;
 using namespace mlir::python::adaptors;
 
+class PyIntMonomial {
+public:
+  PyIntMonomial(MlirIntMonomial intMonomial) : intMonomial(intMonomial) {}
+  PyIntMonomial(int64_t coeff, uint64_t expo)
+      : intMonomial(mlirPolynomialGetIntMonomial(coeff, expo)) {}
+  operator MlirIntMonomial() const { return intMonomial; }
+  MlirIntMonomial get() { return intMonomial; }
+
+  int64_t getCoefficient() {
+    return mlirPolynomialIntMonomialGetCoefficient(this->get());
+  }
+
+  uint64_t getExponent() {
+    return mlirPolynomialIntMonomialGetExponent(this->get());
+  }
+
+private:
+  MlirIntMonomial intMonomial;
+};
+
+#define MLIR_PYTHON_CAPSULE_INT_POLYNOMIAL                                     \
+  MAKE_MLIR_PYTHON_QUALNAME("dialects.polynomial.IntMonomial._CAPIPtr")
+
+static inline MlirIntMonomial
+mlirPythonCapsuleToIntMonomial(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_INT_POLYNOMIAL);
+  MlirIntMonomial intMonomial = {ptr};
+  return intMonomial;
+}
+
+static inline PyObject *
+mlirPythonIntMonomialToCapsule(MlirIntMonomial intMonomial) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(intMonomial),
+                       MLIR_PYTHON_CAPSULE_INT_POLYNOMIAL, nullptr);
+}
+
+static inline bool mlirIntMonomialIsNull(MlirIntMonomial intMonomial) {
+  return !intMonomial.ptr;
+}
+
+namespace pybind11 {
+namespace detail {
+
+/// Casts object <-> MlirIntMonomial.
+template <>
+struct type_caster<MlirIntMonomial> {
+  PYBIND11_TYPE_CASTER(MlirIntMonomial, _("MlirIntMonomial"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToIntMonomial(capsule.ptr());
+    return !mlirIntMonomialIsNull(value);
+  }
+
+  static handle cast(MlirIntMonomial v, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonIntMonomialToCapsule(v));
+    return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("dialects.polynomial"))
+        .attr("IntMonomial")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+} // namespace detail
+} // namespace pybind11
+
 PYBIND11_MODULE(_mlirDialectsPolynomial, m) {
   m.doc() = "MLIR Polynomial dialect";
+
+  py::class_<PyIntMonomial>(m, "IntMonomial", py::module_local())
+      .def(py::init<PyIntMonomial &>())
+      .def(py::init<MlirIntMonomial>())
+      .def(py::init<int64_t, uint64_t>())
+      .def_property_readonly("coefficient", &PyIntMonomial::getCoefficient)
+      .def_property_readonly("exponent", &PyIntMonomial::getExponent)
+      .def("__str__", [](PyIntMonomial &self) {
+        return std::string("<")
+            .append(std::to_string(self.getCoefficient()))
+            .append(", ")
+            .append(std::to_string(self.getExponent()))
+            .append(">");
+      });
 }

mlir/lib/CAPI/Dialect/Polynomial.cpp

@@ -14,3 +14,19 @@ using namespace mlir;
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(polynomial, polynomial,
                                       polynomial::PolynomialDialect)
+
+MlirIntMonomial mlirPolynomialGetIntMonomial(int64_t coeff, uint64_t expo) {
+  return wrap(new mlir::polynomial::IntMonomial(coeff, expo));
+}
+
+int64_t mlirPolynomialIntMonomialGetCoefficient(MlirIntMonomial intMonomial) {
+  return unwrap(intMonomial)
+      ->getCoefficient()
+      .getLimitedValue(/*Limit = UINT64_MAX*/);
+}
+
+uint64_t mlirPolynomialIntMonomialGetExponent(MlirIntMonomial intMonomial) {
+  return unwrap(intMonomial)
+      ->getExponent(/*Limit = UINT64_MAX*/)
+      .getLimitedValue();
+}

mlir/test/python/dialects/polynomial.py

@@ -18,7 +18,10 @@ def constructAndPrintInModule(f):
 @constructAndPrintInModule
 def test_smoke():
     value = Attribute.parse("#polynomial.float_polynomial<0.5 + 1.3e06 x**2>")
-    output = Type.parse("!polynomial.polynomial<ring=<coefficientType=f32>>")
-    res = polynomial.constant(output, value)
+    res = polynomial.constant(value)
     # CHECK: polynomial.constant float<0.5 + 1.3E+6x**2> : <ring = <coefficientType = f32>>
     print(res)
+
+    int_poly = polynomial.IntMonomial(1, 10)
+    # CHECK: <1, 10>
+    print(int_poly)