31b3041e88
This patch adds a pytorch interface to npcomp. This interface is modeled
after pytorch_xla and exposes the MLIR-based flow as a virtual device (similar
to a gpu device or the xla backend). Usage is intended to be something like:
dev = torch_mlir.mlir_device()
t0 = torch.randn((4,4), device=dev)
t1 = torch.randn((4,4), device=dev)
t2 = t0 + t1
t2_mlir = torch_mlir.get_mlir( t2 )
t2_cpu = t2.to('cpu')
In this case t2_cpu would contain the result of the computation, and t2_mlir
contains the mlir description of the computation. Note that this also
properly returns backward paths synthesized by pytorch. There are several
parts of this:
1) A tensor type (implemented by tensor.* and tensor_impl.*)
2) The device modeling (aten_mlir_bridge.*, aten_mlir_device.*, aten_mlir_type*)
3) a temporary IR (implemented by ir.cpp)
There is also a reference lowering directly from the ATen dialect to C
function calls consisting of two parts:
1) The driver that uses the IR to generate MLIR, run Passes and compile the
result using mlir::ExecutionEngine (implemented by jit.cpp and
mlir_gen.cpp)
2) A runtime library implemented by lib/aten_ops.cpp. Most of the operations
are implemented by callbacks into the torch C++ libraries.
Some aspects of this are known to be less than optimal, in particular:
1) There's some function definitions that don't live in the file corresponding
to their declaration.
2) More aspects of this (e.g. the IR) seem like they should be automatically
generated.
3) It's unclear to me how much of the 'IR' is actually necessary, or whether
MLIR could be created on the fly.
Note that this code is licensed in a way similar to pytorch, with the
intention that eventually (when npcomp reaches some maturity) it should be
pushed there. (see frontends/pytorch/LICENSE) The code is also structured
much closer to the pytorch coding style than the LLVM coding style.
|
||
|---|---|---|
| .github/workflows | ||
| build_tools | ||
| cmake/modules | ||
| docs | ||
| external | ||
| frontends | ||
| include/npcomp | ||
| lib | ||
| python | ||
| test | ||
| tools | ||
| .clang-format | ||
| .gitignore | ||
| .gitmodules | ||
| .style.yapf | ||
| CMakeLists.txt | ||
| LICENSE | ||
| README.md | ||
README.md
NPComp - MLIR based compiler toolkit for numerical python programs
This project is participating in the LLVM Incubator process: as such, it is not part of any official LLVM release. While incubation status is not necessarily a reflection of the completeness or stability of the code, it does indicate that the project is not yet endorsed as a component of LLVM.
The NPComp project aims to provide tooling for compiling numerical python programs of various forms to take advantage of MLIR+LLVM code generation and backend runtime systems.
In addition to providing a bridge to a variety of Python based numerical programming frameworks, NPComp also directly develops components for tracing and compilation of generic Python program fragments.
Framework integrations
- PyTorch -- Experimental integration for extracting programs from PyTorch.
Python language compiler tookit
At the core of NPComp are a set of dialects and python support code for tracing (define by run) numerical programs and compiling idiomatic subsets of the Python language. As another interpretation of the name, NPComp also seeks to provide compiler-backed support for Numpy APIs.
See the features doc for a semi-curated status of what is implemented in this area.
Architecture
The compiler is separated into:
- Frontend importer: Translates from various AST levels to corresponding MLIR dialects.
- Frontend compiler: MLIR passes and conversions, mostly operating on the basicpy and numpy dialects.
- Backend compiler and runtime: Some effort has been taken to make this pluggable, but right now, only the IREE Backend exists. There is in-tree work to also build a minimal reference backend directly targeting LLVM.
Repository Layout
The project is roughly split into the following areas of code:
- User-facing Python code
- C++ include and lib trees, following LLVM/MLIR conventions
- LIT testing trees:
- test: Lit/FileCheck tests covering core MLIR based infra
- test/Python/Compiler: Lit test suite that drive the compiler infra from Python
- backend_test: Lit test suites conditionally enabled for each backend
- tools: Scripts and binaries (npcomp-opt, npcomp-run-mlir, etc)
Quick start
LLVM_VERSION=10
export CC=clang-$LLVM_VERSION
export CXX=clang++-$LLVM_VERSION
export LDFLAGS=-fuse-ld=$(which ld.lld-$LLVM_VERSION)
export LLVM_SRC_DIR=/path/to/llvm-project
# Check out last known good commit.
LLVM_COMMIT="$(cat ./build_tools/llvm_version.txt)"
(cd $LLVM_SRC_DIR && git checkout $LLVM_COMMIT)
./build_tools/install_mlir.sh
./build_tools/cmake_configure.sh
# Build and run tests
# ./build_tools/test_all.sh runs all of these commands.
cd build
ninja
ninja check-npcomp
# Setup PYTHONPATH for interactive use
export PYTHONPATH="$(realpath build/python):$(realpath build/iree/bindings/python)"
Interactive Use
The cmake configuration populates symlinks in the build/python directory
mirroring the source layout. This allows edit-run without rebuilding (unless
if files are added/removed).
Configuring the PYTHONPATH as above should be sufficient to run any
interactive tooling (python3, Jupyter/Colab, etc).
Note that running the cmake_configure.sh script will also output a .env
file in the workspace folder with the correct PYTHONPATH set. This allows
tools like VSCode to work by default for debugging.
Notes:
- Python sources are symlinked to the output directory at configure time. Adding sources will require a reconfigure. Editing should not.
- It is a very common issue to have both python 2.7 (aka. "python") and python 3.x (aka. "python3") on a system at a time (and we can only hope that one day this ends). Since the native library at development time binds to a specific version, if you try to run with a different python, you will get an error about the "native" module not being found.
Compiler development
For bash users, adding the following to your .bashrc defines some aliases
that are useful during compiler development, such as shortcuts for builing
and running npcomp-opt.
source $WHERE_YOU_CHECKED_OUT_NPCOMP/tools/bash_helpers.sh