diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt --- a/mlir/lib/Bindings/Python/CMakeLists.txt +++ b/mlir/lib/Bindings/Python/CMakeLists.txt @@ -17,8 +17,11 @@ # The actual extension library produces a shared-object or DLL and has # sources that must be compiled in accordance with pybind11 needs (RTTI and # exceptions). +# TODO: Link the libraries separately once a helper function is available +# to more generically add a pybind11 compliant library. add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE} MainModule.cpp + IRModules.cpp ) target_include_directories(MLIRBindingsPythonExtension PRIVATE @@ -66,5 +69,7 @@ target_link_libraries(MLIRBindingsPythonExtension PRIVATE MLIRIR + MLIRCAPIIR + MLIRCAPIRegistration ${PYEXT_LIBADD} ) diff --git a/mlir/lib/Bindings/Python/IRModules.h b/mlir/lib/Bindings/Python/IRModules.h new file mode 100644 --- /dev/null +++ b/mlir/lib/Bindings/Python/IRModules.h @@ -0,0 +1,53 @@ +//===- IRModules.h - IR Submodules of pybind module -----------------------===// +// +// 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_BINDINGS_PYTHON_IRMODULES_H +#define MLIR_BINDINGS_PYTHON_IRMODULES_H + +#include + +#include "mlir-c/IR.h" + +namespace py = pybind11; + +class PyMlirContext; +class PyMlirModule; + +/// Wrapper around MlirContext. +class PyMlirContext { +public: + PyMlirContext() { context = mlirContextCreate(); } + ~PyMlirContext() { mlirContextDestroy(context); } + /// Parses the module from asm. + PyMlirModule parse(const std::string &module); + + MlirContext context; +}; + +/// Wrapper around MlirModule. +class PyMlirModule { +public: + PyMlirModule(MlirModule module) : module(module) {} + PyMlirModule(PyMlirModule &) = delete; + PyMlirModule(PyMlirModule &&other) { + module = other.module; + other.module.ptr = nullptr; + } + ~PyMlirModule() { + if (module.ptr) + mlirModuleDestroy(module); + } + /// Dumps the module. + void dump(); + + MlirModule module; +}; + +void populateIRSubmodule(py::module &m); + +#endif // MLIR_BINDINGS_PYTHON_IRMODULES_H diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Bindings/Python/IRModules.cpp @@ -0,0 +1,36 @@ +//===- IRModules.cpp - IR Submodules of pybind module ---------------------===// +// +// 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 "IRModules.h" + +//------------------------------------------------------------------------------ +// Context Wrapper Class. +//------------------------------------------------------------------------------ + +PyMlirModule PyMlirContext::parse(const std::string &module) { + auto moduleRef = mlirModuleCreateParse(context, module.c_str()); + return PyMlirModule(moduleRef); +} + +//------------------------------------------------------------------------------ +// Module Wrapper Class. +//------------------------------------------------------------------------------ + +void PyMlirModule::dump() { mlirOperationDump(mlirModuleGetOperation(module)); } + +//------------------------------------------------------------------------------ +// Populates the pybind11 IR submodule. +//------------------------------------------------------------------------------ + +void populateIRSubmodule(py::module &m) { + py::class_(m, "MlirContext") + .def(py::init<>()) + .def("parse", &PyMlirContext::parse, py::keep_alive<0, 1>()); + + py::class_(m, "MlirModule").def("dump", &PyMlirModule::dump); +} diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp --- a/mlir/lib/Bindings/Python/MainModule.cpp +++ b/mlir/lib/Bindings/Python/MainModule.cpp @@ -10,6 +10,7 @@ #include +#include "IRModules.h" #include "mlir/IR/MLIRContext.h" using namespace mlir; @@ -24,4 +25,8 @@ return std::make_tuple(std::string("From the native module"), context.isMultithreadingEnabled()); }); + + // Define and populate IR submodule. + auto irModule = m.def_submodule("ir", "MLIR IR Bindings"); + populateIRSubmodule(irModule); } diff --git a/mlir/test/Bindings/Python/ir_test.py b/mlir/test/Bindings/Python/ir_test.py new file mode 100644 --- /dev/null +++ b/mlir/test/Bindings/Python/ir_test.py @@ -0,0 +1,14 @@ +# RUN: %PYTHON %s | FileCheck %s + +import mlir + +TEST_MLIR_ASM = r""" +module { +} +""" + +ctx = mlir.ir.MlirContext() +module = ctx.parse(TEST_MLIR_ASM) +module.dump() +print(bool(module)) +# CHECK: True