3 releases
0.1.1-rc2 | Feb 24, 2023 |
---|---|
0.1.0 | Jun 22, 2022 |
#3 in #comet
30 downloads per month
Used in comet-rs
110KB
2.5K
SLoC
Comet-rs
This crate is a Rust-based eDSL front end for the COMET compiler.
⚡️ "COMET" / Domain specific COMpiler for Extreme Targets
The COMET compiler consists of a Domain Specific Language (DSL) for sparse and dense tensor algebra computations, a progressive lowering process to map high-level operations to low-level architectural resources, a series of optimizations performed in the lowering process, and various IR dialects to represent key concepts, operations, and types at each level of the multi-level IR. At each level of the IR stack, COMET performs different optimizations and code transformations. Domain-specific, hardware- agnostic optimizations that rely on high-level semantic information are applied at high-level IRs. These include reformulation of high-level operations in a form that is amenable for execution on heterogeneous devices (e.g., rewriting Tensor contraction operations as Transpose-Transpose-GEMM-Transpose) and automatic parallelization of high-level primitives (e.g., tiling for thread- and task-level parallelism).
Through the use of procedural macros, this crate exposes to the user a Rust based eDSL which will be lowered to various dialects and then compiled into a shared library at compile time of the user application. At runtime the shared library is dynamically linked to, exposing the compiled COMET functions, allowing the user to execute them as the would any other function.
External Dependencies
Please follow the build instructions in the COMET repository to install the COMET compiler. (Also included in the COMET installation section below)
Cargo.toml
[dependencies]
comet-rs = "0.1.0"
Required Environment Variables
If you have followed the COMET build instructions exactly, you will only need to set the following envrionment variable:
COMET_DIR=/path/to/COMET/root/dir
If you have changed the build locations from what was listed in the COMET build instructions you will instead need to set the following environment variables:
COMET_BIN_DIR=/path/to/COMET/bin/dir
COMET_LIB_DIR=/path/to/COMET/lib/dir
MLIR_BIN_DIR=/path/to/MLIR/bin/dir
MLIR_LIB_DIR=/path/to/MLIR/lib/dir
Note that as part of the COMET build process we will also build a specific version of MLIR (managed as a git submodule), COMET will only work with this specific commit, so please do not point to a different MLIR version you might have build outside the COMET build process.
Example
COMET uses Einstein mathematical notation and The comet_fn!
macro provides users with an interface to express tensor algebra semantics using a Rust-like eDSL.
use comet_rs::*;
comet_fn! { dense_dense_matrix_multiply, {
let a = Index::with_value(2);
let b = Index::with_value(2);
let c = Index::with_value(2);
let A = Tensor::<f64>::dense([a, b]).fill(2.2);
let B = Tensor::<f64>::dense([b, c]).fill(3.4);
let C = Tensor::<f64>::dense([a, c]).fill(0.0);
C = A * B;
C.print();
}}
fn main() {
dense_dense_matrix_multiply();
}
Operations
We have implemented the following tensor operations (most of which are not valid rust syntax, but are valid COMET eDSL syntax) please refer to the COMET documentation for more in-depth descriptions of each operation.
- Multiplication:
A * B
- Elementwise Multiplication:
A .* B
- Semiring Operations:
@(op1, op2)
- Min:
A @(min) B
- Plus:
A @(+) B
- Mul:
A @(*) B
- Any-Pair:
A @(any,pair) B
- Plus-Mul:
A @(+,*)B
- Plus-Pair:
A @(+,pair) B
- Plus-First:
A @(+,first) B
- Plus-Second:
A @(+,Second) B
- Min-Plus:
A @(min,+) B
- Min-First:
A @(min,first) B
- Min-Second:
A @(min,second) B
- Min:
- Transpose:
B = A.transpose()
Optimizations
We also support the ability to specify various optimizations to be performed by the COMET compiler. please refer to the COMET documentation for more in-depth descriptions of each optimization.
- Permutation TTGT:
BestPermTtgt
- Tensor Algebra to Index Tree:
TaToIt
- Tensor Contraction to TTGT:
TcToTtgt
- Loops:
ToLoops
- Matmult Kernel:
MatMulKernel
- Matmult Tiling:
MatMulTiling
- Dense Transpose:
DenseTranspose
- Workspace:
CompWorkspace
The above optimizations can be passed to the compiler as part of a custom syntax proivded as an argument to the comet_fn
macro.
Example
comet_fn! {function_name, {
eDSL code
},
CometOptions::[TcToTtgt, BestPermTtgt, ToLoop]
}
COMET Output
During evaluation of the comet_fn
macro, the COMET compiler will generate a shared library containing the compiled COMET function.
The shared library will be located in the same directory as the user application in a directory labelled comet_libs
.
This crate handles linking this shared library into your application automatically, but it depends on the library remaining in the comet_libs
directory
Crate Features
By default if a COMET function fails to compile, it will also cause the overall Rust application to fail.
It may be useful in some cases (e.g running the unit tests provided in this crate) to have a failed COMET compilation only emit a warning instead of an error.
In this case the procedure macro simply creates rust code that will print out the COMET compiler error instead of executing the function.
This functionality can be enabled by specifying the comet_errors_as_warnings
feature:
cargo test --no-fail-fast --features comet_errors_as_warnings
COMET Installation Instructions
These commands can be used to setup COMET project: This crate will not work without a successful COMET installation.
-
Install Dependencies cmake and ninja are required to install the COMET compiler.
-
Check out LLVM and COMET repos. COMET contains LLVM as a git submodule. The LLVM repo here includes staged changes to MLIR which may be necessary to support COMET. It also represents the version of LLVM that has been tested. MLIR is still changing relatively rapidly, so feel free to use the current version of LLVM, but APIs may have changed.
$ git clone https://github.com/pnnl/COMET.git
$ cd comet
$ git submodule init
$ git submodule update
Note: The repository is set up so that git submodule update
performs a
shallow clone, meaning it downloads just enough of the LLVM repository to check
out the currently specified commit. If you wish to work with the full history of
the LLVM repository, you can manually "unshallow" the the submodule:
$ cd llvm
$ git fetch --unshallow
- Build and test LLVM/MLIR:
$ cd comet
$ mkdir llvm/build
$ cd llvm/build
$ cmake -G Ninja ../llvm \
-DLLVM_ENABLE_PROJECTS="mlir" \
-DLLVM_TARGETS_TO_BUILD="X86" \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DCMAKE_BUILD_TYPE=DEBUG
$ ninja
$ ninja check-mlir
- Build and test COMET:
$ cd comet
$ mkdir build
$ cd build
$ cmake -G Ninja .. \
-DMLIR_DIR=$PWD/../llvm/build/lib/cmake/mlir \
-DLLVM_DIR=$PWD/../llvm/build/lib/cmake/llvm \
-DLLVM_ENABLE_ASSERTIONS=ON \
-DCMAKE_BUILD_TYPE=DEBUG
$ ninja
$ ninja check-comet-integration # Run the integration tests.
The -DCMAKE_BUILD_TYPE=DEBUG
flag enables debug information, which makes the
whole tree compile slower, but allows you to step through code into the LLVM
and MLIR frameworks.
To get something that runs fast, use -DCMAKE_BUILD_TYPE=Release
or
-DCMAKE_BUILD_TYPE=RelWithDebInfo
if you want to go fast and optionally if
you want debug info to go with it. Release
mode makes a very large difference
in performance.
Dependencies
~4–5.5MB
~102K SLoC