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

MikeLoveRust

发布于 2024-03-06 15:41:47

1030

发布于 2024-03-06 15:41:47

文章被收录于专栏：Rust语言学习交流Rust语言学习交流

[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内核来替换这些操作，并为支持的操作替换专门的cublas内核。

这种方法带来一个简单的库，性能仅受到编译器程序员创造力的限制，而不是模型程序员的限制。

示例代码：

use luminal::prelude::*;

// Setup graph and tensors
let mut cx = Graph::new();
let a = cx.tensor::<R2<3, 1>>()
    .set([[1.0], [2.0], [3.0]]);
let b = cx.tensor::<R2<1, 4>>()
    .set([[1.0, 2.0, 3.0, 4.0]]);

// Do math...
let mut c = a.matmul(b).retrieve();

// Compile and run graph
cx.compile(<(GenericCompiler, CPUCompiler)>::default(), &mut c);
cx.execute();

// Get result
println!("Result: {:?}", c);