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 @@ -19,6 +19,7 @@ # exceptions). add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE} MainModule.cpp + IRModules.cpp ) target_include_directories(MLIRBindingsPythonExtension PRIVATE @@ -66,5 +67,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,44 @@ +//===- 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_PYTHON_IR_MODULES_H +#define MLIR_PYTHON_IR_MODULES_H + +#include +namespace py = pybind11; + +#include "mlir-c/IR.h" + +class PyMlirContext; +class PyMlirModule; + +/// Wrapper around MlirModule. +class PyMlirModule { +public: + 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_PYTHON_IR_MODULES_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,46 @@ +//===- 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. */ +/*============================================================================*/ +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 new PyMlirModule(moduleRef); +} + +/*============================================================================*/ +/* Module Wrapper Class. */ +/*============================================================================*/ +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