作者：chance

Modular 公司在 9 月正式對(duì)外發(fā)布了 Mojo，這是一門面向 AI 領(lǐng)域的新型編程語(yǔ)言，號(hào)稱比 python 快 68000 倍，而且會(huì)“著火”，真有那么猛嗎？跟隨著這篇文章咱來(lái)一探究竟......

為何而來(lái)

在當(dāng)前場(chǎng)景中構(gòu)建統(tǒng)一的統(tǒng)一全球機(jī)器學(xué)習(xí)和人工智能基礎(chǔ)設(shè)施的平臺(tái)時(shí)，整個(gè)技術(shù)棧上的編程過(guò)于復(fù)雜，需要一種創(chuàng)新且可擴(kuò)展的編程模型，能夠針對(duì)加速器和其他在人工智能領(lǐng)域中普遍存在的異構(gòu)系統(tǒng)進(jìn)行編程。這意味著需要一種具有強(qiáng)大的編譯時(shí)元編程能力、集成自適應(yīng)編譯技術(shù)、在整個(gè)編譯流程中具有緩存等特性的編程語(yǔ)言，而這些特性在現(xiàn)有語(yǔ)言中并不支持。

盡管加速器很重要，但最常見且有時(shí)被忽視的“加速器”之一是主機(jī) CPU?，F(xiàn)如今，CPU 擁有許多類似張量核心的加速器模塊和其他 AI 加速單元，但它們也用作處理專用加速器無(wú)法處理的運(yùn)算，例如數(shù)據(jù)加載、前后處理以及與外部系統(tǒng)的集成。因此，很明顯，不能僅僅通過(guò)一種僅適用于特定處理器的“加速器語(yǔ)言”來(lái)推動(dòng) AI 的發(fā)展。

為了解決以上這些問(wèn)題，Mojo 誕生了，開發(fā)者希望用一種語(yǔ)言來(lái)一統(tǒng) AI 的江湖，這種語(yǔ)言需要兼顧 Python 的易用性和 Rust、C++的性能。

面向下一代編譯技術(shù)的語(yǔ)言

當(dāng)意識(shí)到?jīng)]有現(xiàn)有的語(yǔ)言能夠解決人工智能計(jì)算中的挑戰(zhàn)時(shí)，官方開始從頭重新思考如何設(shè)計(jì)和實(shí)現(xiàn)一種編程語(yǔ)言來(lái)解決這些問(wèn)題。由于需要對(duì)各種加速器提供高性能支持，傳統(tǒng)的編譯器技術(shù)如 LLVM 和 GCC 并不適用（基于它們的任何語(yǔ)言和工具都無(wú)法滿足要求）。盡管它們支持各種 CPU 和一些常用的 GPU，但這些編譯器技術(shù)是幾十年前設(shè)計(jì)的，無(wú)法完全支持現(xiàn)代芯片架構(gòu)。如今，專用機(jī)器學(xué)習(xí)加速器的標(biāo)準(zhǔn)技術(shù)是 MLIR。

MLIR 是一個(gè)相對(duì)較新的開源編譯器基礎(chǔ)設(shè)施，最初由 Google 發(fā)起（其負(fù)責(zé)人后來(lái)加入了 Modular），已經(jīng)在機(jī)器學(xué)習(xí)加速器社區(qū)廣泛采用。MLIR 的優(yōu)勢(shì)在于能夠構(gòu)建特定領(lǐng)域的編譯器，特別是對(duì)于那些不是傳統(tǒng) CPU 和 GPU 的奇特領(lǐng)域，如人工智能 ASIC、量子計(jì)算系統(tǒng)、FPGA 和定制芯片。

考慮到 Modular 中構(gòu)建下一代人工智能平臺(tái)的目標(biāo)，已經(jīng)在一些基礎(chǔ)設(shè)施中使用了 MLIR，但是沒(méi)有一種編程語(yǔ)言能夠充分發(fā)揮 MLIR 在整個(gè)技術(shù)棧中的潛力。雖然現(xiàn)在許多其他項(xiàng)目都在使用 MLIR，但 Mojo 是第一個(gè)專門為 MLIR 設(shè)計(jì)的重要語(yǔ)言，這使得 Mojo 在編寫面向 AI 工作負(fù)載的系統(tǒng)級(jí)代碼時(shí)具有獨(dú)特的強(qiáng)大能力。

一個(gè) Python 語(yǔ)言家族的成員

Mojo 的核心使命包括創(chuàng)新編譯器內(nèi)部和對(duì)當(dāng)前和新興加速器的支持，但官方并不認(rèn)為有必要在語(yǔ)法或社區(qū)方面進(jìn)行創(chuàng)新。因此，官方選擇擁抱 Python 生態(tài)系統(tǒng)，因?yàn)樗粡V泛使用，深受人工智能生態(tài)系統(tǒng)的喜愛，并且它是一種非常好的語(yǔ)言。

Mojo 語(yǔ)言有著遠(yuǎn)大的目標(biāo)：官方希望與 Python 生態(tài)系統(tǒng)完全兼容，希望具有可預(yù)測(cè)的低級(jí)性能和低級(jí)控制，并且需要能夠?qū)⒉糠执a部署到加速器上。此外，官方不希望創(chuàng)建一個(gè)碎片化的軟件生態(tài)系統(tǒng)，不希望采用 Mojo 的 Python 用戶像從 Python 2 遷移到 Python 3 那樣痛苦。

幸運(yùn)的是，雖然 Mojo 是一個(gè)全新的代碼庫(kù)，但在概念上官方并非從零開始。擁抱 Python 極大地簡(jiǎn)化了整體的設(shè)計(jì)工作，因?yàn)榇蟛糠终Z(yǔ)法已經(jīng)規(guī)定好了。官方可以將精力集中在構(gòu)建 Mojo 的編譯模型和系統(tǒng)級(jí)編程特性上。官方還從其他語(yǔ)言（如 Rust、Swift、Julia、Zig、Nim 等）以及以前將開發(fā)人員遷移到新編譯器和語(yǔ)言的經(jīng)驗(yàn)中獲益，并利用現(xiàn)有的 MLIR 編譯器生態(tài)系統(tǒng)。

此外，官方?jīng)Q定 Mojo 的長(zhǎng)期目標(biāo)是提供 Python 的超集（即使 Mojo 與現(xiàn)有的 Python 程序兼容），并擁抱 CPython 實(shí)現(xiàn)以支持長(zhǎng)尾生態(tài)系統(tǒng)。如果你是 Python 程序員，官方希望 Mojo 會(huì)讓用戶感到非常容易上手，同時(shí)還提供了開發(fā)安全和高性能系統(tǒng)級(jí)代碼的新工具，否則這些代碼可能需要在 Python 下使用 C 和 C++。

官方并不試圖去證明靜態(tài)是最好的或動(dòng)態(tài)是最好的。相反，官方相信在正確的應(yīng)用場(chǎng)景下，兩者都是好的，因此 Mojo 讓開發(fā)者來(lái)決定何時(shí)使用靜態(tài)或動(dòng)態(tài)。

與 Python 的兼容性

官方計(jì)劃與 Python 生態(tài)系統(tǒng)實(shí)現(xiàn)完全兼容，但實(shí)際上有兩種類型的兼容性，以下是目前在這兩個(gè)方面的情況：

1. 使用 CPython 來(lái)運(yùn)行 Python 代碼，支持導(dǎo)入現(xiàn)有的 Python 模塊并在 Mojo 程序中使用它們，以兼容整個(gè) Python 生態(tài)系統(tǒng)，但是這種方式無(wú)法發(fā)揮 Mojo 的優(yōu)勢(shì)，好處是整個(gè)生態(tài)系統(tǒng)的存在和可用性能夠加速 Mojo 的開發(fā)。
2. 就將任何 Python 代碼遷移到 Mojo 的能力而言，它目前還不是完全兼容的。Mojo 已經(jīng)支持了許多 Python 的核心特性，包括 async/await、錯(cuò)誤處理、可變參數(shù)等等。然而，Mojo 仍然年輕，缺少許多 Python 的其他特性。
3. 
   

開發(fā)環(huán)境配置

方式一：本地搭建

所需環(huán)境：

Ubuntu 20.04/22.04 LTSx86-64 CPU (with SSE4.2 or newer) and a minimum of 8 GiB memoryPython 3.8 - 3.10g++ or clang++ C++ compilerAI助手

1. 搭建基礎(chǔ)鏡像
2. 安裝模塊化 CLI，該工具類似于包管理器來(lái)安裝和更新 Mojo

curl https://get.modular.com | \

MODULAR_AUTH=xxxxxxxxxxxxxxx \

sh -

MODULAR_AUTH 可在
https://developer.modular.com/download 注冊(cè)后獲取

安裝成功界面如下所示：

1. 安裝 mojo sdk

modular install mojo

安裝過(guò)程中遇到了如下報(bào)錯(cuò)：

經(jīng)過(guò)排查后發(fā)現(xiàn)是權(quán)限問(wèn)題，解決方法是加參數(shù)--cap-add=SYS_PTRACE：

docker run --cap-add=SYS_PTRACE

安裝成功界面如下所示：

1. 設(shè)置環(huán)境變量

echo 'export MODULAR_HOME='$HOME/.modular'' >> ~/.bashrc

echo 'export PATH='PATH'' >> ~/.bashrc

source ~/.bashrc

1. 基礎(chǔ)命令

• 查看 Mojo 版本

mojo --version

• 更新 Mojo 版本

modular update mojo

• 更新 modular 工具

sudo apt update

sudo apt install modular

方式二：Modular Playground

Modular 通過(guò) Modular Playground 提供了對(duì) Mojo 的早期訪問(wèn)，這是一個(gè)基于網(wǎng)絡(luò)的 Jupyter Notebook 環(huán)境，可以在上面直接運(yùn)行 Mojo 代碼，網(wǎng)址是

https://playground.modular.com/

方式三：騰訊云 Cloud Studio

騰訊云 Cloud Studio 是騰訊云的面向云端開發(fā)的 IDE 產(chǎn)品。內(nèi)置了 Mojo 鏡像和官方全部 Mojo 示例
https://ide.cloud.tencent.com/

登陸后選擇 Mojo 鏡像，點(diǎn)擊和直接可以編輯、運(yùn)行，也可以按需提高運(yùn)行的資源配置，使用示例如下所示：

代碼運(yùn)行

通過(guò) REPL

1. 命令行輸入 mojo 回車后開啟 REPL 會(huì)話
2. 輸入代碼后，連按兩次回車就會(huì)開始運(yùn)行，如下所示：

運(yùn)行 Mojo 代碼文件

1. 創(chuàng)建代碼文件 hello.mojo
2. 寫入以下代碼保存

1. 運(yùn)行代碼

mojo hello.mojo

運(yùn)行結(jié)果如下所示：

構(gòu)建可執(zhí)行的二進(jìn)制文件

1. 構(gòu)建命令：mojo build hello.mojo -o hello
2. 運(yùn)行：./hello

常用基礎(chǔ)語(yǔ)法

下面來(lái)介紹一些常用的基礎(chǔ)語(yǔ)法，總體來(lái)說(shuō)還是比較易用的

主函數(shù)

構(gòu)建 Mojo 程序需要一個(gè) main()函數(shù)作為程序的入口點(diǎn)，例如：

fn main():    var x: Int = 1    x += 1    print(x)AI助手

如果是構(gòu)建一個(gè) Mojo 的 API 庫(kù)就不需要 main 函數(shù)

引入 python 模塊

Mojo 還不是 python 的完整超集，現(xiàn)在還只支持部分的 python 模塊，引入方法如下所示：

變量

用 var 來(lái)創(chuàng)建可變值，用 let 來(lái)創(chuàng)建不可變值，聲明時(shí)變量類型省略會(huì)自動(dòng)推導(dǎo)，示例如下：

fn do_math():    let a: Int = 1    var b = 2    print(a + b)do_math()AI助手

函數(shù)參數(shù)和返回值

函數(shù)參數(shù)和返回值需要有顯示的類型標(biāo)識(shí)，以下是帶 Int 類型參數(shù)和返回 Int 類型值的例子：

函數(shù)參數(shù)可變性默認(rèn)為不可變的引用，以 borrowed 進(jìn)行修飾，類似于 c++中的常量引用，以上 add 函數(shù)等同于：

fn add(borrowed x: Int, borrowed y: Int) -> Int:    return x + yAI助手

如果希望參數(shù)可變，并且將變動(dòng)同步到函數(shù)外，類似于 c++中的引用傳參，可以用 inout 來(lái)修飾，示例代碼如下：

輸出為：

235AI助手

如果希望在函數(shù)內(nèi)改變傳參，并且不影響函數(shù)外部的變量，可以用 owned 來(lái)修飾，代碼示例如下：

輸出為：

mojomojoAI助手

以上方式傳參 Mojo 會(huì)賦值一份 a 傳遞到 text，類似于 c++中的值傳遞，會(huì)多一次拷貝的消耗，如果希望減少拷貝消耗可以在 a 后面加上^，即調(diào)用語(yǔ)句變?yōu)?let b = set_fire(a^)，這樣 a 中的值會(huì)被轉(zhuǎn)移并且不再被初始化，有點(diǎn)類似 c++中 move 操作，因此由于 a 已經(jīng)被破壞 print(a)將不能正常執(zhí)行會(huì)報(bào)錯(cuò)。

當(dāng)前所有函數(shù)返回值時(shí)都會(huì)創(chuàng)建一個(gè)副本，還沒(méi)類似于 c++中的右值引用延長(zhǎng)返回值聲明周期的操作。

struct 結(jié)構(gòu)體

Mojo 中的 struct 跟 Python 中的 class 類似：它們都支持方法、字段、運(yùn)算符重載、元編程的裝飾器等。它們的區(qū)別如下：

1. Python 類是動(dòng)態(tài)的：它們?cè)试S動(dòng)態(tài)調(diào)用，在運(yùn)行時(shí)動(dòng)態(tài)綁定實(shí)例屬性。
2. Mojo 結(jié)構(gòu)是靜態(tài)的：它們?cè)诰幾g時(shí)綁定 (你不能在運(yùn)行時(shí)添加方法)。

具體示例如下：

加速效果測(cè)評(píng)

矩陣運(yùn)算

1. 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助手

運(yùn)行結(jié)果為 0.0018574928418138128 GFLOP/s

1. mojo 普通版本，后面的 mojo 版本都只是改變了矩陣運(yùn)算函數(shù)，復(fù)用 benchmark

運(yùn)行結(jié)果為 3.0032286709145626 GFLOP/s，是 python 版本的 1616 倍

1. mojo vectorization（向量化計(jì)算）加速版本，使用 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助手

運(yùn)行結(jié)果為 20.56889670260691 GFLOP/s，是 python 的 11073 倍

以上代碼可以用內(nèi)置的向量化函數(shù)來(lái)簡(jiǎn)化，簡(jiǎn)化后代碼如下：

1. mojo 并行化版本，實(shí)用內(nèi)置的 parallelize 函數(shù)

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助手

運(yùn)行結(jié)果為 55.339894628945956 GFLOP/s，是 python 版本的 29792 倍

大模型加速效果測(cè)評(píng)

跑了一下 llama2 的 15M 模型對(duì)比速度差異，具體數(shù)據(jù)如下：

• python 版本

速度為 0.56token/s

• mojo 版本

速度為 322.37token/s

由整體實(shí)驗(yàn)的加速效果來(lái)看官方宣稱的 68000 倍肯定是有些許夸大的，這個(gè) 68000 是對(duì)于特定程序在特定環(huán)境下的最大加速效果，一般代碼優(yōu)化后是達(dá)不到那么大的加速的，但是相比于 python 來(lái)說(shuō)確實(shí)加速了不少，而且 mojo 也還在起步階段，如果它真能達(dá)到它所暢想的目標(biāo)，那還是很有前景的。