在Graphcore IPU上运行Julia

2023-09-29 02:09

作为伦敦大学学院的研究软件开发人员，自2016年以来一直在使用Julia编程语言。我也对尝试在一些有趣的硬件上使用Julia这件事充满好奇心。因此我在Graphcore（拟未） IPUs上尝试实现运行Julia。然而，我不是一名编译器工程师。

您可以尝试在IPU上免费运行用Julia编写的代码，使用Paperspace Gradient上的一些notebook示例。

https://ipu.dev/lmEUrt

您可能也会喜欢观看我在JuliaCon 2023上关于在Graphcore IPU上实现Julia的演示^[1]。

什么是Julia？

Julia是一种现代的、动态的、通用的、经过编译的编程语言。它具有交互性（"类似于Python"），可以在REPL（交互式解释器）或类似Jupyter或Pluto的notebook中使用。Julia拥有一个运行时环境，其中包括即时编译器（JIT编译器）和垃圾收集器（GC），用于自动内存管理。

Julia主要用于数值计算，以其为基础的微分方程求解套件非常受欢迎。

Julia的主要编程范式是多重分派（multiple dispatch），用于解决十分仰仗所有参数的类型和数量的函数。

为何Julia深得人心？

选自Matthijs Cox的《我的目标受众》^[2]

便于探索和易于理解
多重分派使得代码可组合性强
用户自定义的类型与内置类型一样快速和紧凑
代码与数学密切相关
无需切换编程语言以提高性能......
......但如果需要的话，您仍然可以通过简单的外部函数接口（FFI）调用类似C的共享库
MIT许可证：免费和开源的

我的第一个采用Julia运行的IPU程序

我们开发了一个名为IPUToolkit.jl^[3]的包，用于与Poplar SDK进行接口连接：

但我们只是在Julia中编写C++代码，这些都属于常规操作......如果我们想“整活儿”呢？

什么是编译器？

（《我们可以从编译器的设计中学到什么？》^[4]一书中强调前端的编译器流水线的高级图示）

LLVM是一种流行的模块化编译框架，目前唯一存在的Julia实现便是基于LLVM编译器。

通过Julia，我们可以轻松地检查编译过程的每个阶段。

要了解有关Julia编译器的更多详细信息，请观看Valentin Churavy的演讲《编译Julia：使动态程序运行更快》^[5]。

共同因素：LLVM

事实证明，Graphcore为生成IPU的本地代码开发的Poplar编译器也是基于LLVM的。当你为IPU编译代码时，Poplar编译器执行了与上面所见的相同的流程。

这意味着Julia和Poplar编译器实际上可以使用相同的语言：LLVM IR

我们在IPUToolkit.jl中添加了使用Julia生成IPU代码的功能。总的来说，使用这个工具包具有以下目标：

与Poplar SDK进行接口连接，以在Julia中编写IPU程序
探索Julia的元编程能力，以减少IPU程序中的样板代码
利用Julia的代码生成能力，通过Poplar编译器生成IPU的本地代码，使用下图中概述的流程：

在Julia中编写IPU代码块

我们可以使用IPUToolkit.jl包在Julia中编写IPU代码块，生成LLVM IR代码，然后使用Poplar编译器将其编译成本地代码。通过LLVM进行的代码生成基于GPUCompiler.jl^[6]包，这是一个用于为专用目标生成LLVM IR代码的通用框架，尽管其历史名称中包含了GPU，但不限于GPU。

代码块中的代码具有与所有基于GPUCompiler.jl的编译模型相同的限制：

代码必须经过静态推断和编译，不允许动态派发。
不能使用需要Julia运行时的功能，尤其是垃圾收集器。
不能在运行时调用任何其他外部二进制库，例如不能调用BLAS库。

Colossus的目标和性能

LLVM IR并非完全独立于目标。此外，如果已知目标的属性正确，一些细节，例如向量化寄存器的宽度，可以更好地针对目标进行定制。在这个项目的大部分时间里，我实际上是为主机CPU生成LLVM IR，这样的方案“可行”，但并非最优。

最近，我成功地将Julia链接到了Graphcore的LLVM分支，这使我能够直接为IPU（"Colossus"目标）生成代码，但这种组合是高度实验性的，并且会导致一些意外的错误。

计算 𝜫

受到Owain Kenway在多种不同编程语言中的pi_examples^[7]的启发。

与C++程序的比较

我们可以编写一个等效的C++程序来比较性能，以检查Julia代码生成的效果。以下是这个代码块：

在这种情况下，C++和Julia代码块之间的主要性能差异是由于循环展开。一些实验表明，当两种代码都以相似的方式展开循环时，则其表现是完全相配的。这表明Julia可以成为为IPU编写代码的有效前端。

在代码块内进行性能基准测试

IPUToolkit.jl提供了一些用于快速对代码块的部分进行基准测试的宏：@ipucycles、@ipushowcycles和@ipuelapsed（后者在上面已经使用过）。

注意：这不是性能分析的替代品，性能分析仍然是一种宝贵的工具（但JIT使堆栈跟踪复杂化），然而这些宏对于快速反馈可能很有用。

基于周期计数的基准测试对于运行时间超过typemax(UInt32) = 4294967295个周期（取决于您的IPU型号，大约2-3秒）的块不可靠。您需要对您要基准测试的块是否会溢出计数器这件事心中有数。

使用外部包：StaticArrays.jl

我们也可以在代码块内使用第三方包，只要满足上述提到的要求：不进行内存分配，代码完全可推断。StaticArrays.jl^[8]允许您在堆栈上创建数组并对其进行基本的线性代数操作。涉及静态数组的代码通常也很容易被编译器推断。

StaticArrays.jl^[9]虽然并不比使用专门的Popops线性代数例程更高效，但仍然是在Julia中编写的IPU代码块中使用外部包的良好示例。

随机舍入

实数构成了一个连续集合R，但计算机中使用的有限精度数是离散集合F ⊂ R 的一部分。当计算机执行涉及浮点数的操作时，真实结果x ∈ R 会被一个 ^x∈F 的数字近似表示，通常确定性地选择为F中最近的数字：这称为“最近舍入”。

随机舍入是传统确定性舍入的替代舍入模式，它会将一个数字x ∈ R 随机舍入为结果的两个最近浮点数之一 [x]（F中的前一个数字）或 [x]（F中的后一个数字），根据以下规则：

常见的选择是P(x)=1/2，或者更有趣的是，

接下来，我们将始终讨论后者的概率函数P(x)。

（来源：《什么是随机舍入》^[10]，作者Nick Higham）

随机舍入很有用，因为操作的平均结果与数学期望结果相匹配。从统计学角度看，它保留了确定性舍入方案所丢弃的一些信息，从而平滑了由于有限精度而引起的数值舍入误差。这在使用低精度浮点数（如Float16）时尤为重要。相比之下，像最近舍入这样的确定性舍入模式引入了偏差，这种偏差在数字精度较低时更为严重。

IPU是极少数支持支持随机舍入的处理器之一。

让我们在CPU上进行一个练习，使用传统的最近舍入。我们定义一个函数来对一组数字进行简单的顺序求和，因为Julia中的求和函数使用成对求和^[11]，其精度更高。

naive_sum（具有1种方法的通用函数）

Float16(965.5)

False

现在让我们编写一个IPU程序——使用随机舍入多次计算x_sr的总和：

(879.0, 919.0)

899.90252

4.683982498528412

Float16(900.0)

true

在IPU上求解微分方程

若想在IPU上使用外部包，我们可以寻找类似的针对GPU的解决方案。如果它们已经被设计得足够通用，可以在不同的后端上工作，那么它们很有可能也可以在IPU上使用。

举个例子，DiffEqGPU.jl^[12]是SciML生态系统中的微分方程求解器套件，可以在不同类型的GPU上运行（Nvidia、AMD、Intel、Metal），但它提供了一些基本功能，我们也可以在IPU上复用它们。

自动微分

在数学中，与积分相反，计算表达式的导数是一个机械过程。通过代码中自动计算数学函数的精确导数，岂不美哉？

自动微分是一组在计算机程序中自动计算函数导数的技术。

Enzyme^[13]是一种在LLVM级别上运行的源代码转换自动微分引擎：它分析源代码的LLVM IR，并根据内置的微分规则对所有指令进行微分。它在LLVM级别上工作的事实意味着它主要是前端语言无关的，可用于结合多种语言或使用并行化框架来获取源代码的派生（《通过编译器增强实现多种并行范式的可扩展自动微分》^[14]在2022年超级计算大会上获得了最佳学生论文奖）。此外，Enzyme在编译时生成代码，它不会在目标系统上运行：它是用于IPU程序的完美候选者。

Rosenbrock 函数

这是人们最喜欢的优化问题函数实例。

rosenbrock (generic function with 2 methods)

虽然找到谷值是小意思，但得出全局的最小值还是很困难。

Rosenbrock 函数是一个多项式函数，很容易计算它的梯度：

我们可以让Enzyme计算第二个参数的偏导数：

使用Enzyme最小化函数

我们可以通过运行IPU程序，在不同的起始点上对Rosenbrock函数进行大规模网格搜索，然后绘制在满足终止条件时停止之前所需的迭代次数。具体而言，我们将使用Adam^[15]优化方法，该方法需要函数的梯度作为输入进行优化，而我们将使用Enzyme在主机系统上在编译时自动计算它。

结尾

我们取得了那些成就？

我们开发了第一个（据我们所知）非Graphcore官方支持的IPU的第三方编程模型，使用了LLVM编译器框架，这使得这一切成为可能。
Julia允许我们在IPU上使用高级语言进行通用编程（尽管仍然使用低级的Poplar功能），而不局限于机器学习领域。
提出了使用第三方包的复杂程序（例如，使用Enzyme进行微分方程求解和零运行时成本的自动微分），展示了代码重用的可能性，远远超过了在C++中可以实现的范围。
借助Julia的内省和元编程能力，我们能够简化编写IPU程序的某些方面。
我们可以使用单精度和半精度浮点数，后者包括随机舍入。
首次使用了Graphcore的LLVM分支。
总体而言，我们在整个项目中并不关心性能，主要目标是使Julia与IPU之间的基本接口正常工作。然而，我们展示了至少在π程序的具体示例中，由Julia生成的LLVM IR的性能（针对主机CPU！）与本机C++代码块的性能竞争力。

探索Julia的限制和未来的工作方向

对于在IPU上重用Julia代码存在许多限制（无运行时：无编译，无垃圾回收等等；无外部二进制库），但这些限制不特定于IPU，与其他卸载技术（如GPU）共同存在。
Julia不支持本机的8位浮点数。
只有Poplar/Poplibs库的一个子集已被打包。
我们目前使用一个轻量级的C++外壳来定义顶点➡完全定义LLVM IR中的代码块。
针对Colossus后端的目标是非常实验性的（有时优化传递优化过于强烈）➡解决与Colossus后端的集成问题。
无法找到一种与GPUArrays.jl集成的方法➡探索其他编程模型（KernelAbstracts.jl，GPUArrays.jl？）
目前无法访问tile级多线程➡是否存在表达tile级多线程的可能？

在云端尝试

请访问https://github.com/JuliaIPU/JuliaIpuDemo^[16]，按照README.md中的说明，了解如何在云端使用Paperspace免费体验在IPU上运行Julia。