NVIDIA 初创加速计划,免费加速您的创业启动 了解详情
写点什么

Mojo 对 Rust:Mojo 真能比 Rust 还快?

  • 2024-04-01
    北京
  • 本文字数:5974 字

    阅读完需:约 20 分钟

Mojo对Rust:Mojo 真能比Rust 还快?

Mojo 基于 MLIR 中最新的编译器技术构成而成,所谓 MLIR 则是 LLVM 的演变产物(与 Rust 同样关注底层),因此速度表现更好。只要程序员技术水平达标,又有充分优化的意愿,也完全可以让代码跑得飞快。Mojo 语言的目标在于既满足 Python 开发者的需求,又提供一系列新的代码优化技巧来充分发掘硬件设备的性能极限。

博文与性能基准

上周末,Netflix 工程师兼 Rust 倡导者 @ThePrimeagen 发布了一段视频:用 Mojo 以超越 Rust50%的速度解析 DNA 序列。这篇博文引发了一些争议,毕竟 Rust 被定位为 AI 领域主导语言的潜在继任者(目前话事的主要是 Python 和 C++)。下面来看 @ThePrimeagen 对于 Mojo 和 Rust 在 AI 编程领域的未来展望:


如果 Mojo 正式加入战局,那我相信 Mojo 无疑将最终胜出。Mojo 取胜的原因,在于无需改变任何开发者已经熟知的范式。只需要稍加学习,就能获得惊人的性能表现。首先 Mojo 的编译速度很快,而且使用感受跟大家已经熟悉的语言非常接近,性能也跟 Rust 不相上下。唯一的问题就是怎么让更多人接受它。


在发表评论后,在业界颇有声望的 Rust 贡献者兼《Rust:从入门到生产(Zero to Production in Rust)》一书的作者 Luca Palmieri 在 X 上回应称:



昨天在 Mojo vs Rust 上看到了 @ThePrimeagen 的直播:他说的没错。如果 Mojo 能够全面落地,那么对于从事 AI 工作的朋友们来说,我们再也不用在“userspace”里看到 Rust 了。Mojo 的价值主张对于熟悉 Python 的机器学习工程师和数据科学家们都是种福音。

Mojo:我们的目标

Mojo 的目标是让 Python 开发者能够直观轻松地加以掌握。正如 Mohamed 所展示,他在几周之内就以业余项目的形式学会了 Mojo 并使用到 SIMD 优化算法。虽然本文将着重讨论性能层面的差异,但 @ThePrimeagen 和 Luca Palmieri 提出的观点同样非常重要。对于关注 AI 开发的朋友们来说,目前确实存在三种语言选择其一的难题,而且硬件层面的 CPU+GPU 可编程性也极为关键。需要强调的是,Mojo 真正的目标是进一步增强当前全球最流行的 AI 语言 Python,并为世界各地的开发人员提供难以置信的性能表现、硬件可移植性与可编程性。

Mojo 真比其他语言更快吗?

@ThePrimeagen 提出了一个重要问题:Rust 向来以强大的底层性能而闻名,那 Mojo 要如何实现比 Rust(和 C++)更好的开箱即用性能?


新人用户在初次加入 Discord 时,最常提出的问题就是 Mojo 能比某其他语言快多少。其实任何基准测试的结果都会受到各种因素的影响,我们不可能单凭一项测试结果就认定某语言比另一种语言更快。更科学的提法,应该是与某语言相比,Mojo 相应的开销是多少。Mojo 的一大核心目标就是帮助开发者发掘硬件设备上的性能极限,同时以符合人体工学的方式为 Python 开发者提供熟悉的使用感受。


与 Python 等动态语言相比,编译语言允许开发者去除不必要的 CPU 指令,例如将对象分配到堆、引用计数和定期垃圾收集等。Mojo 从 C++、Rust 和 Swift 中汲取了经验教训与最佳实践,可通过直接访问机器的方式回避此类开销。

Mojo 对决 Rust

Mojo 与 Rust 都允许开发者在较低层级进行优化。以 Rust 为例,大家当然也可以把所有内容都打包在 Arc、Mutex 或者 Box 当中,以避免与借用检查器发生冲突,但这也会相应牺牲掉一部分性能。如果我们是在编写应用程序代码,这点性能差异可能没什么重大影响;但对于库或者其他性能敏感的代码,由此带来的开销可能会迅速增加。具体如何选择,取决于程序员对减少开销和优化性能的关注程度。


这两种语言都可以使用 LLVM 来优化代码生成,也都允许使用内联汇编(当然,相信没人会真这么做),所以理论上二者在传统硬件上的性能潜力基本相当。


但真正的问题在于:在不清楚编译器工作细节的情况下,惯用/常规 Mojo 代码的性能跟非汇编语言专家编写的普通 Rust 代码相比,到底孰优孰劣?

默认情况下靠借用减少 memcpy

在新用户学习 Rust 时,遇到的第一个陷阱往往是函数参数默认通过移动来获取对象。也就是说当我们将某些内容传递给函数并尝试重用时,会收到编译器错误提示:


Rust


fn bar(foo: String){}

fn main(){

let foo = String::from("bar");

bar(foo);

dbg!(foo);

}


Output


5 | let foo = String::from("bar");

| --- move occurs because `foo` has type `String`, which does not implement the `Copy` trait

6 | bar(foo);

| --- value moved here

7 | dbg!(foo);

| ^^^^^^^^^ value used here after move


dbg!这行会引发编译器错误,因为我们已经将 foo 移至 bar 函数当中。在 Rust 这边,这意味着 foo 会对字符串指针、大小和容量进行 memcpy。在某些情况下,memcpy 可以被 LLVM 优化掉,但也并非永远如此。而且除非大家明确了解 Rust/LLVM 编译器的工作方式,否则实际情况将难以预测。


Mojo 则针对标准用例简化了这一概念:


Mojo


# foo is an immutable reference by default

fn bar(foo: String):

pass

fn main():

let foo = String("foo")

bar(foo)

print(foo)


Output


foo


默认情况下只会借用 Mojo 参数:与 Rust 相比,Mojo 的学习曲线不仅更平缓,而且由于不存在隐式 memcpy,其效率也会更高。如果想要实现类似 Rust 的行为,则可以将参数更改为 owned:


Mojo


fn bar(owned foo: String):

foo += "bar"

fn main():

let foo = String("foo")

bar(foo)

print(foo)


Output


foo


这仍然有效!因为 String 实现了一个复制构造函数,所以它可以被移动至 bar 中并留下一个副本。从底层看,整个过程仍在通过引用传递以获取最大效率,且保证只在 foo 发生变化时才创建相应副本。


要完全重现 Rust 默认的移动对象并失去所有权,大家需要使用^传递运算符:


Mojo


fn bar(owned foo: String):

foo += "bar" # Ok to mutate a uniquely owned value

fn main():

let foo = String("foo")

bar(foo ^)

print(foo) # error: foo is uninit because it was transferred above


折腾到这里,我们终于在移动后尝试使用 foo 时触发了编译器错误——没错,想骗过 Mojo 中的借用检查器还真不轻松!这样的默认设计明显更好,不仅效率更高,而且不会妨碍拥有动态编程背景的工程师。默认情况下,他们仍然会获得预期行为,同时尽量提升代码的性能表现。

无需使用 Pin

在 Rust 当中,不存在值同一性的概念。对于指向其自身成员的自引用结构,一旦对象移动,则该数据可能会因继续指向内存中的旧位置而变得无效。这会造成复杂性激增,特别是在异步 Rust 的部分,其中 future 需要自我引用并存储状态,因此必须用 Pin 打包 Self 以保证它不会移动。但在 Mojo 这边,对象带有一个标识,因此引用 self.foo 将始终返回内存中的正确位置,无需程序员承担任何额外复杂性。总之,Mojo 在设计上帮助程序员回避掉了很多复杂因素。

基于最先进的编译器技术

Rust 于 2006 年启动,Swift 则诞生于 2010 年,二者主要构建在 LLVM IR 之上。Mojo 则亮相于 2022 年,基于 MLIR 构建而成——MLIR 是比 Rust 上使用的 LLVM IR 方法更加现代的“下一代”编译器堆栈。这里还有一段历史:我们的 CEO Chris Lattner 于 2000 年 12 月在大学里创立了 LLVM,并从其多年来的演变和发展中学到了很多。他随后加入谷歌领导 MLIR 的开发,旨在支持公司的 TPU 及其他 AI 加速器项目。接下来,他继续利用从 LLVM IR 中学到的知识开启了下一步探索。


Mojo 是首个充分利用到 MLIR 先进特性的编程语言,既可以生成优化度更高的 CPU 代码,也能支持 GPU 和其他加速器,而且统计速度也比 Rust 快得多。这是目前其他语言无法实现的优势,也是 AI 和编译器爱好者们痴迷 Mojo 的核心原因。他们能够针对奇特的硬件建立起奇特的抽象,而我们普通开发者则可以通过 Python 式的语法轻松加以使用。

出色的 SIMD 人体工学设计

CPU 通过特殊的寄存器与指令来同时处理多位数据,这就是 SIMD(单指令、多数据)。但从历史上看,此类代码的编写体验在人体工学层面来看非常丑陋且难以使用。这些特殊指令已经存在多年,但大多数代码仍未针对其进行过优化。所以谁能解决这种复杂性并编写出可移植的 SIMD 优化算法,谁就能在市场上脱颖而出,例如 simd_json。


Mojo 的原语在设计之初就考虑到了 SIMD 优先:UInt8 实际上是一个 SIMD[DType.uint8, 1],即 1 元素的 SIMD。以这种方式表示它不会产生性能开销,同时允许程序员轻松将其用于 SIMD 优化。例如,我们可以将文本拆分成 64 字节块,将其表示为 SIMD[DType.uint8, 64],再将其与单个换行符进行比较,从而找到每个换行符的索引。由于机器上的 SIMD 寄存器可以同时计算 512 位数据的运算,因此这种操作就能将此类运算的性能提高 64 倍!


或者举个更简单的例子,假设大家有一个 SIMD[DType.float64, 8](2, 4, 6, 8, 16, 32, 64, 128),那么只需简单将其乘以 Float64(2),就能轻松提高性能。与单独将每个元素相乘比较,这种方法在大多数机器上能够将性能提高 8 倍。


LLVM(也就是 Rust)具有自动向量化优化通道,但由于无法更改 SIMD 的内存布局和其他重要细节,所以其性能表现永远达不到理论层面的开发优化极限。但 Mojo 在设计之初就考虑到 SIMD 特性,因此编写 SIMD 优化的感受与编写普通代码非常相似。

Eager Destruction 急切销毁

Rust 的设计灵感来自 C++的 RAII(资源获取即初始化),就是说一旦对象超出范围,应用程序开发者不必分心释放内存,编程语言本身会自行处理。这是个非常好的范例,能在保障动态语言人体工学的前提下回避垃圾收集机制带来的性能缺陷。


Mojo 则更进一步,它不会等待末尾作用域,而在最后一次使用对象时释放内存。这对 AI 场景非常有利,因为提前释放对象意味着提前释放 GPU 张量,因此可以在等量 GPU RAM 中拟合更大的模型。这是 Mojo 的独特优势,程序员无需费心设计即可获得最佳性能。Rust 借用检查器最初会将全部内容的生命周期延长至其作用域的末尾,借此匹配解构函数(destructor)的行为,但这会给用户带来一些令人困惑的后果。Rust 随后添加了一些非词汇生命周期功能以简化开发者的工作。但凭借 Mojo 中的急切销毁(eager destructor)机制,这种简化效果可以直接实现,而且其与对象的实际销毁方式保持一致,因此不会引发令人难以理解的极端情况。


Rust 中的另一种开销来自 Drop 的实现方式。它使用 Drop Flags 标记跟踪是否应该在运行时删除对象。Rust 在某些情况下能够实现优化,但 Mojo 可通过明确定义消除一切情况下的额外开销。

尾调用优化 (TCO)

更新:社区讨论中指出,对于以下原始示例,Mojo 可以正确优化所有内容,而 Rust 则因存在潜在 bug 而导致实现速度慢上许多。生成的程序集还显示,Rust 会执行某种形式的 TCO,即使对于堆分配的对象也会执行。考虑到这些,我更新了以下示例并调整了本章节的具体内容。


由于 Mojo 具有急切销毁机制,因此 MLIR 和 LLVM 能够更高效地执行尾调用优化。以下示例将两种语言中的递归函数与堆分配的动态向量进行了比较。请注意,这里只是简单示例,强调以尽可能少的代码演示二者间的差异。


首先运行 cargo new rust,而后对./rust/src/main.rs 做如下编辑:

./rust/src/main.rs

fn recursive(x: usize){

if x == 0 {

return;

}

let mut stuff = Vec::with_capacity(x);

for i in 0..x {

stuff.push(i);

}

recursive(x - 1)

}

fn main() {

recursive(50_000);

}


之后运行:


Bash


cd rust

cargo build --release

cd target/release

hyperfine ./rust


在 M2 Mac 上的运行结果如下:


Output


Benchmark 1: ./rust

Time (mean ± σ): 2.119 s ± 0.031 s [User: 1.183 s, System: 0.785 s]

Range (min … max): 2.081 s … 2.172 s 10 runs


我们可以在同一文件夹中使用单一文件运行 mojo 版本,这里将其命名为 mojo.mojo:


fn recursive(x: Int):if x == 0:returnvar stuff = DynamicVectorIntfor i in range(x):stuff.push_back(i)recursive(x - 1)


fn main():recursive(50_000)


之后运行:


Bash


mojo build mojo.mojo

hyperfine ./mojo


Output


Benchmark 1: ./mojo

Time (mean ± σ): 620.6 ms ± 5.6 ms [User: 605.2 ms, System: 2.1 ms]

Range (min … max): 613.9 ms … 632.4 ms 10 runs


编译器必须在适当时机调用析构函数。对 Rust 来说,也就是在值超出范围的时候。在递归函数中,Vec 拥有一个析构函数,需要在每次函数调用后运行。也就是说该函数的堆栈帧无法如尾调用优化所需要的那样被丢弃或覆盖。而由于 Mojo 拥有急切销毁机制,因此不存在这一限制,能够通过堆分配的对象更有效地实现 TCO 优化。


使用 valgrind --tool=massif 分析两个版本的程序,能帮助我们更深入地理解此行为。这里切换至 Linux 云实例来运行本实验,在 10 GB 峰值分配内存之下,Rust 版本的平均运行时间为 0.067 秒;而在 1.5 MB 的峰值分配内存下,Mojo 版本的成绩则为 1.189 秒!如前所述,内存是 AI 应用场景下的重要资源,而急切销毁显然能帮助程序员在无需特别设计的情况下获取最佳性能。


感兴趣的朋友也可以亲自尝试运行以上基准测试。如果大家还没有部署 Mojo,可以点击此处安装(https://developer.modular.com/download)。

总结

我们对 Rust 高度赞赏,Mojo 的设计也在很大程度上其启发。Rust 拥有系统编程语言领域最出色的高级人体工学设计,但正如 @ThePrimeagen 所指出,它在 AI 应用领域存在两大问题:


  1. 编译速度慢,而 AI 特别强调实验与快速迭代;

  2. 大多数有 Python 经验的 AI 研究人员不愿花时间从零开始学习一门新语言。我们团队的成员曾试图在谷歌通过“Swift for TensorFlow”解决这个问题,但同样由于 AI 研究者不愿学习全新且编译速度较慢的语言,这套方案没能流行起来。我们也很喜欢 Python/C++/Rust/Swift/Julia 等语言,但它们都是拥有长期历史包袱的传统语言,所以轻装上阵的 Mojo 就成了应对这些古老挑战的唯一方法。


Mojo 能够为系统工程师提供最佳性能,但距离为 Python 程序员提供符合期待的所有动态功能还有很长的路要走。就当前来讲,如果大家需要开发生产级别的应用程序,那么 Rust 仍是个不错的选择。但如果各位好奇心旺盛并更多面向未来,希望掌握一门可能在未来 50 年内对 AI 发展有所助益的语言,那不妨给 Mojo 个机会!我们将逐步将各种 AI 库添加至 Mojo 附带的软件包中,努力通过更多杀手级应用向世界展示 Mojo 的卓越能力。


最后期待大家加入 Mojo 社区大家庭,相关资源链接整理如下:



原文链接:

https://www.modular.com/blog/mojo-vs-rust-is-mojo-faster-than-rust

2024-04-01 15:224422

评论

发布
暂无评论

系统化服务构建-软件工程分层

图南日晟

微服务 软件工程 架构设计

我的时间管理之路(附工具集合及使用心得)

YoungZY

App 时间管理

要和竞争对手做比较吗?

邓瑞恒Ryan

创业 战略管理

docker19.03读取NVIDIA显卡

首富手记

Docker Dockerfile

df 和 ls 命令执行夯主

首富手记

生产力

【转载】如何在团队中做好Code Review?

北纬32°

谈谈控制感(10):怎么做一个靠谱的人

史方远

职场 心理 成长

字符与编码

引花眠

计算机基础 utf-8

关于用户体验的一些思考

AR7

android 产品开发

孩子,我们在睡前一起来阅读 15 分钟的好书,让彼此都带着好的故事入眠。

叶小鍵

正确阅读 托马斯·奥本 Doug Antin 蒂·泰德罗克

短视频时代下的知识摄取

Neco.W

学习 知识管理 知识体系 短视频

Flink 完美搭档:数据存储层上的 Pravega

Apache Flink

大数据 flink 流计算 实时计算

ARTS week 1

丽子

Java开发工具与HelloWorld

编号94530

Java eclipse Hello World ! IDEA 开发工具

写给产品经理的信(2):产品设计能力怎样进阶

punkboy

产品 个人成长 产品经理 产品设计 进阶

技术工作中的颜值

N维空间的尘埃

世界那么大,你有偏见吗?

谢锐 | Frozen

创业 技术管理

我们都可能陷入经济困境

董一凡

生活

阿里的OceanBase上天了,但你还不会用Explain看SQL的查询计划吗?

Super~琪琪

MySQL 数据库 后台开发 后端

怎么控制老板不断加需求?

kimmking

已发表的技术文章-大数据方面

绝影-大数据

美国播客节目《指数视角》专访李飞飞:疫情、 AI 伦理、人才培养

神经星星

人工智能 程序员 李飞飞 硅谷 AI 伦理

如何在团队中做好Code Review

Ken

团队协作 代码审查 Code Review 代码质量

ARTS打卡 第1周

引花眠

ARTS 打卡计划

重新开始,被自己搞砸的生活

小天同学

个人感想 日常思考

实战 Java8-CompletableFuture

子路无倦

Java 多线程 java8 CompletableFuture

Java运算符实际运用

凌轩

Java 编程语言

自制操作系统

贾献华

C#刷遍Leetcode面试题系列连载(1) - 入门与工具简介

Python名人堂

C# .net 算法 LeetCode

不要抱怨,也别憋屈

孙苏勇

职场 随笔杂谈

这个名字,你不能再读错了

小天同学

历史 科普

Mojo对Rust:Mojo 真能比Rust 还快?_编程语言_Jack Clayton_InfoQ精选文章