Tensor Comprehensions (TC) 是一个功能完整的 C++ 库，能够自动地使用 Halide、ISL 以及 NVRTC 或 LLVM 来合成高性能的机器学习内核。TC 还提供了与 Caffe2 和 PyTorch 的基本集成。更多细节请参阅我们在 arXiv 上发表的论文。

该库设计为高度可移植、与机器学习框架无关，仅需要一个具备内存分配、卸载和同步功能的简单张量库。

目前，我们已经将 TC 与 Caffe2 和 PyTorch 集成。

一个简单的例子

以下示例展示了该库一个简短但强大的功能：能够按需为特定尺寸即时编译（JIT）高性能的机器学习内核。

import tensor_comprehensions as tc
import torch
lang = """
def tensordot(float(N, C1, C2, H, W) I0, float(N, C2, C3, H, W) I1) -> (O) {
    O(n, c1, c3, h, w) +=! I0(n, c1, c2, h, w) * I1(n, c2, c3, h, w)
}
"""
N, C1, C2, C3, H, W = 32, 512, 8, 2, 28, 28
tensordot = tc.define(lang, name="tensordot")
I0, I1 = torch.randn(N, C1, C2, H, W).cuda(), torch.randn(N, C2, C3, H, W).cuda()
best_options = tensordot.autotune(I0, I1, cache=True)
out = tensordot(I0, I1, options=best_options)

在 2-GPU P100 系统上进行几代 autotuning 后，我们得到类似以下的结果：

在 C++ 中，一个最小的自动调优示例类似于以下代码：

TEST(TensorDot, SimpleAutotune) {
  // 1. 定义并设置 TC 编译单元，使用 ATen 张量支持的 CUDA 内存管理。
  std::string tc = R"TC(
def tensordot(float(N, C1, C2, H, W) I0,
              float(N, C2, C3, H, W) I1)  -> (O)
{
    O(n, c1, c3, h, w) +=! I0(n, c1, r_c2, h, w) * I1(n, r_c2, c3, h, w)
}
  )TC";

  // 2. 分配具有随机数据的张量。
  at::Tensor I0 = at::CUDA(at::kFloat).rand({32,  8, 16, 17, 25});
  at::Tensor I1 = at::CUDA(at::kFloat).rand({32, 16, 2, 17, 25});

  // 3. 使用进化搜索进行自动调优，从一个朴素的选项开始。
  auto naiveOptions = Backend::MappingOptionsType::makeNaiveMappingOptions();
  tc::aten::ATenAutotuner<tc::CudaBackend, tc::autotune::GeneticSearch>
      geneticAutotuneATen(tc);
  auto bestOption =
      geneticAutotuneATen.tune("tensordot", {I0, I1}, {naiveOptions});

  // 4. 分配输出张量后，使用最佳选项编译并运行 TC。
  auto pExecutor =
      tc::aten::compile<Backend>(tc, "tensordot", {I0, I1}, bestOption[0]);
  auto outputs = tc::aten::prepareOutputs(tc, "tensordot", {I0, I1});
  auto timings = tc::aten::profile(*pExecutor, {I0, I1}, outputs);
  std::cout << "tensordot size I0: " << I0.sizes() << ", "
            << "size I1: " << I1.sizes()
            << " ran in: " << timings.kernelRuntime.toMicroSeconds() << "us\n";
}

请注意，对于一个给定的 TC，我们只需要自动调优一次即可获得合理的映射选项，这些选项可以应用于其他问题尺寸，如下面的代码片段所示：

// 5. 在其他内核上复用自动调优得到的最佳选项
for (auto sizes : std::vector<std::pair<at::IntList, at::IntList>>{
         {{4, 9, 7, 16, 14}, {4, 7, 3, 16, 14}},
         {{8, 5, 11, 10, 10}, {8, 11, 16, 10, 10}},
     }) {
  at::Tensor I0 = makeATenTensor<Backend>(sizes.first);
  at::Tensor I1 = makeATenTensor<Backend>(sizes.second);
  auto pExecutor =
      tc::aten::compile<Backend>(tc, "tensordot", {I0, I1}, bestOption[0]);
  auto outputs = tc::aten::prepareOutputs(tc, "tensordot", {I0, I1});
  auto timings = tc::aten::profile(*pExecutor, {I0, I1}, outputs);
  std::cout << "tensordot size I0: " << I0.sizes() << ", "
            << "size I1: " << I1.sizes()
            << " ran in: " << timings.kernelRuntime.toMicroSeconds()
            << "us\n";
}

将所有部分组合起来，你可能会看到类似以下的输出：

> build$ ./examples/example_simple
[==========] Running 1 test from 1 test case.
[----------] Global test environment set-up.
[----------] 1 test from TensorDot
[ RUN      ] TensorDot.SimpleAutotune
Generation 0    Jobs(Compiled, GPU)/total  (10, 10)/10   (best/median/worst)us: 226/4238/7345
Generation 1    Jobs(Compiled, GPU)/total  (10, 10)/10   (best/median/worst)us: 220/221/233
Generation 2    Jobs(Compiled, GPU)/total  (10, 10)/10   (best/median/worst)us: 220/221/234
tensordot size I0: [16, 8, 16, 17, 25], size I1: [16, 16, 2, 17, 25] ran in: 239us
tensordot size I0: [4, 9, 7, 16, 14], size I1: [4, 7, 3, 16, 14] ran in: 56us
tensordot size I0: [8, 5, 11, 10, 10], size I1: [8, 11, 16, 10, 10] ran in: 210us
[       OK ] TensorDot.SimpleAutotune (27812 ms)
[----------] 1 test from TensorDot (27812 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test case ran. (27812 ms total)
[  PASSED  ] 1 test.

我们尚未精确量化所获得的峰值性能比例，但在自动调优后，获得 80% 以上的共享内存峰值带宽并不罕见。寄存器级别的深度优化仍在进行中，但 TC 目前的形式已经弥合了研究需求与生产需求之间的生产力差距。这也是我们为何兴奋地与整个社区分享，并将这项协作努力开源的原因。

文档

通用文档：你可以在这里找到关于 Tensor Comprehensions 的详细信息。

C++ API：我们还提供了 C++ API 的文档，可以在这里找到。

安装

二进制包

我们提供了 conda 包，以便于安装和使用 TC 二进制文件。请参阅我们的文档此处获取安装说明。

从源码构建

你可以在这里找到文档，其中包含了通过 Docker、conda 包或在非 conda 环境中构建 TC 的说明。

交流

• 电子邮件：tensorcomp@fb.com
• GitHub issues：错误报告、功能请求、安装问题、RFC、想法等。

行为准则

详情请参阅 CODE_OF_CONDUCT.md 文件。

许可证

Tensor Comprehensions 采用宽松的 Apache v2.0 许可证分发，详情请参阅 LICENSE 文件。

贡献

详情请参阅 CONTRIBUTING.md 文件。