diff --git a/include/npcomp/Backend/IREE/PythonModule.h b/include/npcomp/Backend/IREE/PythonModule.h new file mode 100644 index 000000000..6d8a85623 --- /dev/null +++ b/include/npcomp/Backend/IREE/PythonModule.h @@ -0,0 +1,25 @@ +//===- PythonModule.h - IREE python bindings ------------------------------===// +// +// 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 NPCOMP_BACKEND_IREE_PYTHON_MODULE_H +#define NPCOMP_BACKEND_IREE_PYTHON_MODULE_H + +#include "npcomp/Python/PybindUtils.h" + +namespace mlir { +namespace npcomp { +namespace python { + +/// Defines an "iree" module with backend support definitions. +void defineBackendIREEModule(py::module m); + +} // namespace python +} // namespace npcomp +} // namespace mlir + +#endif // NPCOMP_BACKEND_IREE_PYTHON_MODULE_H diff --git a/lib/Backend/IREE/CMakeLists.txt b/lib/Backend/IREE/CMakeLists.txt new file mode 100644 index 000000000..c8c943728 --- /dev/null +++ b/lib/Backend/IREE/CMakeLists.txt @@ -0,0 +1,22 @@ +################################################################################ +# NPCOMPBackendIREEPythonModule +################################################################################ + +set(PYBIND_SOURCES + PythonModule.cpp +) +set_source_files_properties( + ${PYBIND_SOURCES} + PROPERTIES COMPILE_FLAGS + "-frtti -fexceptions") + +add_library(NPCOMPBackendIREEPythonModule + ${PYBIND_SOURCES} +) + +target_link_libraries(NPCOMPBackendIREEPythonModule + NPCOMPPythonCommon + iree_compiler_Dialect_Flow_Transforms_Transforms + iree_compiler_Dialect_HAL_Transforms_Transforms + iree_compiler_Dialect_VM_Transforms_Transforms +) diff --git a/lib/Backend/IREE/PythonModule.cpp b/lib/Backend/IREE/PythonModule.cpp new file mode 100644 index 000000000..4c79b39d5 --- /dev/null +++ b/lib/Backend/IREE/PythonModule.cpp @@ -0,0 +1,88 @@ +//===- PythonModule.cpp - IREE python bindings ----------------------------===// +// +// 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 "npcomp/Backend/IREE/PythonModule.h" + +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h" +#include "iree/compiler/Dialect/HAL/Transforms/Passes.h" +#include "iree/compiler/Dialect/VM/Target/Bytecode/BytecodeModuleTarget.h" +#include "iree/compiler/Dialect/VM/Transforms/Passes.h" +#include "npcomp/Python/MlirIr.h" +#include "npcomp/Python/MlirPass.h" + +using namespace mlir; + +namespace { + +class Blob { +public: + Blob(std::string contents) : contents(contents) {} + std::string contents; +}; + +} // namespace + +/// Defines an "iree" module with backend support definitions. +void mlir::npcomp::python::defineBackendIREEModule(py::module m) { + py::class_(m, "Blob", py::buffer_protocol()) + .def_buffer([](Blob &self) -> py::buffer_info { + return py::buffer_info( + static_cast(&self.contents.front()), // Pointer to buffer + sizeof(uint8_t), // Size of one scalar + py::format_descriptor::value, // Python struct-style + // format + 1, // Number of dimensions + {self.contents.size()}, // Buffer dimensions + {self.contents.size()} // Strides + ); + }); + + m.def("build_flow_transform_pass_pipeline", + [](PyPassManager &pm) { + mlir::iree_compiler::IREE::Flow::buildFlowTransformPassPipeline( + pm.passManager); + }, + py::arg("pm"), + py::doc("Builds a pass pipeline for top-level Flow import")); + m.def("build_hal_transform_pass_pipeline", + [](PyPassManager &pm, std::vector targetBackends) { + mlir::iree_compiler::IREE::HAL::TargetOptions options; + if (targetBackends.empty()) { + options.targets = + mlir::iree_compiler::IREE::HAL::getRegisteredTargetBackends(); + } else { + options.targets = std::move(targetBackends); + } + iree_compiler::IREE::HAL::buildHALTransformPassPipeline( + pm.passManager, options); + }, + py::arg("pm"), py::arg("target_backends") = std::vector(), + py::doc("Builds a pass pipeline for top-level Flow import")); + m.def("build_vm_transform_pass_pipeline", + [](PyPassManager &pm) { + mlir::iree_compiler::IREE::VM::buildVMTransformPassPipeline( + pm.passManager); + }, + py::arg("pm"), py::doc("Builds the VM transformation pipeline")); + m.def("translate_to_vm_bytecode", [](PyModuleOp &module) { + // TODO: Make the options parameterizable. + mlir::iree_compiler::IREE::VM::BytecodeTargetOptions options; + std::string contents; + llvm::raw_string_ostream out(contents); + if (failed(mlir::iree_compiler::IREE::VM::translateModuleToBytecode( + module.moduleOp, options, out))) { + // TODO: Merge diagnostic captures in. + throw py::raisePyError(PyExc_RuntimeError, + "Error translating module (see stderr)"); + } + + out.flush(); + return Blob(std::move(out.str())); + }); +} diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index d0b7eeb2d..1b403dd59 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -3,6 +3,10 @@ add_subdirectory(Dialect) add_subdirectory(E2E) add_subdirectory(Python) +if(NPCOMP_ENABLE_IREE) + add_subdirectory(Backend/IREE) +endif() + ################################################################################ # Setup the initialization target. # This includes conditional dependencies based on whether features are enabled. diff --git a/lib/Python/CMakeLists.txt b/lib/Python/CMakeLists.txt index 01abd3eb2..cf4ce11d7 100644 --- a/lib/Python/CMakeLists.txt +++ b/lib/Python/CMakeLists.txt @@ -1,4 +1,8 @@ -add_library(NPCOMPPythonCommon +################################################################################ +# NPCOMPPythonCommon +################################################################################ + +set(PYBIND_SOURCES MlirInit.cpp MlirIr.cpp MlirPass.cpp @@ -6,6 +10,14 @@ add_library(NPCOMPPythonCommon PybindUtils.cpp CoreDialects.cpp ) +set_source_files_properties( + ${PYBIND_SOURCES} + PROPERTIES COMPILE_FLAGS + "-frtti -fexceptions") + +add_library(NPCOMPPythonCommon + ${PYBIND_SOURCES} +) target_link_libraries(NPCOMPPythonCommon pybind11::module diff --git a/python_native/CMakeLists.txt b/python_native/CMakeLists.txt index cfb38109f..1eb6049ff 100644 --- a/python_native/CMakeLists.txt +++ b/python_native/CMakeLists.txt @@ -15,6 +15,10 @@ set(NPCOMP_PYEXT_LINK_MODE SHARED) set(NPCOMP_PYEXT_LIBADD ${PYTHON_LIBRARIES}) +if(NPCOMP_ENABLE_IREE) + list(APPEND NPCOMP_PYEXT_LIBADD NPCOMPBackendIREEPythonModule) +endif() + # TODO(laurenzo): Add a config setting to control this. # set(NPCOMP_PYEXT_LINK_MODE MODULE) # set(NPCOMP_PYEXT_LIBADD "") diff --git a/python_native/NpcompModule.cpp b/python_native/NpcompModule.cpp index c3c4a005b..21612bfed 100644 --- a/python_native/NpcompModule.cpp +++ b/python_native/NpcompModule.cpp @@ -12,9 +12,12 @@ #include "npcomp/Python/MlirInit.h" #include "npcomp/Python/NpcompModule.h" #include "npcomp/Python/PybindUtils.h" - #include "llvm/Support/CommandLine.h" +#ifdef NPCOMP_ENABLE_IREE +#include "npcomp/Backend/IREE/PythonModule.h" +#endif + namespace mlir { namespace npcomp { namespace python { @@ -73,6 +76,15 @@ PYBIND11_MODULE(_npcomp, m) { // Outer "_npcomp" module auto npcomp_dialect = m.def_submodule("dialect", "NPComp custom dialects"); defineNpcompDialect(npcomp_dialect); + + // Optional backend modules. + auto backend_m = m.def_submodule("backend", "Backend support"); + (void)backend_m; + +#ifdef NPCOMP_ENABLE_IREE + auto iree_m = backend_m.def_submodule("iree", "IREE backend support"); + defineBackendIREEModule(iree_m); +#endif } } // namespace python