[mlir][polynomial] python bindings

Author: makslevental

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

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";
+}

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
 

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)

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

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

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 *

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)