【Rust 日报】2024-03-03 高性能深度学习库luminal

[new lib] luminal

Luminal是一个深度学习库，它使用可组合的编译器来实现高性能。

当前的机器学习库往往很庞大复杂，因为它们试图直接将高级操作映射到底层手工编写的内核上，并且专注于立刻执行（eager模式）。像PyTorch这样的库包含数十万行代码，几乎不可能由一个单一的程序员理解所有内容，更不用说进行大规模重构了。

但是，是否有必要如此复杂呢？机器学习模型往往是由几个简单的运算符构成的静态数据流图。这使得我们可以拥有一个非常简单的核心，只支持几个基本操作，并使用它们来构建复杂的神经网络。然后，我们可以编写编译器，在构建图之后修改它，根据我们运行的后端替换更有效的运算符。

Luminal将这种方法推向了极致，仅支持11种基本操作（primops）：

一元操作 - Log2、Exp2、Sin、Sqrt、Recip

二元操作 - Add、Mul、Mod、LessThan

其他操作 - SumReduce、MaxReduce、Contiguous

每个复杂操作都可以归结为这些基本操作，因此当执行 a - b 时，实际上会将 add(a, mul(b, -1)) 写入执行图中。或者当执行 a.matmul(b) 时，实际上在图上放置的是 sum_reduce(mul(reshape(a), reshape(b)))。

构建完图后，迭代式编译器可以修改图，将primops替换为更有效的操作，具体取决于它运行的设备。例如，在Nvidia卡上，会实时编写高效的Cuda内核来替换这些操作