Tensor Comprehensions将数学符号快速转换为高性能机器学习代码

Tensor Comprehensions 是 Facebook AI  研究院开源的  C++  库及数学语言，功能齐全，能有效填补研究人员于数学运算领域的沟通鸿沟，并基于各种硬件后端上大规模运行工程模型。

Tensor Comprehensions 采用了 Just-In-Time 的编译自动生成机器学习社区所需的高性能代码，并被设计为高度可移植的。通过 Tensor Comprehensions，研究人员能够以数学符号的方式进行编写，系统能够根据需求进行编译调整，并输出专业的代码。

示例：

  #include <ATen/ATen.h>
  #include "tc/aten/aten_compiler.h"
  #include "tc/core/mapping_options.h"

  // 1. Define and setup the TC compilation unit with CUDA memory management backed by ATen.
  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, c2, h, w) * I1(n, c2, c3, h, w)
  })TC";

  // 2. Allocate tensors with random data
  at::Tensor I0 = at::CUDA(at::kFloat).rand({32, 512, 8, 28, 28});
  at::Tensor I1 = at::CUDA(at::kFloat).rand({32,   8, 2, 28, 28});
  std::vector<at::Tensor> outputs;

  // 3. Run autotuning with evolutionary search starting from a naive option
  auto options = tc::MappingOptions::makeNaiveMappingOptions();
  auto bestOption = autotune(cacheFilename, tc, "TensorDot", {I0, I1}, options, {options});

  // 4. Compile and run the TC with the best option.
  tc::ATenCompilationUnit atCompl;
  atCompl.define(tc);
  auto handle = atCompl.compile("TensorDot", {I0, I1}, bestOption);
  atCompl.run("TensorDot", {I0, I1}, outputs, handle);

  // 5. Perform precision checks against an ATen reference implementation
  check({I0, I1}, outputs, [&I0, &I1](){ return ...; });