[Mlir-commits] [mlir] [mlir][polynomial] python bindings (PR #93109)

Maksim Levental llvmlistbot at llvm.org
Tue May 28 17:36:20 PDT 2024


https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/93109

>From af85c8ceec4f4732755e3ca4a498f28829dd1deb Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 18:41:16 -0500
Subject: [PATCH 1/7] [mlir][polynomial] python bindings

---
 mlir/include/mlir-c/Dialect/Polynomial.h      | 25 +++++++++++++++++++
 .../lib/Bindings/Python/DialectPolynomial.cpp | 23 +++++++++++++++++
 mlir/lib/CAPI/Dialect/CMakeLists.txt          | 10 ++++++++
 mlir/lib/CAPI/Dialect/Polynomial.cpp          | 17 +++++++++++++
 mlir/python/CMakeLists.txt                    | 25 +++++++++++++++++--
 mlir/python/mlir/dialects/Polynomial.td       | 14 +++++++++++
 mlir/python/mlir/dialects/polynomial.py       |  6 +++++
 mlir/test/python/dialects/polynomial.py       | 24 ++++++++++++++++++
 8 files changed, 142 insertions(+), 2 deletions(-)
 create mode 100644 mlir/include/mlir-c/Dialect/Polynomial.h
 create mode 100644 mlir/lib/Bindings/Python/DialectPolynomial.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Polynomial.cpp
 create mode 100644 mlir/python/mlir/dialects/Polynomial.td
 create mode 100644 mlir/python/mlir/dialects/polynomial.py
 create mode 100644 mlir/test/python/dialects/polynomial.py

diff --git a/mlir/include/mlir-c/Dialect/Polynomial.h b/mlir/include/mlir-c/Dialect/Polynomial.h
new file mode 100644
index 0000000000000..6381c0d9336f4
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Polynomial.h
@@ -0,0 +1,25 @@
+//===-- mlir-c/Dialect/Polynomial.h - C API for LLVM --------------*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_POLYNOMIAL_H
+#define MLIR_C_DIALECT_POLYNOMIAL_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Polynomial, polynomial);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_POLYNOMIAL_H
diff --git a/mlir/lib/Bindings/Python/DialectPolynomial.cpp b/mlir/lib/Bindings/Python/DialectPolynomial.cpp
new file mode 100644
index 0000000000000..66863aa28375c
--- /dev/null
+++ b/mlir/lib/Bindings/Python/DialectPolynomial.cpp
@@ -0,0 +1,23 @@
+//===- DialectPolynomial.cpp - 'polynomial' dialect submodule -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Polynomial.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+#include <pybind11/pybind11.h>
+#include <vector>
+
+namespace py = pybind11;
+using namespace llvm;
+using namespace mlir;
+using namespace mlir::python::adaptors;
+
+PYBIND11_MODULE(_mlirDialectsPolynomial, m) {
+  m.doc() = "MLIR Polynomial dialect";
+}
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index 4e141b60ff8cc..d2e330da33763 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -225,6 +225,16 @@ add_mlir_upstream_c_api_library(MLIRCAPIQuant
   MLIRQuantDialect
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIPolynomial
+  Polynomial.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRPolynomialDialect
+)
+
+
 add_mlir_upstream_c_api_library(MLIRCAPIOpenMP
   OpenMP.cpp
 
diff --git a/mlir/lib/CAPI/Dialect/Polynomial.cpp b/mlir/lib/CAPI/Dialect/Polynomial.cpp
new file mode 100644
index 0000000000000..330fd05bb0521
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Polynomial.cpp
@@ -0,0 +1,17 @@
+//===- Polynomial.cpp - C Interface for Polynomial dialect
+//--------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Polynomial.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialDialect.h"
+
+using namespace mlir;
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(polynomial, polynomial,
+                                      polynomial::PolynomialDialect)
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index d8f2d1989fdea..c3c3890e5debb 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -162,14 +162,22 @@ declare_mlir_dialect_python_bindings(
   GEN_ENUM_BINDINGS)
 
 declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformPDLExtensionOps.td
   SOURCES
     dialects/transform/pdl.py
   DIALECT_NAME transform
   EXTENSION_NAME transform_pdl_extension)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/Polynomial.td
+  SOURCES
+    dialects/polynomial.py
+  DIALECT_NAME polynomial)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
@@ -537,6 +545,19 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
     MLIRCAPIQuant
 )
 
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Polynomial.Pybind
+  MODULE_NAME _mlirDialectsPolynomial
+  ADD_TO_PARENT MLIRPythonSources.Dialects.polynomial
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectPolynomial.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIPolynomial
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
   MODULE_NAME _mlirDialectsNVGPU
   ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
diff --git a/mlir/python/mlir/dialects/Polynomial.td b/mlir/python/mlir/dialects/Polynomial.td
new file mode 100644
index 0000000000000..9567ea6e3c023
--- /dev/null
+++ b/mlir/python/mlir/dialects/Polynomial.td
@@ -0,0 +1,14 @@
+//===-- PolynomialOps.td - Entry point for PolynomialOps bind ------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_POLYNOMIAL_OPS
+#define PYTHON_BINDINGS_POLYNOMIAL_OPS
+
+include "mlir/Dialect/Polynomial/IR/Polynomial.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/polynomial.py b/mlir/python/mlir/dialects/polynomial.py
new file mode 100644
index 0000000000000..dd9bee4aa3bd9
--- /dev/null
+++ b/mlir/python/mlir/dialects/polynomial.py
@@ -0,0 +1,6 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._mlir_libs._mlirDialectsPolynomial import *
+from ._polynomial_ops_gen import *
\ No newline at end of file
diff --git a/mlir/test/python/dialects/polynomial.py b/mlir/test/python/dialects/polynomial.py
new file mode 100644
index 0000000000000..a405518671d6a
--- /dev/null
+++ b/mlir/test/python/dialects/polynomial.py
@@ -0,0 +1,24 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import polynomial
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: TEST: test_smoke
+ at 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)
+    # CHECK: polynomial.constant {value = #polynomial.float_polynomial<0.5 + 1.3E+6x**2>} : <ring = <coefficientType = f32>>
+    print(res)

>From fd6c12353721ce6dc5c4a5c33c2902a24353ec2b Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 21:00:04 -0500
Subject: [PATCH 2/7] Update mlir/include/mlir-c/Dialect/Polynomial.h

Co-authored-by: Jeremy Kun <kun.jeremy at gmail.com>
---
 mlir/include/mlir-c/Dialect/Polynomial.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir-c/Dialect/Polynomial.h b/mlir/include/mlir-c/Dialect/Polynomial.h
index 6381c0d9336f4..52a3df5b475de 100644
--- a/mlir/include/mlir-c/Dialect/Polynomial.h
+++ b/mlir/include/mlir-c/Dialect/Polynomial.h
@@ -1,4 +1,4 @@
-//===-- mlir-c/Dialect/Polynomial.h - C API for LLVM --------------*- C -*-===//
+//===-- mlir-c/Dialect/Polynomial.h - C API for Polynomial --------*- C -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM
 // Exceptions.

>From 923b95249771d88af68650d965890ea2446c81da Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 21:00:14 -0500
Subject: [PATCH 3/7] Update mlir/lib/CAPI/Dialect/Polynomial.cpp

Co-authored-by: Jeremy Kun <kun.jeremy at gmail.com>
---
 mlir/lib/CAPI/Dialect/Polynomial.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/lib/CAPI/Dialect/Polynomial.cpp b/mlir/lib/CAPI/Dialect/Polynomial.cpp
index 330fd05bb0521..2fda7b0036207 100644
--- a/mlir/lib/CAPI/Dialect/Polynomial.cpp
+++ b/mlir/lib/CAPI/Dialect/Polynomial.cpp
@@ -1,5 +1,4 @@
-//===- Polynomial.cpp - C Interface for Polynomial dialect
-//--------------------------===//
+//===- Polynomial.cpp - C Interface for Polynomial dialect ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

>From 243d468531a88437a01024aa888d198e8cae3204 Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 21:00:21 -0500
Subject: [PATCH 4/7] Update mlir/python/mlir/dialects/polynomial.py

Co-authored-by: Jeremy Kun <kun.jeremy at gmail.com>
---
 mlir/python/mlir/dialects/polynomial.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/python/mlir/dialects/polynomial.py b/mlir/python/mlir/dialects/polynomial.py
index dd9bee4aa3bd9..5bd8084d5e5ef 100644
--- a/mlir/python/mlir/dialects/polynomial.py
+++ b/mlir/python/mlir/dialects/polynomial.py
@@ -3,4 +3,4 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._mlir_libs._mlirDialectsPolynomial import *
-from ._polynomial_ops_gen import *
\ No newline at end of file
+from ._polynomial_ops_gen import *

>From bc2f2b2c5539831425c59b2e4d3dba31b451738a Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 21:00:27 -0500
Subject: [PATCH 5/7] Update mlir/python/mlir/dialects/Polynomial.td

Co-authored-by: Jeremy Kun <kun.jeremy at gmail.com>
---
 mlir/python/mlir/dialects/Polynomial.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/python/mlir/dialects/Polynomial.td b/mlir/python/mlir/dialects/Polynomial.td
index 9567ea6e3c023..b38cb05e47447 100644
--- a/mlir/python/mlir/dialects/Polynomial.td
+++ b/mlir/python/mlir/dialects/Polynomial.td
@@ -1,4 +1,4 @@
-//===-- PolynomialOps.td - Entry point for PolynomialOps bind ------------*- tablegen -*-===//
+//===-- Polynomial.td - Entry point for Polynomial bind ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

>From 0c7792879e0a5a1aca739e88e3267c515323844d Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Wed, 22 May 2024 21:00:35 -0500
Subject: [PATCH 6/7] Update mlir/test/python/dialects/polynomial.py

Co-authored-by: Jeremy Kun <kun.jeremy at gmail.com>
---
 mlir/test/python/dialects/polynomial.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/python/dialects/polynomial.py b/mlir/test/python/dialects/polynomial.py
index a405518671d6a..0efec166fb66c 100644
--- a/mlir/test/python/dialects/polynomial.py
+++ b/mlir/test/python/dialects/polynomial.py
@@ -20,5 +20,5 @@ 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)
-    # CHECK: polynomial.constant {value = #polynomial.float_polynomial<0.5 + 1.3E+6x**2>} : <ring = <coefficientType = f32>>
+    # CHECK: polynomial.constant float<0.5 + 1.3E+6x**2> : <ring = <coefficientType = f32>>
     print(res)

>From 3e758c1f8638e51f5cdb851be3dade2aac639c40 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 28 May 2024 19:36:07 -0500
Subject: [PATCH 7/7] bind intmonomial (in order to later bind *polynomial)

---
 mlir/include/mlir-c/Dialect/Polynomial.h      | 44 ++++++++++++++
 .../lib/Bindings/Python/DialectPolynomial.cpp | 59 +++++++++++++++++++
 mlir/lib/CAPI/Dialect/Polynomial.cpp          | 16 +++++
 mlir/test/python/dialects/polynomial.py       |  4 ++
 4 files changed, 123 insertions(+)

diff --git a/mlir/include/mlir-c/Dialect/Polynomial.h b/mlir/include/mlir-c/Dialect/Polynomial.h
index 52a3df5b475de..40b507c554a47 100644
--- a/mlir/include/mlir-c/Dialect/Polynomial.h
+++ b/mlir/include/mlir-c/Dialect/Polynomial.h
@@ -10,14 +10,58 @@
 #ifndef MLIR_C_DIALECT_POLYNOMIAL_H
 #define MLIR_C_DIALECT_POLYNOMIAL_H
 
+#include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/IR.h"
+#include "mlir/CAPI/Wrap.h"
+#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+#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_DECLARE_CAPI_DIALECT_REGISTRATION(Polynomial, polynomial);
 
+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);
+
+#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, NULL);
+}
+
+static inline bool mlirIntMonomialIsNull(MlirIntMonomial intMonomial) {
+  return !intMonomial.ptr;
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/lib/Bindings/Python/DialectPolynomial.cpp b/mlir/lib/Bindings/Python/DialectPolynomial.cpp
index 66863aa28375c..6dab082a796eb 100644
--- a/mlir/lib/Bindings/Python/DialectPolynomial.cpp
+++ b/mlir/lib/Bindings/Python/DialectPolynomial.cpp
@@ -18,6 +18,65 @@ 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;
+};
+
+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(">");
+      });
 }
diff --git a/mlir/lib/CAPI/Dialect/Polynomial.cpp b/mlir/lib/CAPI/Dialect/Polynomial.cpp
index 2fda7b0036207..18e974c0a2efe 100644
--- a/mlir/lib/CAPI/Dialect/Polynomial.cpp
+++ b/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();
+}
\ No newline at end of file
diff --git a/mlir/test/python/dialects/polynomial.py b/mlir/test/python/dialects/polynomial.py
index 0efec166fb66c..ea31271ba5ece 100644
--- a/mlir/test/python/dialects/polynomial.py
+++ b/mlir/test/python/dialects/polynomial.py
@@ -22,3 +22,7 @@ def test_smoke():
     res = polynomial.constant(output, 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)



More information about the Mlir-commits mailing list