Tensor Compute Primitives

Mid-level intermediate representation for machine learning programs.

🚧 This project is under active development (WIP).

Project Communication

• For general discussion use #mlir-tcp channel on the LLVM Discord
• For feature request or bug report file a detailed issue on GitHub

Developer Guide

To build TCP using Bazel, follow these steps:

1. (Optional) For a quick start, launch an interactive docker container with clang (and lld) pre-installed:

./docker/run_docker.sh

2. You can now build tcp-opt by running:

bazel build //:tcp-opt

3. To run TCP lit and aot compile tests:

bazel test //...

We welcome contributions to mlir-tcp. When authoring new TCP ops with dialect conversions from/to Torch and Linalg, please include lit tests for dialect and conversions, as well as aot_compile generated e2e integration tests. Lastly, please finalize your PR with clang-format, black and bazel buildifier to ensure the C++/python sources and BUILD files are formatted consistently:

# clang-format
find . -type f -name "*.cpp" -o -name "*.h" | xargs clang-format -i

# black
black .

# buildifer
bazel run //tools/buildifier:buildifier

To enable clangd (for code completion, navigation and insights), generate the compilation database using bazel-compile-commands-extractor:

bazel build //...

bazel run //tools/clangd:refresh_compile_commands

When run successfully, a compile_commands.json is generated at the workspace root (and refreshed upon re-runs). If you're using VSCode, just hit CMD+SHIFT+P and select clangd: Restart language server to start clangd. Note that this only works for non-docker builds at the moment.

When bumping upstream dependencies (LLVM, Torch-MLIR, StableHLO), you may validate the set of "green commits" by running the corresponding third-party tests:

bazel test @llvm-project//mlir/...
bazel test @torch-mlir//...
bazel test @stablehlo//...

The following CI workflows are automatically triggered anytime upstream dependencies (deps.bzl) are updated:

• 
• 
• 

To use newer torch-mlir and/or torch python packages in our hermetic python sandbox, just regenerate requirements_lock.txt as follows:

truncate -s 0 requirements_lock.txt
bazel run //tools/pip:requirements.update

Debugging Guide

Below are some standard techniques for debugging your compilation process, assuming you've reduced it to a form that can be reproduced with tcp-opt. For MLIR-specific debugging tips, refer here.

printf debugging

Printing to stdout/stderr works as usual:

op.emitWarning() << "HERE: " << myVariable;      // preferred for op/loc diagnostics

llvm::errs() << "HERE: " << myVariable << "\n";  // alternative

You can also hook into the LLVM_DEBUG macro:

#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "foo"
LLVM_DEBUG(llvm::dbgs() << "This only shows up when -debug or -debug-only=foo is provided.\n");
#undef DEBUG_TYPE

#define DEBUG_TYPE "bar"
LLVM_DEBUG(llvm::dbgs() << "This only shows up when -debug or -debug-only=bar is provided.\n");
#undef DEBUG_TYPE

Then run with the -debug-only=foo,bar flag to cuts out messages that aren't associated with the passed DEBUG_TYPEs.

bazel run //:tcp-opt -- --some-pass `pwd`/test.mlir -debug-only=foo,bar

gdb debugging

To debug tcp-opt with gdb:

bazel build --config=gdb //:tcp-opt

gdb --args bazel-bin/tcp-opt -h

For help with gdb commands please refer to gdb cheat sheet.

aot_compile debugging

Refer this README for a step-by-step guide to debugging an end-to-end compilation pipeline using the AOT Compile framework.

Enable llvm-symbolizer

If you get a stack dump without any symbol names:

Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  tcp-opt   0x000055ac1c9c0c1d
1  tcp-opt   0x000055ac1c9c110b
2  tcp-opt   0x000055ac1c9be846
3  tcp-opt   0x000055ac1c9c1855
4  libc.so.6 0x00007f7011c6a520
...

Do this and re-run:

bazel build @llvm-project//llvm:llvm-symbolizer
export LLVM_SYMBOLIZER_PATH=`pwd`/bazel-bin/external/llvm-project/llvm/llvm-symbolizer