将TVM集成到PyTorch中。

* TVM：深度学习编译器

Github项目链接：

https://github.com/pytorch/tvm

构建

你需要在这个PR的基础上构建PyTorch： https://github.com/pytorch/pytorch/pull/18588

cd pytorchgit fetch origin pull/18588/head:tvm_devgit checkout tvm_devpython setup.py install

然后，你需要单独构建本仓库：

# Make sure the right llvm-config is in your PATHpython setup.py install

测试

python setup.py test

使用

这个包非常显然地挂钩到 PyTorch 的 JIT 中，因此适用相同的工具（可查看 @torch.jit.script，torch.jit.trace 和 graph_for）。以下是使用示例：

from tvm import relay # This imports all the topi operatorsimport torch_tvm
torch_tvm.enable()
# The following function will be compiled with TVM@torch.jit.scriptdef my_func(a, b, c):    return a * b + c

如果要禁用JIT挂钩，请使用 torch_tvm.disable() 。

代码布局

• register.cpp：设置pybind绑定并调用TVM后端的注册。

• compiler.{h,cpp}：用TVM编译PyTorch JIT图的主要逻辑。

• operators.{h,cpp}:从JIT IR映射到TVM操作符的位置。

TODO