[mlir][polynomial] python bindings #93109

makslevental · 2024-05-22T23:43:08Z

Python bindings for the polynomial dialect. TODO: bind some types and attrs.

llvmbot · 2024-05-22T23:43:39Z

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

Python bindings for the polynomial dialect. TODO: bind some types and attrs.

Full diff: https://github.com/llvm/llvm-project/pull/93109.diff

8 Files Affected:

(added) mlir/include/mlir-c/Dialect/Polynomial.h (+25)
(added) mlir/lib/Bindings/Python/DialectPolynomial.cpp (+23)
(modified) mlir/lib/CAPI/Dialect/CMakeLists.txt (+10)
(added) mlir/lib/CAPI/Dialect/Polynomial.cpp (+17)
(modified) mlir/python/CMakeLists.txt (+23-2)
(added) mlir/python/mlir/dialects/Polynomial.td (+14)
(added) mlir/python/mlir/dialects/polynomial.py (+6)
(added) mlir/test/python/dialects/polynomial.py (+24)

diff --git a/mlir/include/mlir-c/Dialect/Polynomial.h b/mlir/include/mlir-c/Dialect/Polynomial.h
new file mode 100644
index 0000000000000..c34032607484c
--- /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..ee5275cb536a3
--- /dev/null
+++ b/mlir/lib/Bindings/Python/DialectPolynomial.cpp
@@ -0,0 +1,23 @@
+//===- DialectQuant.cpp - 'quant' 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
+@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)

github-actions · 2024-05-22T23:45:53Z

✅ With the latest revision this PR passed the C/C++ code formatter.

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

mlir/lib/CAPI/Dialect/Polynomial.cpp

mlir/python/mlir/dialects/polynomial.py

mlir/python/mlir/dialects/Polynomial.td

mlir/test/python/dialects/polynomial.py

j2kun

LGTM, but perhaps you can summarize what additional steps you took after we got off stream?

makslevental · 2024-05-23T02:10:22Z

LGTM, but perhaps you can summarize what additional steps you took after we got off stream?

yea will do - good thing all the gnashing of teeth is on video so I can pinpoint with precise accuracy the most painful parts lol.

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

ftynse · 2024-05-30T13:57:13Z

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

+#include "mlir/CAPI/Wrap.h"
+#include "mlir/Dialect/Polynomial/IR/Polynomial.h"


We shouldn't include headers from include/mlir inside include/mlir-c.

ftynse · 2024-05-30T13:57:52Z

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

+#define DEFINE_C_API_STRUCT(name, storage)                                     \
+  struct name {                                                                \
+    storage *ptr;                                                              \
+  };                                                                           \
+  typedef struct name name


Including mlir-c/IR.h should bring this macro. Normally we don't have it redefined in dialects.

We define and undef in all usages at the moment. Given the complexity it would seem better to avoid relying on macro def across headers.

ftynse

header layering looks wrong

jpienaar · 2024-06-01T13:36:22Z

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

+#define DEFINE_C_API_STRUCT(name, storage)                                     \
+  struct name {                                                                \
+    storage *ptr;                                                              \
+  };                                                                           \
+  typedef struct name name


We define and undef in all usages at the moment. Given the complexity it would seem better to avoid relying on macro def across headers.

jpienaar · 2024-06-01T13:36:38Z

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

+
+#undef DEFINE_C_API_STRUCT
+
+DEFINE_C_API_PTR_METHODS(MlirIntMonomial, mlir::polynomial::IntMonomial);


Don't think this will compile with C compiler either

jpienaar · 2024-06-01T13:39:30Z

mlir/test/python/dialects/polynomial.py

+# CHECK-LABEL: TEST: test_smoke
+@constructAndPrintInModule
+def test_smoke():
+    value = Attribute.parse("#polynomial.float_polynomial<0.5 + 1.3e06 x**2>")


I was thinking this is the one where you'd rely on the builders rather than the generic parse (e.g., polynomial.FloatPolynomail) so that you'd be able to have a better error experience Python side/make errors less prone.

This was just a first pass/attempt/"just make it pass" impl. But what would a builder look like here? Since 0.5 + 1.3e06 x**2 is an arbitrary polynomial..?

jpienaar · 2024-06-01T13:43:34Z

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

+                                                                uint64_t expo);
+
+MLIR_CAPI_EXPORTED int64_t
+mlirPolynomialIntMonomialGetCoefficient(MlirIntMonomial intMonomial);


Is this purely for Python usage? (getting the feeling its making C shims for feasible C++ accessors)

I mean... yes? Is there a way around adding all of these? Actually this is exactly why I paused dev on this PR (because it immediately started to look like reproducing the entire class hierarchy in C and then pybind).

makslevental · 2024-06-01T16:51:35Z

FYI I've changed to draft because I'm still iterating (so the oddities @ftynse and @jpienaar pointed out were known).

makslevental requested review from ftynse and stellaraccident as code owners May 22, 2024 23:43

makslevental requested a review from j2kun May 22, 2024 23:43

llvmbot added mlir:python MLIR Python bindings mlir labels May 22, 2024

makslevental force-pushed the poly_pybind branch 2 times, most recently from 6356028 to 104bfc8 Compare May 23, 2024 00:26

j2kun reviewed May 23, 2024

View reviewed changes

mlir/include/mlir-c/Dialect/Polynomial.h Outdated Show resolved Hide resolved

j2kun reviewed May 23, 2024

View reviewed changes

mlir/lib/CAPI/Dialect/Polynomial.cpp Outdated Show resolved Hide resolved

j2kun reviewed May 23, 2024

View reviewed changes

mlir/python/mlir/dialects/polynomial.py Outdated Show resolved Hide resolved

j2kun reviewed May 23, 2024

View reviewed changes

mlir/python/mlir/dialects/Polynomial.td Outdated Show resolved Hide resolved

j2kun reviewed May 23, 2024

View reviewed changes

mlir/test/python/dialects/polynomial.py Outdated Show resolved Hide resolved

j2kun approved these changes May 23, 2024

View reviewed changes

makslevental force-pushed the poly_pybind branch from 8f52769 to 0c77928 Compare May 24, 2024 20:16

makslevental and others added 7 commits May 28, 2024 19:36

[mlir][polynomial] python bindings

47e4668

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

db010de

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

Update mlir/lib/CAPI/Dialect/Polynomial.cpp

3437b4c

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

Update mlir/python/mlir/dialects/polynomial.py

d72e361

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

Update mlir/python/mlir/dialects/Polynomial.td

70a2a34

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

Update mlir/test/python/dialects/polynomial.py

0bdb03f

Co-authored-by: Jeremy Kun <kun.jeremy@gmail.com>

bind intmonomial (in order to later bind *polynomial)

1a82557

makslevental force-pushed the poly_pybind branch from 3e758c1 to 1a82557 Compare May 29, 2024 01:03

ftynse reviewed May 31, 2024

View reviewed changes

ftynse requested changes May 31, 2024

View reviewed changes

jpienaar reviewed Jun 1, 2024

View reviewed changes

makslevental marked this pull request as draft June 1, 2024 16:49

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][polynomial] python bindings #93109

[mlir][polynomial] python bindings #93109

makslevental commented May 22, 2024

llvmbot commented May 22, 2024

github-actions bot commented May 22, 2024 •

edited

Loading

j2kun left a comment

makslevental commented May 23, 2024

ftynse May 30, 2024

ftynse May 30, 2024

jpienaar Jun 1, 2024

ftynse left a comment

jpienaar Jun 1, 2024

jpienaar Jun 1, 2024

jpienaar Jun 1, 2024

makslevental Jun 1, 2024

jpienaar Jun 1, 2024

makslevental Jun 1, 2024

makslevental commented Jun 1, 2024

		#include "mlir/CAPI/Wrap.h"
		#include "mlir/Dialect/Polynomial/IR/Polynomial.h"


		#undef DEFINE_C_API_STRUCT

		DEFINE_C_API_PTR_METHODS(MlirIntMonomial, mlir::polynomial::IntMonomial);

[mlir][polynomial] python bindings #93109

Are you sure you want to change the base?

[mlir][polynomial] python bindings #93109

Conversation

makslevental commented May 22, 2024

llvmbot commented May 22, 2024

github-actions bot commented May 22, 2024 • edited Loading

j2kun left a comment

Choose a reason for hiding this comment

makslevental commented May 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ftynse left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

makslevental commented Jun 1, 2024

github-actions bot commented May 22, 2024 •

edited

Loading