The Torch-MLIR project aims to provide first class support from the PyTorch ecosystem to the MLIR ecosystem.
 
 
 
 
 
 
Go to file
Stella Laurenzo 504e3c4946 Fixup local ndarray<->tensor transforms to preserve shape.
* Preserving shape across the copy ops makes more thing shaped by default.
* Inference of ndarray types will now preserve the shape when specializing the dtype.
2020-07-05 17:45:45 -07:00
backend_test/iree Refactor environment.py into components. 2020-06-28 16:52:25 -07:00
build_tools Bump llvm and iree versions. 2020-07-02 11:24:05 -07:00
cmake/modules Fix some compiler option and warning levels. 2020-07-04 17:38:01 -07:00
docs Fix broken link 2020-06-29 16:49:23 -07:00
include/npcomp Allow the ndarray type to carry a shape. 2020-07-05 17:34:03 -07:00
lib Fixup local ndarray<->tensor transforms to preserve shape. 2020-07-05 17:45:45 -07:00
pytest Fixup local ndarray<->tensor transforms to preserve shape. 2020-07-05 17:45:45 -07:00
python Some fixes to get npcomp building and passing on windows. 2020-07-01 21:28:04 -07:00
python_native Fix some compiler option and warning levels. 2020-07-04 17:38:01 -07:00
runtime Add bare-bones npcomp-run-mlir. 2020-05-28 18:37:11 -07:00
test Allow the ndarray type to carry a shape. 2020-07-05 17:34:03 -07:00
tools Rework README, add docs, and move some scripts to better locations accordingly. 2020-06-20 11:56:13 -07:00
.clang-format Add stub numpy dialect. 2020-04-26 17:20:58 -07:00
.gitignore Add MLIRContext.dense_elements_attr to create an attribute from a python buffer/array. 2020-05-08 17:36:07 -07:00
.style.yapf Introduce a Target class and use it to define generic 32 and 64bit variants. 2020-06-13 14:43:10 -07:00
CMakeLists.txt Fix some compiler option and warning levels. 2020-07-04 17:38:01 -07:00
LICENSE Add license 2020-04-26 20:10:24 -07:00
README.md Some fixes to get npcomp building and passing on windows. 2020-07-01 21:28:04 -07:00
contributing.md Add contributing.md 2020-04-26 20:17:16 -07:00

README.md

NPComp - An aspirational MLIR based numpy compiler

NPComp aims to be an idiomatic subset of the Python language, suitable for extracting isolated, statically typed programs from a running Python session. It is inspired by many projects that have come before it, including:

  • PyPy/RPython
  • Numba
  • Pythran
  • TorchScript
  • Autograph

As the name implies, NPComp also seeks to provide compiler-backed support for Numpy APIs.

The project spawned out of both LLVM's MLIR project and The IREE Project and seeks to use the MLIR and IREE tooling to enable progressive lowering of high level compute dominant sub-programs in a way that preserves high level semantic information that is expected to be useful for exploiting parallelism, generating high performance code, and enabling portability and deployment to a range of devices. Some of these goals overlap with existing projects, and to a first approximation, the experiment with NPComp is to determine whether rebasing on the MLIR tooling and ML backends like IREE produce a lift.

Before getting too excited, keep in mind that this project barely exists: it is very new and doesn't do anything useful yet :) We are using it as a testing ground for some new ideas and infrastructure improvement, and depending on how things turn out, may end up carrying it forward or breaking it up for parts.

See the features doc for a semi-curated status of what is implemented.

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:

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/python_native):$(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