4 个版本

0.1.1-rc3	2023 年 2 月 24 日
0.1.0	2022 年 6 月 22 日

#409 在编程语言中

自定义许可

54KB
339 行

Comet-rs

该软件包是 COMET 编译器的基于 Rust 的 eDSL 前端。

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

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

通过使用过程宏，该软件包向用户暴露了一个基于 Rust 的 eDSL，该 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使用爱因斯坦数学符号，并且《code>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。这个包将自动将这个共享库链接到你的应用程序中，但它依赖于库仍然位于 comet_libs 目录下。

包特性

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

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

COMET 安装说明

以下命令可以用来设置 COMET 项目：此包在没有成功安装 COMET 的情况下无法工作。

安装依赖项 要安装 COMET 和 LLVM/MLIR，需要安装以下依赖项

检出 LLVM 和 COMET 仓库。 COMET 将 LLVM 作为 git 子模块。此处的 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 仓库的完整历史记录，你可以手动将子模块“unshallow”。

构建和测试 LLVM/MLIR

$ 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

构建和测试 COMET

$ cd ../../
$ 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 模式在性能上有很大的差异。

依赖项

~4–6MB
~105K SLoC