- 授权协议: Apache
- 开发语言: C/C++
- 操作系统: 跨平台
- 软件首页: https://facebookresearch.github.io/TensorComprehensions/
- 软件文档: https://github.com/facebookresearch/TensorComprehensions/
- 官方下载: https://github.com/facebookresearch/TensorComprehensions/
软件介绍
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 ...; });
