一门号称比Python快68000倍的新型AI编程语言
作者：chance
Modular 公司在 9 月正式对外发布了 Mojo，这是一门面向 AI 领域的新型编程语言，号称比 python 快 68000 倍，而且会“着火”，真有那么猛吗？跟随着这篇文章咱来一探究竟......
首先来解释为什么说会着火，因为这门语言的标准文件后缀可以是.mojo或者.🔥，你没看错，就是一个emoji。AI助手
为何而来
在当前场景中构建统一的统一全球机器学习和人工智能基础设施的平台时，整个技术栈上的编程过于复杂，需要一种创新且可扩展的编程模型，能够针对加速器和其他在人工智能领域中普遍存在的异构系统进行编程。这意味着需要一种具有强大的编译时元编程能力、集成自适应编译技术、在整个编译流程中具有缓存等特性的编程语言，而这些特性在现有语言中并不支持。
尽管加速器很重要，但最常见且有时被忽视的“加速器”之一是主机 CPU。现如今，CPU 拥有许多类似张量核心的加速器模块和其他 AI 加速单元，但它们也用作处理专用加速器无法处理的运算，例如数据加载、前后处理以及与外部系统的集成。因此，很明显，不能仅仅通过一种仅适用于特定处理器的“加速器语言”来推动 AI 的发展。
为了解决以上这些问题，Mojo 诞生了，开发者希望用一种语言来一统 AI 的江湖，这种语言需要兼顾 Python 的易用性和 Rust、C++的性能。
面向下一代编译技术的语言
当意识到没有现有的语言能够解决人工智能计算中的挑战时，官方开始从头重新思考如何设计和实现一种编程语言来解决这些问题。由于需要对各种加速器提供高性能支持，传统的编译器技术如 LLVM 和 GCC 并不适用（基于它们的任何语言和工具都无法满足要求）。尽管它们支持各种 CPU 和一些常用的 GPU，但这些编译器技术是几十年前设计的，无法完全支持现代芯片架构。如今，专用机器学习加速器的标准技术是 MLIR。
MLIR 是一个相对较新的开源编译器基础设施，最初由 Google 发起（其负责人后来加入了 Modular），已经在机器学习加速器社区广泛采用。MLIR 的优势在于能够构建特定领域的编译器，特别是对于那些不是传统 CPU 和 GPU 的奇特领域，如人工智能 ASIC、量子计算系统、FPGA 和定制芯片。
考虑到 Modular 中构建下一代人工智能平台的目标，已经在一些基础设施中使用了 MLIR，但是没有一种编程语言能够充分发挥 MLIR 在整个技术栈中的潜力。虽然现在许多其他项目都在使用 MLIR，但 Mojo 是第一个专门为 MLIR 设计的重要语言，这使得 Mojo 在编写面向 AI 工作负载的系统级代码时具有独特的强大能力。
一个 Python 语言家族的成员
Mojo 的核心使命包括创新编译器内部和对当前和新兴加速器的支持，但官方并不认为有必要在语法或社区方面进行创新。因此，官方选择拥抱 Python 生态系统，因为它被广泛使用，深受人工智能生态系统的喜爱，并且它是一种非常好的语言。
Mojo 语言有着远大的目标：官方希望与 Python 生态系统完全兼容，希望具有可预测的低级性能和低级控制，并且需要能够将部分代码部署到加速器上。此外，官方不希望创建一个碎片化的软件生态系统，不希望采用 Mojo 的 Python 用户像从 Python 2 迁移到 Python 3 那样痛苦。
幸运的是，虽然 Mojo 是一个全新的代码库，但在概念上官方并非从零开始。拥抱 Python 极大地简化了整体的设计工作，因为大部分语法已经规定好了。官方可以将精力集中在构建 Mojo 的编译模型和系统级编程特性上。官方还从其他语言（如 Rust、Swift、Julia、Zig、Nim 等）以及以前将开发人员迁移到新编译器和语言的经验中获益，并利用现有的 MLIR 编译器生态系统。
此外，官方决定 Mojo 的长期目标是提供 Python 的超集（即使 Mojo 与现有的 Python 程序兼容），并拥抱 CPython 实现以支持长尾生态系统。如果你是 Python 程序员，官方希望 Mojo 会让用户感到非常容易上手，同时还提供了开发安全和高性能系统级代码的新工具，否则这些代码可能需要在 Python 下使用 C 和 C++。
官方并不试图去证明静态是最好的或动态是最好的。相反，官方相信在正确的应用场景下，两者都是好的，因此 Mojo 让开发者来决定何时使用静态或动态。
与 Python 的兼容性
官方计划与 Python 生态系统实现完全兼容，但实际上有两种类型的兼容性，以下是目前在这两个方面的情况：
使用 CPython 来运行 Python 代码，支持导入现有的 Python 模块并在 Mojo 程序中使用它们，以兼容整个 Python 生态系统，但是这种方式无法发挥 Mojo 的优势，好处是整个生态系统的存在和可用性能够加速 Mojo 的开发。 就将任何 Python 代码迁移到 Mojo 的能力而言，它目前还不是完全兼容的。Mojo 已经支持了许多 Python 的核心特性，包括 async/await、错误处理、可变参数等等。然而，Mojo 仍然年轻，缺少许多 Python 的其他特性。
开发环境配置
方式一：本地搭建
所需环境：
Ubuntu 20.04/22.04 LTS
x86-64 CPU (with SSE4.2 or newer) and a minimum of 8 GiB memory
Python 3.8 - 3.10
g++ or clang++ C++ compilerAI助手
搭建基础镜像 安装模块化 CLI，该工具类似于包管理器来安装和更新 Mojo
curl https://get.modular.com | \
MODULAR_AUTH=xxxxxxxxxxxxxxx \
sh -
MODULAR_AUTH 可在 https://developer.modular.com/download 注册后获取
安装成功界面如下所示：
安装 mojo sdk
modular install mojo
安装过程中遇到了如下报错：
经过排查后发现是权限问题，解决方法是加参数--cap-add=SYS_PTRACE：
docker run --cap-add=SYS_PTRACE
安装成功界面如下所示：
设置环境变量
echo 'export MODULAR_HOME="$HOME/.modular"' >> ~/.bashrc
echo 'export PATH="PATH"' >> ~/.bashrc
source ~/.bashrc
基础命令
查看 Mojo 版本
mojo --version
更新 Mojo 版本
modular update mojo
更新 modular 工具
sudo apt update
sudo apt install modular
方式二：Modular Playground
Modular 通过 Modular Playground 提供了对 Mojo 的早期访问，这是一个基于网络的 Jupyter Notebook 环境，可以在上面直接运行 Mojo 代码，网址是
https://playground.modular.com/
方式三：腾讯云 Cloud Studio
腾讯云 Cloud Studio 是腾讯云的面向云端开发的 IDE 产品。内置了 Mojo 镜像和官方全部 Mojo 示例 https://ide.cloud.tencent.com/
登陆后选择 Mojo 镜像，点击和直接可以编辑、运行，也可以按需提高运行的资源配置，使用示例如下所示：
代码运行
通过 REPL
命令行输入 mojo 回车后开启 REPL 会话 输入代码后，连按两次回车就会开始运行，如下所示：
运行 Mojo 代码文件
创建代码文件 hello.mojo 写入以下代码保存
fn main():
print("Hello, chance!")AI助手
运行代码
mojo hello.mojo
运行结果如下所示：
构建可执行的二进制文件
构建命令：mojo build hello.mojo -o hello 运行：./hello
常用基础语法
下面来介绍一些常用的基础语法，总体来说还是比较易用的
主函数
构建 Mojo 程序需要一个 main()函数作为程序的入口点，例如：
fn main():
var x: Int = 1
x += 1
print(x)AI助手
如果是构建一个 Mojo 的 API 库就不需要 main 函数
引入 python 模块
Mojo 还不是 python 的完整超集，现在还只支持部分的 python 模块，引入方法如下所示：
from python import Python
let np = Python.import_module("numpy")
ar = np.arange(15).reshape(3, 5)
print(ar)
print(ar.shape)AI助手
变量
用 var 来创建可变值，用 let 来创建不可变值，声明时变量类型省略会自动推导，示例如下：
fn do_math():
let a: Int = 1
var b = 2
print(a + b)
do_math()AI助手
函数参数和返回值
函数参数和返回值需要有显示的类型标识，以下是带 Int 类型参数和返回 Int 类型值的例子：
fn add(x: Int, y: Int) -> Int:
return x + y
z = add(1, 2)
print(z)AI助手
函数参数可变性默认为不可变的引用，以 borrowed 进行修饰，类似于 c++中的常量引用，以上 add 函数等同于：
fn add(borrowed x: Int, borrowed y: Int) -> Int:
return x + yAI助手
如果希望参数可变，并且将变动同步到函数外，类似于 c++中的引用传参，可以用 inout 来修饰，示例代码如下：
fn add_inout(inout x: Int, inout y: Int) -> Int:
x += 1
y += 1
return x + y
var a = 1
var b = 2
c = add_inout(a, b)
print(a)
print(b)
print(c)AI助手
输出为：
2
3
5AI助手
如果希望在函数内改变传参，并且不影响函数外部的变量，可以用 owned 来修饰，代码示例如下：
fn set_fire(owned text: String) -> String:
text += "🔥"
return text
fn mojo():
let a: String = "mojo"
let b = set_fire(a)
print(a)
print(b)
mojo()AI助手
输出为：
mojo
mojo🔥AI助手
以上方式传参 Mojo 会赋值一份 a 传递到 text，类似于 c++中的值传递，会多一次拷贝的消耗，如果希望减少拷贝消耗可以在 a 后面加上^，即调用语句变为 let b = set_fire(a^)，这样 a 中的值会被转移并且不再被初始化，有点类似 c++中 move 操作，因此由于 a 已经被破坏 print(a)将不能正常执行会报错。
当前所有函数返回值时都会创建一个副本，还没类似于 c++中的右值引用延长返回值声明周期的操作。
struct 结构体
Mojo 中的 struct 跟 Python 中的 class 类似：它们都支持方法、字段、运算符重载、元编程的装饰器等。它们的区别如下：
Python 类是动态的：它们允许动态调用，在运行时动态绑定实例属性。 Mojo 结构是静态的：它们在编译时绑定 (你不能在运行时添加方法)。
具体示例如下：
struct MyPair:
var first: Int
var second: Int
fn __init__(inout self, first: Int, second: Int):
self.first = first
self.second = second
fn dump(self):
print(self.first, self.second)
let mine = MyPair(2, 4)
mine.dump()
AI助手
加速效果测评
矩阵运算
python 版本
def matmul_python(C, A, B):
for m in range(C.rows):
for k in range(A.cols):
for n in range(C.cols):
C[m, n] += A[m, k] * B[k, n]
def benchmark_matmul_python(M, N, K):
A = PyMatrix(list(np.random.rand(M, K)), M, K)
B = PyMatrix(list(np.random.rand(K, N)), K, N)
C = PyMatrix(list(np.zeros((M, N))), M, N)
secs = timeit(lambda: matmul_python(C, A, B), number=2) / 2
gflops = ((2 * M * N * K) / secs) / 1e9
print(gflops, "GFLOP/s")
return gflopsAI助手
运行结果为 0.0018574928418138128 GFLOP/s
mojo 普通版本，后面的 mojo 版本都只是改变了矩阵运算函数，复用 benchmark
fn matmul_naive(C: Matrix, A: Matrix, B: Matrix, _rt: Runtime):
for m in range(C.rows):
for k in range(A.cols):
for n in range(C.cols):
C[m, n] += A[m, k] * B[k, n]
fn benchmark[
func: fn (Matrix, Matrix, Matrix, Runtime) -> None
](M: Int, N: Int, K: Int, base_gflops: Float64, str: String):
var C = Matrix(M, N)
C.zero()
var A = Matrix(M, K)
var B = Matrix(K, N)
with Runtime() as rt:
@always_inline
@parameter
fn test_fn():
_ = func(C, A, B, rt)
let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000
# Prevent the matrices from being freed before the benchmark run
_ = (A, B, C)
let gflops = ((2 * M * N * K) / secs) / 1e9
let speedup: Float64 = gflops / base_gflops
# print(gflops, "GFLOP/s", speedup, " speedup")
print(str)
print(gflops, "GFLOP/s <>", speedup.to_int(), "x speedup over Python")AI助手
运行结果为 3.0032286709145626 GFLOP/s，是 python 版本的 1616 倍
mojo vectorization（向量化计算）加速版本，使用 SIMD 向量类型
alias nelts = simdwidthof[DType.float32]() # The SIMD vector width.
fn matmul_vectorized_0(C: Matrix, A: Matrix, B: Matrix, _rt: Runtime):
for m in range(C.rows):
for k in range(A.cols):
for nv in range(0, C.cols, nelts):
C.store[nelts](
m, nv, C.load[nelts](m, nv) + A[m, k] * B.load[nelts](k, nv)
)
# Handle remaining elements with scalars.
for n in range(nelts * (C.cols // nelts), C.cols):
C[m, n] += A[m, k] * B[k, n]AI助手
运行结果为 20.56889670260691 GFLOP/s，是 python 的 11073 倍
以上代码可以用内置的向量化函数来简化，简化后代码如下：
fn matmul_vectorized_1(C: Matrix, A: Matrix, B: Matrix, _rt: Runtime):
for m in range(C.rows):
for k in range(A.cols):
@parameter
fn dot[nelts: Int](n: Int):
C.store[nelts](
m, n, C.load[nelts](m, n) + A[m, k] * B.load[nelts](k, n)
)
vectorize[nelts, dot](C.cols)AI助手
mojo 并行化版本，实用内置的 parallelize 函数
fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):
@parameter
fn calc_row(m: Int):
for k in range(A.cols):
@parameter
fn dot[nelts: Int](n: Int):
C.store[nelts](
m, n, C.load[nelts](m, n) + A[m, k] * B.load[nelts](k, n)
)
vectorize[nelts, dot](C.cols)
parallelize[calc_row](rt, C.rows)AI助手
运行结果为 55.339894628945956 GFLOP/s，是 python 版本的 29792 倍
大模型加速效果测评
跑了一下 llama2 的 15M 模型对比速度差异，具体数据如下：
python 版本
速度为 0.56token/s
mojo 版本
速度为 322.37token/s
由整体实验的加速效果来看官方宣称的 68000 倍肯定是有些许夸大的，这个 68000 是对于特定程序在特定环境下的最大加速效果，一般代码优化后是达不到那么大的加速的，但是相比于 python 来说确实加速了不少，而且 mojo 也还在起步阶段，如果它真能达到它所畅想的目标，那还是很有前景的。