Add _npcomp.backend.iree module.

* Populates with builders for the various path pipelines and translator.

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

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
+)

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_<Blob>(m, "Blob", py::buffer_protocol())
+      .def_buffer([](Blob &self) -> py::buffer_info {
+        return py::buffer_info(
+            static_cast<void *>(&self.contents.front()), // Pointer to buffer
+            sizeof(uint8_t),                             // Size of one scalar
+            py::format_descriptor<uint8_t>::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<std::string> 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<std::string>(),
+        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()));
+  });
+}

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.

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

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 "")

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