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,10 @@ # 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: linking the libraries separately. add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE} MainModule.cpp + IRModules.cpp ) target_include_directories(MLIRBindingsPythonExtension PRIVATE @@ -66,5 +68,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,45 @@ +//===- 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 +namespace py = pybind11; + +#include "mlir-c/IR.h" + +class PyMlirContext; +class PyMlirModule; + +/// Wrapper around MlirModule. +class PyMlirModule { +public: + PyMlirModule(PyMlirModule &&pyModule) noexcept; + PyMlirModule(MlirModule moduleRef); + ~PyMlirModule(); + // Dump the module. + void dump(); + + MlirModule module; +}; + +/// Wrapper around MlirContext. +class PyMlirContext { +public: + PyMlirContext(); + ~PyMlirContext(); + // Parses the module from asm. + PyMlirModule parse(const std::string &module); + + MlirContext context; +}; + +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,53 @@ +//===- 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" +#include "mlir/IR/Module.h" + +/*============================================================================*/ +/* Context Wrapper Class. */ +/*============================================================================*/ +PyMlirContext::PyMlirContext() { context = mlirContextCreate(); } +PyMlirContext::~PyMlirContext() { mlirContextDestroy(context); } +// Parses the module from asm. +PyMlirModule PyMlirContext::parse(const std::string &module) { + auto moduleRef = mlirModuleCreateParse(context, module.c_str()); + return PyMlirModule(moduleRef); +} + +/*============================================================================*/ +/* Module Wrapper Class. */ +/*============================================================================*/ +// Move constructor of PyMlirModule. +PyMlirModule::PyMlirModule(PyMlirModule &&pyModule) noexcept { + // Create a deep copy to avoid double free the object. + auto ptr = mlir::ModuleOp::getFromOpaquePointer(pyModule.module.ptr).clone(); + this->module = {ptr}; +} +PyMlirModule::PyMlirModule(MlirModule moduleRef) : module(moduleRef) {} +PyMlirModule::~PyMlirModule() { mlirModuleDestroy(module); } +// Get operation and dump. +void PyMlirModule::dump() { mlirOperationDump(mlirModuleGetOperation(module)); } + +// Implementation of IR submodule. +void populateIRSubmodule(py::module &m) { + /*==========================================================================*/ + /* Context bindings. */ + /*==========================================================================*/ + py::class_(m, "MlirContext") + .def(py::init<>()) + // Parses the module from asm. + .def("parse", &PyMlirContext::parse, py::keep_alive<0, 1>()); + + /*==========================================================================*/ + /* Module bindings. */ + /*==========================================================================*/ + py::class_(m, "MlirModule") + // Dump the module. + .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