Comet-rs

该crate是COMET编译器的基于Rust的eDSL前端。

⚡️ "COMET" / 针对极端目标的领域特定编译器

COMET编译器由稀疏和稠密张量代数计算的领域特定语言(dsl)、将高级操作映射到低级架构资源的过程化降级过程、在降级过程中执行的一系列优化以及表示多级IR中各个级别的关键概念、操作和类型的各种IR方言组成。在IR堆栈的每个级别，COMET执行不同的优化和代码转换。依赖高级语义信息的特定领域、硬件无关的优化应用于高级IR。这包括将高级操作重新表述为适合在异构设备上执行的形式（例如，将张量收缩操作重写为转置-转置-GEMM-转置）以及自动并行化高级原语（例如，为线程和任务级并行性进行分块）。

通过使用过程宏，此crate在用户应用编译时向用户公开了一个基于Rust的eDSL，它将被降低到各种方言，然后编译为共享库。在运行时，共享库被动态链接，暴露编译后的COMET函数，允许用户像执行其他函数一样执行它们。

外部依赖

• COMET

请遵循COMET仓库中的构建说明来安装COMET编译器。（也包括在下面的COMET安装部分中）

Cargo.toml

[dependencies]
comet-rs = "0.1.0"

所需环境变量

如果您严格按照COMET构建说明进行，您只需设置以下环境变量

COMET_DIR=/path/to/COMET/root/dir

如果您已将构建位置从COMET构建说明中列出的位置更改，则需要设置以下环境变量

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

请注意，作为COMET构建过程的一部分，我们还将构建MLIR的特定版本（作为git子模块管理），COMET只能与这个特定的提交一起工作，因此请不要指向您在COMET构建过程之外构建的任何不同的MLIR版本。

示例

COMET使用爱因斯坦数学符号，comet_fn! 宏为用户提供了一个接口，使用类似Rust的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();
}

操作

我们实现了以下张量操作（其中大多数不是有效的Rust语法，但却是有效的COMET eDSL语法），请参阅COMET文档以获取每个操作的更详细描述。

• 乘法： A * B
• 逐元素乘法： A .* B
• 半群运算： @(op1, op2)
  • 最小值： A @(min) B
  • 加法： A @(+) B
  • 乘法： A @(*) B
  • 任意对： A @(any,pair) B
  • 加乘： A @(+,*)B
  • 加对： A @(+,pair) B
  • 加第一： A @(+,first) B
  • 加第二： A @(+,Second) B
  • 最小加： A @(min,+) B
  • 最小第一： A @(min,first) B
  • 最小第二： A @(min,second) B
• 转置： B = A.transpose()

优化

我们还可以指定COMET编译器执行的各种优化。请参阅COMET文档以获取每个优化的更详细描述。

• Permutation TTGT： BestPermTtgt
• Tensor Algebra to Index Tree： TaToIt
• Tensor Contraction to TTGT： TcToTtgt
• Loops： ToLoops
• 矩阵乘法内核：MatMulKernel
• 矩阵乘法分块：MatMulTiling
• 密集转置：DenseTranspose
• 工作空间：CompWorkspace

上述优化可以作为自定义语法的一部分传递给comet_fn宏的参数。

示例

comet_fn! {function_name, {
        eDSL code
},
CometOptions::[TcToTtgt, BestPermTtgt, ToLoop]
}

COMET 输出

在评估comet_fn宏期间，COMET编译器将生成一个包含编译好的COMET函数的共享库。共享库将位于与用户应用程序相同的目录中，目录标签为comet_libs。这个crate会自动将此共享库链接到您的应用程序中，但它依赖于该库保持在comet_libs目录中。

crate 特性

默认情况下，如果COMET函数编译失败，也会导致整个Rust应用程序失败。在某些情况下（例如运行此crate中提供的单元测试）可能需要仅发出警告而不是错误。在这种情况下，procedure宏会简单地创建Rust代码，该代码会打印出COMET编译器的错误而不是执行函数。此功能可以通过指定comet_errors_as_warnings特性来启用。

cargo test --no-fail-fast --features comet_errors_as_warnings

COMET 安装说明

以下命令可用于设置COMET项目：如果没有成功安装COMET，此crate将无法工作。

1. 安装依赖项 安装COMET编译器需要cmake和ninja。

2. 检出LLVM和COMET仓库 COMET包含作为git子模块的LLVM。这里的LLVM仓库包含了对MLIR的阶段性更改，这可能对于支持COMET是必要的。它还代表了经过测试的LLVM版本。MLIR仍在相对快速地变化，因此您可以自由使用LLVM的当前版本，但API可能已经更改。

$ git clone https://github.com/pnnl/COMET.git
$ cd comet
$ git submodule init
$ git submodule update

注意： 仓库已设置为git submodule update执行浅克隆，这意味着只下载足够的LLVM仓库来检出当前指定的提交。如果您希望使用LLVM仓库的完整历史，您可以手动“取消浅克隆”子模块。

$ cd llvm
$ git fetch --unshallow

3. 构建和测试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

4. 构建和测试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.

-DCMAKE_BUILD_TYPE=DEBUG标志启用了调试信息，这使得整个树编译速度变慢，但允许您在LLVM和MLIR框架中逐步执行代码。

要获得运行速度快的版本，请使用-DCMAKE_BUILD_TYPE=Release或如果需要同时包含调试信息，请使用-DCMAKE_BUILD_TYPE=RelWithDebInfo。在Release模式下，性能会有很大差异。