comet_rs/lib.rs
1//! Comet-rs is a Rust-based eDSL frontend for the COMET compiler.
2//!
3//! The COMET compiler consists of a Domain Specific Language (DSL) for sparse and dense tensor algebra computations,
4//! a progressive lowering process to map high-level operations to low-level architectural resources,
5//! a series of optimizations performed in the lowering process, and various IR dialects to represent key concepts, operations,
6//! and types at each level of the multi-level IR. At each level of the IR stack,
7//! COMET performs different optimizations and code transformations.
8//! Domain-specific, hardware- agnostic optimizations that rely on high-level semantic information are applied at high-level IRs.
9//! These include reformulation of high-level operations in a form that is amenable for execution on heterogeneous devices
10//! (e.g., rewriting Tensor contraction operations as Transpose-Transpose-GEMM-Transpose)
11//! and automatic parallelization of high-level primitives (e.g., tiling for thread- and task-level parallelism).
12//!
13//! Through the use of procedural macros, this crate exposes to the user a Rust based eDSL
14//! which will be lowered to various IR dialects and then compiled into a shared library at compile time of the user application
15//! 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.
16//!
17//! # External Dependencies
18//! - [COMET](https://github.com/pnnl/COMET)
19//!
20//! Please follow the build instructions in the COMET repository to install the COMET compiler.
21//! This crate will not work without a successful COMET installation.
22//!
23//! # Cargo.toml
24//! ```toml
25//! [dependencies]
26//! comet-rs = "0.1.0"
27//!```
28//! # Required Environment Variables
29//! ```text
30//! COMET_BIN_DIR = /path/to/COMET/bin/dir
31//! COMET_LIB_DIR = /path/to/COMET/lib/dir
32//! MLIR_BIN_DIR = /path/to/MLIR/bin/dir
33//! MLIR_LIB_DIR = /path/to/MLIR/lib/dir
34//! ```
35//! *Note that as part of the COMET build process we will also build a specific version of MLIR (managed as a git submodule),
36//! 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.*
37//!
38//! # Example
39//!
40//! ```
41//! use comet_rs::*;
42//!
43//! comet_fn! { dense_dense_matrix_multiply, {
44//! let a = Index::with_value(2);
45//! let b = Index::with_value(2);
46//! let c = Index::with_value(2);
47//!
48//! let A = Tensor::<f64>::dense([a, b]).fill(2.2);
49//! let B = Tensor::<f64>::dense([b, c]).fill(3.4);
50//! let C = Tensor::<f64>::dense([a, c]).fill(0.0);
51//! C = A * B;
52//! C.print();
53//! }}
54//!
55//! fn main() {
56//! dense_dense_matrix_multiply();
57//! }
58//! ```
59//! # Operations
60//!
61//! We have implemented the following tensor operations (most of which are not valid rust syntax, but are valid COMET eDSL syntax)
62//! please refer to the [COMET documentation](https://pnnl-comet.readthedocs.io/en/latest/operations.html) for more in-depth descriptions of each operation.
63//! - Multiplication: `A * B`
64//! - Elementwise Multiplication: `A .* B`
65//! - Semiring Operations: `@(op1, op2)`
66//! - Min: ` A @(min) B`
67//! - Plus: `A @(+) B`
68//! - Mul: `A @(*) B`
69//! - Any-Pair: `A @(any,pair) B`
70//! - Plus-Mul: `A @(+,*)B `
71//! - Plus-Pair: `A @(+,pair) B`
72//! - Plus-First: `A @(+,first) B`
73//! - Plus-Second: `A @(+,Second) B`
74//! - Min-Plus: `A @(min,+) B`
75//! - Min-First: `A @(min,first) B`
76//! - Min-Second: `A @(min,second) B`
77//! - Transpose: `B = A.transpose()`
78//!
79//! # Optimizations
80//! We also support the ability to specify various optimizations to be performed by the COMET compiler.
81//! please refer to the [COMET documentation](https://pnnl-comet.readthedocs.io/en/latest/optimizations.html) for more in-depth descriptions of each optimization.
82//! - [Permutation TTGT](https://pnnl-comet.readthedocs.io/en/latest/passes/PermTTGT.html): `BestPermTtgt`
83//! - [Tensor Algebra to Index Tree](https://pnnl-comet.readthedocs.io/en/latest/passes/TAtoIT.html): `TaToIt`
84//! - [Tensor Contraction to TTGT](https://pnnl-comet.readthedocs.io/en/latest/passes/TC.html): `TcToTtgt`
85//! - [Loops](https://pnnl-comet.readthedocs.io/en/latest/passes/loops.html): `ToLoops`
86//! - [Matmult Kernel](https://pnnl-comet.readthedocs.io/en/latest/passes/mkernel.html): `MatMulKernel`
87//! - [Matmult Tiling](https://pnnl-comet.readthedocs.io/en/latest/passes/tiling.html): `MatMulTiling`
88//! - [Dense Transpose](https://pnnl-comet.readthedocs.io/en/latest/passes/transpose.html): `DenseTranspose`
89//! - [Workspace](https://pnnl-comet.readthedocs.io/en/latest/passes/workspace.html): `CompWorkspace`
90//!
91//! The above optimizations can be passed to the compiler as part of a custom syntax proivded as an argument to the `comet_fn` macro.
92//!
93//! #### Example
94//! ```
95//! comet_fn! {function_name, {
96//! eDSL code
97//! },
98//! CometOptions::[TcToTtgt, BestPermTtgt, ToLoop]
99//! }
100//! ```
101//!
102//! # COMET Output
103//! During evaluation of the `comet_fn` macro, the COMET compiler will generate a shared library containing the compiled COMET function.
104//! The shared library will be located in the same directory as the user application in a directory labelled `comet_libs`.
105//! This crate handles linking this shared library into your application automatically, but it depends on the library remainng in the `comet_libs` directory.
106
107
108use comet_rs_impl;
109
110/// 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.
111///
112/// # Syntax
113/// * `function_name` - is the (user defined) name of the COMET function and can be called directly from the rust program as `function_name()`.
114/// * `eDSL code` - is the eDSL code to be compiled by the COMET compiler.
115/// * `CometOptions` - are optional optimizations to pass to the COMET compiler.
116/// ```
117/// comet_rs::comet_fn! {function_name, {
118/// eDSL code
119/// }[,]
120/// [CometOptions::[options],]
121/// }
122/// ```
123/// # Example
124/// ```
125/// comet_rs::comet_fn! { print_dense, {
126/// let i = Index::with_value(4);
127/// let j = Index::with_value(4);
128///
129/// let A = Tensor::<f64>::dense([i, j]).fill(2.3);
130///
131/// A.print();
132///}}
133///
134///comet_rs::comet_fn! { sum_coo, {
135/// let i = Index::new();
136/// let j = Index::new();
137///
138/// let A = Tensor::<f64>::coo([i, j]).fill_from_file("../../../integration_test/data/test.mtx");
139/// let a = A.sum();
140/// a.print();
141///}}
142///
143///comet_rs::comet_fn! { ccsd_t1_4_ttgt_bestperm, {
144/// let i, c = Index::with_value(2);
145/// let m, a = Index::with_value(4);
146///
147/// let v = Tensor::<f64>::dense([c, i, m, a]).fill(2.3);
148/// let t2 = Tensor::<f64>::dense([m, c]).fill(3.4);
149/// let i0 = Tensor::<f64>::dense([i, a]).fill(0.0);
150/// i0 = v * t2;
151/// i0.print();
152///},
153///CometOption::[BestPermTtgt, TcToTtgt]
154///}
155///
156///fn main() {
157/// print_dense();
158/// sum_coo();
159/// ccsd_t1_4_ttgt_bestperm();
160///}
161/// ```
162
163pub use comet_rs_impl::comet_fn; //,main,function};
164
165use std::sync::Arc;
166
167use std::collections::{HashMap, HashSet};
168
169#[doc(hidden)]
170pub use inventory;
171
172use ctor::ctor;
173
174#[doc(hidden)]
175#[ctor]
176pub static COMET_FUNCS: HashMap<
177 &'static str,
178 (
179 libloading::os::unix::Symbol<unsafe extern "C" fn()>,
180 Arc<libloading::Library>,
181 ),
182> = {
183 let mut shared_libs = HashSet::new();
184 for lib in crate::inventory::iter::<CometLib> {
185 shared_libs.insert(lib.name.clone());
186 }
187 let mut libs = vec![];
188 for lib in shared_libs {
189 libs.push(Arc::new(libloading::Library::new(lib).unwrap()));
190 }
191 let mut func_names = HashSet::new();
192 for func in crate::inventory::iter::<CometFunc> {
193 func_names.insert(func.name.clone());
194 }
195 let mut map = HashMap::new();
196 for func in func_names {
197 let func_ptr = libs
198 .iter()
199 .filter_map(
200 |lib| match lib.get::<unsafe extern "C" fn()>(func.as_bytes()) {
201 Ok(func_ptr) => Some((func_ptr.into_raw(), lib.clone())),
202 Err(_) => None,
203 },
204 )
205 .last()
206 .expect("function not found");
207 map.insert(func.clone(), func_ptr);
208 }
209 map
210};
211
212#[doc(hidden)]
213#[derive(Debug)]
214pub struct CometFunc {
215 pub name: &'static str,
216}
217
218#[doc(hidden)]
219#[derive(Debug)]
220pub struct CometLib {
221 pub name: &'static str,
222}
223
224crate::inventory::collect!(CometFunc);
225crate::inventory::collect!(CometLib);
226
227pub mod index;
228pub mod tensor;
229
230