PyTorch 2.0的全新编译器后端:Torch Inductor
目录
- 介绍
- Torch Inductor的设计原则
- Torch Inductor的技术
- Mine by Run Loop Level AR
- 编译器堆栈概述
- Torch Inductor在GPU上的结果
- Torch Inductor在CPU上的结果
- 结论
- 相关资源
介绍
大家好,我是Jason Ansel。今天我将介绍Torch Inductor,这是PyTorch 2.0的新编译器后端。在设计Torch Inductor时,我们遵循了三个关键原则。首先,Torch Inductor是PyTorch原生的,这意味着Inductor使用与PyTorch Eager非常相似的抽象,可以准确捕捉PyTorch的所有行为。其次,我们秉承以Python为中心的原则,完全使用Python编写Inductor,使其更易于进行开发和扩展。最后,我们希望Inductor具有广泛的适用性,因此在早期主要关注处理复杂的运算符和优化,以确保设计的普适性和可扩展性。
Torch Inductor的设计原则
Python First
Torch Inductor将Python放在首位。这意味着我们使用Python完全编写Inductor,因此更容易进行hack和扩展。借助Python的强大功能,我们可以使用Inductor中的代码执行和分析IR,简化设计过程,加快开发速度。
通用性
Torch Inductor旨在实现通用性而非深度。我们早期将重点放在处理复杂运算符和优化上,以确保设计的通用性和扩展性。Torch Inductor同时支持动态形状和步幅,使用符号数学库Senpai来推理形状并生成不依赖于特定输入尺寸的代码。此外,我们还借鉴了用户在高性能计算方面的实践经验,使用了现代编程语言Triton的解决方案,使得我们在生成易于理解、可修改的Triton代码时更加得心应手。
Torch Inductor的技术
Torch Inductor使用了三种关键技术:
定义按运行循环的AR
Inductor的核心编译IR是在循环级别上完成的,并且是一个Python可调用的函数。通过执行此IR来进行代码生成和分析,这在编译器领域是一种新颖的技巧。实际上,这种方法早在PyTorch中就被广泛使用。使用定义按运行循环的IR,我们可以轻松编写降低操作,为运行时环境生成更高效的代码。
动态形状和步幅的支持
Torch Inductor从一开始就支持动态形状和步幅。为了实现这一目标,我们使用了符号数学库Senpai来推理和生成不依赖于特定输入尺寸的代码。这使得Torch Inductor能够适应各种不同的输入尺寸,提高代码的通用性和可扩展性。
Triton和C++代码生成
我们希望借鉴先进的编程语言,因此选择了现代编程语言Triton作为Torch Inductor生成代码的目标语言。Triton是一种介于Cuda和其他领域特定语言之间的高级语言,不仅易于理解,还能够轻松编写高性能的代码。对于CPU,我们生成C++代码,而对于Triton,我们生成易于理解和修改的代码。这种做法使Torch Inductor成为一个通用的编译器,可以适应各种不同的硬件后端。
Mine by Run Loop Level AR
Mine by Run Loop Level AR是Torch Inductor中一个重要的概念。在这种编译模式下,我们使用类似Python的IR来定义操作和相应的计算逻辑。下面是一个permute和add操作的IR示例:
def permute_and_add(coordinates):
x, y = ops.load(coordinates)
z = ops.add(x, y)
return z
在使用这个IR时,如果我们希望进行分析,我们可以替换ops
函数来记录loads;如果我们想要进行代码生成,我们可以替换ops
函数为打印Triton代码的函数,并执行此函数来生成代码。这种设计使得编写lowerings非常容易,特别是对于构建组合操作的情况,我们可以利用Python语言的特性来简化降低操作的编写。
编译器堆栈概述
Torch Inductor的编译器堆栈由以下几个阶段组成:
-
图形分解:使用AOT Autograd和PrimTorch将操作符集分解为更小的操作符集,约250个操作符。AOT Autograd会分别捕捉前向计算和反向计算的图形。
-
图形降低:图形降低将这些原始操作符(约250个)降低为Inductor的Loop Level IR。在这个阶段,Inductor的Loop Level IR只包含约50个操作符,大大简化了图形并减少了复杂性。
-
调度:在调度阶段,我们决定哪些操作要合并,进行内存规划和切片等其他优化。这一阶段对性能的提升非常重要。
-
代码生成:代码生成分为两部分。第一部分是后端代码生成,根据需要生成Triton代码或C++代码。第二部分是封装代码生成,这些代码用于将多个不同的内核调用连接在一起,取代了编译器堆栈中的解释器部分。
Torch Inductor在GPU上的结果
我们对Torch Inductor在GPU上的性能进行了测试,下面是一些结果:
- 对于大型实际基准测试套件,Torch Inductor在GPU上的性能提升达到了1.86倍的几何平均速度提升。我们对这个结果感到非常激动,并希望在未来进一步改进性能。
Torch Inductor在CPU上的结果
我们与Intel Pi torch团队合作,在CPU上对Torch Inductor进行了测试,下面是一些结果:
- 对于CPU推理,Torch Inductor在性能上提升了1.26倍的几何平均速度。这不仅表明Torch Inductor支持CPU,也证明了其通用性。我们不仅想构建一个只支持GPU的编译器,还希望能够适应各种不同的硬件后端,这也是我们选择使用C++和Triton的原因。
结论
通过介绍Torch Inductor的设计原则、技术、编译器堆栈和性能结果,我们可以看出Torch Inductor是一个旨在提高PyTorch性能和通用性的强大工具。它以PyTorch为核心,通过Python编写,支持动态形状和步幅,并借鉴现代编程语言Triton,使得高性能编程变得更加简单。我们对Torch Inductor的未来发展充满期待,并希望它能够为PyTorch用户带来更好的体验。
相关资源