[mpact][passes] Setup MPACT python pipeline interface. (#20)
* [mpact][passes] WIPs...
* [mpact][passes] use mpact pass manager
* update readme
diff --git a/CMakeLists.txt b/CMakeLists.txt
index a9964e3..38d9f26 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -57,6 +57,18 @@
endif()
endfunction()
+list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules)
+list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake)
+list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/build_tools/cmake)
+
+include(TableGen)
+include(AddLLVM)
+include(AddMLIR)
+include(AddMLIRPython)
+
+include(MLIRDetectPythonEnv)
+mlir_configure_python_dev_packages()
+
add_subdirectory(include)
add_subdirectory(lib)
add_subdirectory(tools)
diff --git a/README.md b/README.md
index b415e3f..8704d35 100644
--- a/README.md
+++ b/README.md
@@ -38,7 +38,7 @@
Also make sure to set the Python paths as follows.
```shell
-export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/build/tools/mpact/python_packages/mpact
+export PYTHONPATH=`pwd`/build/tools/mpact/python_packages/mpact
```
### Install build requirements
diff --git a/include/mpact-c/Registration.h b/include/mpact-c/Registration.h
new file mode 100644
index 0000000..877b06c
--- /dev/null
+++ b/include/mpact-c/Registration.h
@@ -0,0 +1,27 @@
+/*===-- mpact-c/Registration.h - Registration functions -----*- C -*-===*\
+|* *|
+|* Part of the MPACT 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 MPACT_C_REGISTRATION_H
+#define MPACT_C_REGISTRATION_H
+
+#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/** Registers all passes for symbolic access with the global registry. */
+MLIR_CAPI_EXPORTED void mpactRegisterAllPasses(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MPACT_C_REGISTRATION_H
diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt
new file mode 100644
index 0000000..f45c469
--- /dev/null
+++ b/lib/CAPI/CMakeLists.txt
@@ -0,0 +1,12 @@
+add_mlir_public_c_api_library(MPACTCAPI
+ Registration.cpp
+
+ ENABLE_AGGREGATION
+
+ LINK_LIBS PUBLIC
+ MLIRIR
+ MLIRSupport
+ MPACTTransformPasses
+)
+
+mpact_target_includes(MPACTCAPI)
diff --git a/lib/CAPI/Registration.cpp b/lib/CAPI/Registration.cpp
new file mode 100644
index 0000000..51234fb
--- /dev/null
+++ b/lib/CAPI/Registration.cpp
@@ -0,0 +1,20 @@
+//===- Registration.cpp - C Interface for MLIR Registration ---------------===//
+//
+// 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
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mpact-c/Registration.h"
+
+#include "mlir/CAPI/IR.h"
+#include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Transforms/Passes.h"
+#include "mpact/Transforms/Passes.h"
+
+MLIR_CAPI_EXPORTED void mpactRegisterAllPasses() {
+ mlir::mpact::registerTransformPasses();
+}
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
index e31af32..9e7dcd1 100644
--- a/lib/CMakeLists.txt
+++ b/lib/CMakeLists.txt
@@ -1 +1,2 @@
+add_subdirectory(CAPI)
add_subdirectory(Transforms)
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index 4b3953a..12c46bd 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -2,16 +2,40 @@
# The MPACT Compiler Python Modules
#-------------------------------------------------------------------------------
+# Disables generation of "version soname" (i.e. libFoo.so.<version>), which
+# causes pure duplication as part of Python wheels.
+set(CMAKE_PLATFORM_NO_VERSIONED_SONAME ON)
+
+
# The directory at which the Python import tree begins.
set(MPACT_PYTHON_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mpact")
+# We vendor our own MLIR instance in the `mpact` namespace.
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=mpact.")
+
declare_mlir_python_sources(MPACTPythonSources)
+declare_mlir_python_sources(MPACTPythonExtensions)
declare_mlir_python_sources(MPACTPythonSources.PublicAPI
ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}"
ADD_TO_PARENT MPACTPythonSources
SOURCES
mpactbackend.py
+ )
+
+#-------------------------------------------------------------------------------
+# Extensions
+#-------------------------------------------------------------------------------
+
+declare_mlir_python_extension(MPACTPythonExtensions.Main
+ MODULE_NAME _mpact
+ ADD_TO_PARENT MPACTPythonExtensions
+ SOURCES
+ MPACTModule.cpp
+ EMBED_CAPI_LINK_LIBS
+ MPACTCAPI
+ PRIVATE_LINK_LIBS
+ LLVMSupport
)
declare_mlir_python_sources(MPACTPythonSources.SampleModels
@@ -25,12 +49,36 @@
# Python Modules
#-------------------------------------------------------------------------------
+# To compiler TorchMLIRPythonExtension (which registers torch dialect and
+# related passes).
+# TODO: is there a cleaner way?
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../externals/torch-mlir/include)
+
set(_source_components
+ MLIRPythonSources
+ MLIRPythonExtension.Core
+ MLIRPythonExtension.RegisterEverything
+
+ # We need the FxImporter from torch-mlir
+ TorchMLIRPythonSources.Importers
+ TorchMLIRPythonExtensions
+
MPACTPythonSources
+ MPACTPythonExtensions
+)
+
+add_mlir_python_common_capi_library(MPACTAggregateCAPI
+ INSTALL_COMPONENT MPACTPythonModules
+ INSTALL_DESTINATION python_packages/mpact/mpact/_mlir_libs
+ OUTPUT_DIRECTORY "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact/_mlir_libs"
+ RELATIVE_INSTALL_ROOT ".."
+ DECLARED_SOURCES ${_source_components}
)
add_mlir_python_modules(MPACTPythonModules
ROOT_PREFIX "${MPACT_PYTHON_PACKAGES_DIR}/mpact/mpact"
INSTALL_PREFIX "python_packages/mpact/mpact"
DECLARED_SOURCES ${_source_components}
+ COMMON_CAPI_LINK_LIBS
+ MPACTAggregateCAPI
)
diff --git a/python/MPACTModule.cpp b/python/MPACTModule.cpp
new file mode 100644
index 0000000..5751287
--- /dev/null
+++ b/python/MPACTModule.cpp
@@ -0,0 +1,17 @@
+//===-- MPACTModule.cpp ------------------------------------*- cpp -*-===//
+//
+// Part of the MPACT 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
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mpact-c/Registration.h"
+
+PYBIND11_MODULE(_mpact, m) {
+ mpactRegisterAllPasses();
+
+ m.doc() = "mpact main python extension";
+}
diff --git a/python/mpact/mpactbackend.py b/python/mpact/mpactbackend.py
index 6df6137..bb8cc3b 100644
--- a/python/mpact/mpactbackend.py
+++ b/python/mpact/mpactbackend.py
@@ -1,28 +1,135 @@
+# Initialize mpact python extension.
+import mpact._mlir_libs._mpact
+
+import abc
import ctypes
+from enum import Enum
+from io import StringIO
import numpy as np
import os
+import sys
+import tempfile
import torch
-from typing import Any, Callable, Optional, Tuple, Dict
+from typing import Any, Callable, Optional, Tuple, Dict, TypeVar, Union
-from torch_mlir import ir
-from torch_mlir.compiler_utils import run_pipeline_with_repro_report
-from torch_mlir.dialects import torch as torch_d
-from torch_mlir.execution_engine import *
-from torch_mlir.extras.fx_importer import FxImporter, SparsityMeta
-from torch_mlir.ir import *
-from torch_mlir.passmanager import *
-from torch_mlir.runtime import *
-
-from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import (
- LinalgOnTensorsBackend,
-)
+from mpact import ir
+from mpact.ir import Module
+from mpact.dialects import torch as torch_d
+from mpact.execution_engine import *
+from mpact.extras.fx_importer import FxImporter, SparsityMeta
+from mpact.ir import *
+from mpact.passmanager import *
+from mpact.runtime import *
# One time set up of support library and optimization level.
SUPPORT_LIB = os.getenv("SUPPORT_LIB", default=None)
SHARED_LIBS = [] if SUPPORT_LIB is None else [SUPPORT_LIB]
OPT_LEVEL = int(os.getenv("OPT_LEVEL", default=2))
+# A type shared between the result of `LinalgOnTensorsBackend.compile` and the
+# input to `LinalgOnTensorsBackend.load`. Each backend will likely have a
+# different definition of this type.
+CompiledArtifact = TypeVar("CompiledArtifact")
+
+# A wrapper around a backend-specific loaded program representation
+# that uniformly translates the `x.method(...)` interface expected of
+# Torch modules into appropriate lower-level operations.
+Invoker = TypeVar("Invoker")
+
+
+class LinalgOnTensorsBackend(abc.ABC):
+ """The interface to an linalg-on-tensors backend.
+
+ Backends are recommended to raise meaningful exceptions in case of error,
+ ideally with easy reproduction instructions.
+ """
+
+ @abc.abstractmethod
+ def compile(self, module: Module) -> CompiledArtifact:
+ """Compile the provided MLIR module into a compiled artifact.
+
+ The module adheres to the linalg-on-tensors backend contract
+ (see the VerifyLinalgOnTensorsBackendContract pass).
+
+ The compiled artifact can be any type, but must be correctly
+ interpreted by the `load` method.
+ """
+
+ @abc.abstractmethod
+ def load(self, artifact: CompiledArtifact) -> Invoker:
+ """Load the compiled artifact into a uniformly invokable form.
+
+ The compiled artifact is the result of a previous call to `compile`.
+
+ See the description of `Invoker` for the requirements on the returned
+ type.
+ """
+
+
+def get_module_name_for_debug_dump(module):
+ """Gets a name suitable for a debug dump.
+
+ The name is not guaranteed to be unique.
+ """
+ if not "torch.debug_module_name" in module.operation.attributes:
+ return "UnnammedModule"
+ return StringAttr(module.operation.attributes["torch.debug_module_name"]).value
+
+
+class MPACTCompilerError(Exception):
+ pass
+
+
+def run_pipeline_with_repro_report(
+ module, pipeline: str, description: str, enable_ir_printing: bool = False
+):
+ """Runs `pipeline` on `module`, with a nice repro report if it fails."""
+ module_name = get_module_name_for_debug_dump(module)
+ original_stderr = sys.stderr
+ try:
+ sys.stderr = StringIO()
+ asm_for_error_report = module.operation.get_asm(
+ large_elements_limit=10, enable_debug_info=True
+ )
+ # Lower module in place to make it ready for compiler backends.
+ with module.context as ctx:
+ pm = PassManager.parse(pipeline)
+ if enable_ir_printing:
+ ctx.enable_multithreading(False)
+ pm.enable_ir_printing()
+ pm.run(module.operation)
+ except Exception as e:
+ # TODO: More robust.
+ # - don't arbitrarily clutter up /tmp. When a test suite has many
+ # tests, this can be a big disk cost (also, /tmp/ is frequently a
+ # RAM fs, which increases worries about capacity).
+ # - don't have colliding filenames (hard to do without cluttering
+ # up /tmp)
+ # - if we do have have colliding filenames, writes should at least
+ # avoid being racy.
+ filename = os.path.join(tempfile.gettempdir(), module_name + ".mlir")
+ with open(filename, "w") as f:
+ f.write(asm_for_error_report)
+ debug_options = "-mlir-print-ir-after-all -mlir-disable-threading"
+ # Put something descriptive here even if description is empty.
+ description = description or f"{module_name} compile"
+
+ message = f"""\
+ {description} failed with the following diagnostics:
+ {sys.stderr.getvalue()}
+
+ python exception: {e}
+
+ The error can be reproduced with:
+ $ mpact-opt -pass-pipeline='{pipeline}' {filename}
+ Add '{debug_options}' to get the IR dump for debugging purpose.
+ """
+ trimmed_message = "\n".join([m.lstrip() for m in message.split("\n")])
+ raise MPACTCompilerError(trimmed_message) from None
+ finally:
+ sys.stderr = original_stderr
+
def assert_arg_type_is_supported(ty):
SUPPORTED = [
@@ -153,6 +260,10 @@
"func.func(linalg-generalize-named-ops)",
"func.func(linalg-fuse-elementwise-ops)",
"convert-shape-to-std",
+ # Propagate sparse encodings before sparsifier mini-pipeline.
+ # TODO: the following pass currently contains no pattern. Will be
+ # added as needed.
+ "func.func(sparse-encoding-propagation)",
# MLIR Sparsifier mini-pipeline.
"sparse-assembler{direct-out}",
"sparsification-and-bufferization",
diff --git a/test/lit.cfg.py b/test/lit.cfg.py
index 8dc4ba4..910acca 100644
--- a/test/lit.cfg.py
+++ b/test/lit.cfg.py
@@ -64,7 +64,7 @@
config.mpact_obj_root,
]
tools = [
- "torch-mlir-opt",
+ "mpact-opt",
ToolSubst("%PYTHON", config.python_executable, unresolved="ignore"),
]
@@ -74,7 +74,6 @@
"PYTHONPATH",
[
os.path.join(config.mpact_obj_root, "python_packages/mpact"),
- os.path.join(config.torch_mlir_obj_root, "python_packages/torch_mlir"),
],
append_path=True,
)
