-
Notifications
You must be signed in to change notification settings - Fork 1
Torch优化手段调研
why Torch Inductor： PyTorch需要一种与PyTorch的抽象紧密匹配的本地编译器。
什么是Torch Inductor： TorchInductor是PyTorch的一个新编译器，它能够表示所有PyTorch的功能，并且以一种通用方式构建，以便能够支持训练和多个后端目标。TorchInductor能够通过拥有与torch.Tensor和torch.Storage一对一映射的TensorBox和StorageBox的概念来表示别名和突变。它能够通过拥有一个符号步长的张量来处理视图，这个张量直接从本地的torch.Tensor步长表示映射而来，这使得视图易于处理。PyTorch的其他部分也以类似的方式处理，即在后端映射PyTorch的数据模型。设计理念是一种薄薄的、易于黑客攻击的方式，将PyTorch符号性地映射到更底层的后端，并启用快速实验、不同后端之间的自动调优以及更高层次的优化，如内存规划。
Lower 到哪些后端：
Triton：是一种新的编程语言，它提供了比 CUDA 更高的生产力，同时能够以清晰简单的代码超越高度优化的库如 cuDNN 的性能。它由 OpenAI 的 Philippe Tillet 开发，并在整个行业中获得巨大的采用和吸引力。Triton 支持 NVIDIA GPU，并且作为手动编写 CUDA 内核的替代品，其受欢迎程度正在迅速增长。
C++/OpenMP：是一种广泛采用的编写并行内核的规范。OpenMP 提供了一个工作共享的并行执行模型，并支持 CPU。C++ 也是一个有趣的目标，因为它是一种高度可移植的语言，并且能够支持导出到更多独特的边缘设备和硬件架构。
Lower Line：
-
PyTorch → TorchDynamo → TorchInductor → Triton → NVIDIA GPU
-
PyTorch → TorchDynamo → TorchInductor → OpenMP → CPU
有用的功能：
TorchInductor使用了一种基于define-by-run的循环级中间表示（IR）。IR的许多部分是Python可调用对象，它们接受SymPy表达式作为输入。
# example:
x.permute(1, 0) + x[2, :]
# to
def inner_fn(index: List[sympy.Expr]):
i1, i0 = index
tmp0 = ops.load("x", i1 + i0*size1)
tmp1 = ops.load("x", 2*size1 + i0)
return ops.add(tmp0, tmp1)
torchinductor.ir.Pointwise(
device=torch.device("cuda"),
dtype=torch.float32,
inner_fn=inner_fn,
ranges=[size0, size1],
)“Define-by-run” 是一种编程范式，它允许程序在运行时动态定义计算图。这种范式最著名的实现是在Python深度学习框架PyTorch中。在传统的"define-and-run"（先定义后运行）范式中，如TensorFlow的早期版本，用户首先定义一个静态的计算图，然后提交给引擎执行。这个静态图在执行过程中不会改变。
总结： Python 的动态语言特性，将给予了 Inductor 通过 Define-By-Run 便利地 JIT 化简 PyTorch 非 SSA 特性，非 Tensor 计算和动态 Shape/Layout 的能力。