From 47e4668c32281d6c2de81e61e4e8272f4641d6f8 Mon Sep 17 00:00:00 2001 From: makslevental 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 +#include + +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 +@constructAndPrintInModule +def test_smoke(): + value = Attribute.parse("#polynomial.float_polynomial<0.5 + 1.3e06 x**2>") + output = Type.parse("!polynomial.polynomial>") + res = polynomial.constant(output, value) + # CHECK: polynomial.constant {value = #polynomial.float_polynomial<0.5 + 1.3E+6x**2>} : > + print(res) From db010de5171ec294888d051843f16e1cd8485a94 Mon Sep 17 00:00:00 2001 From: Maksim Levental 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 --- 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 3437b4cce709f8542603463d7ae99d9b762aa0de Mon Sep 17 00:00:00 2001 From: Maksim Levental 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 --- 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 d72e361d598a3ad73e78fcf29f78561c5d987173 Mon Sep 17 00:00:00 2001 From: Maksim Levental 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 --- 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 70a2a34f0aaf847c151ed3d0cbefc6d519d278eb Mon Sep 17 00:00:00 2001 From: Maksim Levental 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 --- 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 0bdb03f45c2c3b50c441263455a23c06a45756bc Mon Sep 17 00:00:00 2001 From: Maksim Levental 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 --- 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>") res = polynomial.constant(output, value) - # CHECK: polynomial.constant {value = #polynomial.float_polynomial<0.5 + 1.3E+6x**2>} : > + # CHECK: polynomial.constant float<0.5 + 1.3E+6x**2> : > print(res) From 1a8255761dbac6ed4d118d0a7ad6e2994df46673 Mon Sep 17 00:00:00 2001 From: makslevental 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 | 24 +++++- .../lib/Bindings/Python/DialectPolynomial.cpp | 79 +++++++++++++++++++ mlir/lib/CAPI/Dialect/Polynomial.cpp | 16 ++++ mlir/test/python/dialects/polynomial.py | 7 +- 4 files changed, 123 insertions(+), 3 deletions(-) diff --git a/mlir/include/mlir-c/Dialect/Polynomial.h b/mlir/include/mlir-c/Dialect/Polynomial.h index 52a3df5b475de..31fb817cf534c 100644 --- a/mlir/include/mlir-c/Dialect/Polynomial.h +++ b/mlir/include/mlir-c/Dialect/Polynomial.h @@ -10,7 +10,8 @@ #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" { @@ -18,6 +19,27 @@ extern "C" { 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 diff --git a/mlir/lib/Bindings/Python/DialectPolynomial.cpp b/mlir/lib/Bindings/Python/DialectPolynomial.cpp index 66863aa28375c..f0832d400e8c5 100644 --- a/mlir/lib/Bindings/Python/DialectPolynomial.cpp +++ b/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 { + 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(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_(m, "IntMonomial", py::module_local()) + .def(py::init()) + .def(py::init()) + .def(py::init()) + .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..d7e308a6c2b54 100644 --- a/mlir/test/python/dialects/polynomial.py +++ b/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>") - res = polynomial.constant(output, value) + res = polynomial.constant(value) # CHECK: polynomial.constant float<0.5 + 1.3E+6x**2> : > print(res) + + int_poly = polynomial.IntMonomial(1, 10) + # CHECK: <1, 10> + print(int_poly)