Crate comet_rs

Comet-rs is a Rust-based eDSL frontend for the COMET compiler.

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 IR 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

• COMET

Please follow the build instructions in the COMET repository to install the COMET compiler. This crate will not work without a successful COMET installation.

Cargo.toml

[dependencies]
comet-rs = "0.1.0"

Required 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 mave have build outside the COMET build process.

Example

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
• 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 remainng in the comet_libs directory.

Modules

index

Range-based index label constructs (e.g., [i, j]) represent the range of indices expressed through a scalar, a range, or a range with increment. Index labels can be used both for constructing a tensor or for representing a tensor operation. In a tensor construction, index labels are used to represent each dimension size. In the context of a tensor operation, they represent slicing information of the tensor object where the operation will be applied.

tensor

A tensor object (e.g., A) refers to a multi-dimensional array of arithmetic values that can be accessed by using indexing values.

Macros

comet_fn

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.