Tensor Comprehensions - 将数学符号快速转换为高性能机器学习代码
Apache
跨平台
C/C++
软件简介
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 ...; });