最近（2024年3月29日）庵芭，號稱比Python快6.8萬倍的Mojo編程語言開源啦雀监！6.8萬倍？你敢相信這個(gè)數(shù)字是真的嗎好乐？不過瓦宜，就連Mojo官網(wǎng)都把這個(gè)結(jié)果貼了出來（見下圖）临庇，這就很難讓你不對這個(gè)數(shù)字引起好奇昵慌。很顯然淮蜈，Mojo官方的結(jié)果難免有“自賣自夸”的嫌疑礁芦，但至少說明在某些特殊的場景下確實(shí)得到了這個(gè)數(shù)字，官網(wǎng)不會造假肖方。那么未状，究竟是什么原因讓Mojo能比Python快這么多呢？下面我們就揭開這一神秘的面紗艰垂，也借此機(jī)會讓我們進(jìn)一步了解Mojo這門比較火的編程語言猜憎。
Mojo簡介
Mojo編程語言是由Modular公司開發(fā)的搔课，旨在為人工智能領(lǐng)域提供統(tǒng)一的編程框架爬泥。它是基于Python語法的超集，結(jié)合了Python的易用性和C語言的性能踩官，支持多核境输、向量單元和加速器單元等硬件功能嗅剖。Mojo能夠?qū)Υ罅康图堿I硬件進(jìn)行編程，模型擴(kuò)展性更強(qiáng)区匣，為開發(fā)者提供卓越的性能體驗(yàn)。Mojo的創(chuàng)始人是Chris Lattner莲绰，他是Swift語言的創(chuàng)始人姑丑，也參與了LLVM和Clang的開發(fā)。他與Google的機(jī)器學(xué)習(xí)產(chǎn)品經(jīng)理Tim Davis共同創(chuàng)立了Modular公司震肮，并在2022年推出了Mojo語言戳晌。
特殊的例子
可以猜想痴柔，這個(gè)6.8萬倍的結(jié)果是在一個(gè)特殊的例子上完成，具體來說豪嚎，它計(jì)算和繪制了Mandelbrot集谈火，就是下面的分圖案糯耍。這是一個(gè)非常簡單但是又非常耗費(fèi)計(jì)算資源的例子，測試者也給出了選擇這個(gè)作為例子的理由：
- 簡單表達(dá)：只有很少的代碼
- 純計(jì)算：曼德勃羅集沒有內(nèi)存開銷
- 容易并行
- 可向量化
所以6.8萬倍的第1個(gè)秘密就是這個(gè)計(jì)算場景非常適合發(fā)揮Mojo的所有優(yōu)勢啦租，這是經(jīng)典的以己之長比別人之短荒揣。
# 代碼示例：下面函數(shù)中z是復(fù)數(shù)
MAX_ITERS = 1000
def mandelbrot_kernel(c):
z = c
nv = 0
for i in range(MAX_ITERS):
if abs(z) > 2:
break
z = z*z + c
nv += 1
return nv
編譯語言vs解釋語言
眾所周知系任，Python是解釋型語言俩滥，性能上天然會有一些劣勢贺奠。Mojo是雖然語法上兼容Python（很多寫法上是一樣的），但卻是一個(gè)編譯語言挂据。除此之外崎逃，Mojo除了像Python一樣支持動態(tài)類型（在運(yùn)行的時(shí)候才知道變量的類型），還支持另一種靜態(tài)類型的寫法（見下面代碼示例）勒葱，當(dāng)使用靜態(tài)類型的時(shí)候編譯器可以提前對代碼做出很多針對性的優(yōu)化巴柿，提升性能篮洁。6.8萬倍的第2個(gè)秘密就是這Mojo是一門支持靜態(tài)類型的編譯語言。
fn mandelbrot_2(c: ComplexFloat64) -> Int:
var z = c
var nv = 0
for i in range(1, MAX_ITERS):
if z.squared_norm() > 4:
break
z = z.squared_add(c)
nv += 1
return nv
向量化
前面兩個(gè)秘密其實(shí)還談不上多神秘瓦阐，很容易理解和想到篷牌。我認(rèn)為接下來談到的這個(gè)才算是Mojo真正厲害的地方枷颊。
正如宣傳所說，Mojo是面向人工智能的語言信卡，人工智能計(jì)算的特點(diǎn)是什么题造？大量的向量計(jì)算界赔。于是Mojo對向量計(jì)算進(jìn)行了針對性的優(yōu)化，并且這種優(yōu)化深入到了底層硬件咐低。為此袜腥，Mojo內(nèi)置了SIMD類型。
單指令多數(shù)據(jù)（SIMD）是一種并行處理技術(shù)鲤屡，內(nèi)置于許多現(xiàn)代CPU执俩、GPU和定制加速器中。SIMD允許您一次對多個(gè)數(shù)據(jù)執(zhí)行單個(gè)操作尝丐。例如衡奥，如果您想對數(shù)組中的每個(gè)元素求平方根矮固，可以使用SIMD來并行化工作。
Mojo中的SIMD類型就是專門負(fù)責(zé)針對不同的CPU/GPU進(jìn)行這種優(yōu)化的盹兢，具體是實(shí)現(xiàn)細(xì)節(jié)在這里就不展開了守伸。在原作者測試的機(jī)器上尼摹，CPU具有512bit長的向量寄存器，這意味著CPU可以一次操作512/64=8個(gè)雙精度浮點(diǎn)數(shù)玄呛，理論上可以實(shí)現(xiàn)8x的加速和二，實(shí)測結(jié)果是實(shí)現(xiàn)了6x以上的加速儿咱。此外场晶，原作者在SIMD的基礎(chǔ)上還進(jìn)行了進(jìn)一步針對CPU的指令的優(yōu)化：現(xiàn)代 x86 系統(tǒng)具有多個(gè)融合乘加（FMA）單元诗轻，使其能夠在每個(gè)時(shí)鐘周期執(zhí)行多個(gè) FMA。這一優(yōu)化也將速度再原有基礎(chǔ)上提升了一倍多吏颖，不過這一技巧很難適用于所有的計(jì)算場景，不多討論疚俱。
鑒于Mojo內(nèi)置了SIMD數(shù)據(jù)類型呆奕，所以實(shí)現(xiàn)上面的優(yōu)化并不算復(fù)雜衬吆，這一向量化加速的技術(shù)還真是非常適合人工智能計(jì)算的場景呢逊抡。6.8萬倍的第3個(gè)秘密就是SIMD向量化加速。
多線程加速
Python實(shí)是單線程的拇勃，如果要利用多核CPU的特性還需要一些特殊的處理孝凌，很不方便胎许。Mojo是原生支持多線程的，可以很方便利用多核CPU的特性钩述。用多核CPU對比Python的單核CPU牙勘，這不是作弊嗎所禀？確實(shí)色徘，不過在這里我們先不談公平問題，先看看在Mojo中是如何實(shí)現(xiàn)多線程加速的横腿。
fn compute_row(chunk_idx:Int):
let y = chunk_size * chunk_idx
let cy = min_y + y * scale_y
@parameter
fn compute_vector[simd_width:Int](w:Int):
let cx = min_x + iota[DType.float64, simd_width]() * scale_x
output.simd_store[simd_width](Index(h,w),
mandelbrot_kernel(
ComplexSIMD[DType.float64,
simd_width](cx,cy))
vectorize[num_ports * simd_width, compute_vector](width)
# !!! 重點(diǎn)代碼在這里
with Runtime(num_cores()) as rt:
parallelize[compute_row](rt, height)
實(shí)事求是的說耿焊，在Mojo中實(shí)現(xiàn)并行確實(shí)方便的多啊器腋！無論如何還是要為這一特性點(diǎn)個(gè)贊蒂培。所以6.8萬倍的第4個(gè)秘密就是多線程并行加速榜苫。
原作者測試的機(jī)器具有88個(gè)CPU垂睬，通過這一“作弊行為”，直接將性能在原有的基礎(chǔ)上提升了 30 倍钳枕，效果那是相當(dāng)明顯??赏壹。但是你可能好奇蝌借，為什么沒有提升到88倍呢？
負(fù)載均衡和數(shù)據(jù)傾斜
我相信“負(fù)責(zé)均衡”和“數(shù)據(jù)傾斜”這兩個(gè)概念至少有一個(gè)你是比較熟悉的自晰，通俗點(diǎn)講它們都反應(yīng)了一個(gè)問題：分工不均酬荞，活都讓少數(shù)人干了瞧哟。這就是上面的例子中為什么88核CPU只實(shí)現(xiàn)了30倍加速的原因：計(jì)算在88個(gè)CPU中并不是均勻分布的勤揩。那么如何進(jìn)一步優(yōu)化呢？
方法其實(shí)也并不復(fù)雜凿傅，我們可以把任務(wù)進(jìn)一步拆分成更小的單元数苫，拆分的越細(xì)聪舒，平均分配給每個(gè)cpu之后越不容易產(chǎn)生“分工不均”的現(xiàn)象（如果你寫過Spark，應(yīng)該聽過這個(gè)最佳實(shí)踐：任務(wù)/Task的數(shù)量最好是Executor/Core數(shù)量的2-3倍虐急，而不是等于）箱残。值得慶幸的是，Mojo 包含一個(gè)高性能并發(fā)運(yùn)行時(shí)止吁，因此我們不必自己創(chuàng)建線程池或進(jìn)行循環(huán)選擇和執(zhí)行被辑。Mojo 的運(yùn)行時(shí)包含高級功能，可以充分利用像這樣的多核系統(tǒng)敬惦。
# 只需要對原來代碼做很少改動
with Runtime(num_cores()) as rt:
let partition_factor = 16 # Is autotuned.
parallelize[compute_row](rt, height, partition_factor * num_cores())
6.8萬倍的第5個(gè)秘密就是負(fù)載均衡盼理。
總結(jié)
以上每一個(gè)優(yōu)化都會使得性能提升幾倍到數(shù)十倍不等，這些數(shù)字相乘之后確實(shí)就得到了6.8萬這一聳人聽聞的數(shù)字俄删。一路學(xué)習(xí)了解下來，我覺得Mojo確實(shí)是一門相當(dāng)不錯(cuò)的編程語言畴椰，同時(shí)Mojo團(tuán)隊(duì)也是很懂營銷半铩！
關(guān)注【黑客悟理】斜脂，不錯(cuò)過任何奇奇怪怪的知識
參考資料
- https://www.modular.com/blog/how-mojo-gets-a-35-000x-speedup-over-python-part-1
- https://www.modular.com/blog/how-mojo-gets-a-35-000x-speedup-over-python-part-2
- https://www.modular.com/blog/mojo-a-journey-to-68-000x-speedup-over-python-part-3
- https://mojocn.org/
如果你喜歡我的文章抓艳，歡迎到我的個(gè)人網(wǎng)站關(guān)注我，非常感謝帚戳！
