揭开AI编程语言Mojo比Pyhon快6.8万倍的5个秘密!

20240403133200

最近(2024年3月29日),号称比Python快6.8万倍的Mojo编程语言开源啦!6.8万倍?你敢相信这个数字是真的吗?不过,就连Mojo官网都把这个结果贴了出来(见下图),这就很难让你不对这个数字引起好奇。很显然,Mojo官方的结果难免有“自卖自夸”的嫌疑,但至少说明在某些特殊的场景下确实得到了这个数字,官网不会造假。那么,究竟是什么原因让Mojo能比Python快这么多呢?下面我们就揭开这一神秘的面纱,也借此机会让我们进一步了解Mojo这门比较火的编程语言。

20240403200358

Mojo简介

Mojo编程语言是由Modular公司开发的,旨在为人工智能领域提供统一的编程框架。它是基于Python语法的超集,结合了Python的易用性和C语言的性能,支持多核向量单元加速器单元等硬件功能。Mojo能够对大量低级AI硬件进行编程,模型扩展性更强,为开发者提供卓越的性能体验。Mojo的创始人是Chris Lattner,他是Swift语言的创始人,也参与了LLVM和Clang的开发。他与Google的机器学习产品经理Tim Davis共同创立了Modular公司,并在2022年推出了Mojo语言。

特殊的例子

可以猜想,这个6.8万倍的结果是在一个特殊的例子上完成,具体来说,它计算和绘制了Mandelbrot集,就是下面的分图案。这是一个非常简单但是又非常耗费计算资源的例子,测试者也给出了选择这个作为例子的理由:

  • 简单表达:只有很少的代码
  • 纯计算:曼德勃罗集没有内存开销
  • 容易并行
  • 可向量化

所以6.8万倍的第1个秘密就是这个计算场景非常适合发挥Mojo的所有优势,这是经典的以己之长比别人之短

20240403201325
# 代码示例:下面函数中z是复数
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(很多写法上是一样的),但却是一个编译语言。除此之外,Mojo除了像Python一样支持动态类型(在运行的时候才知道变量的类型),还支持另一种静态类型的写法(见下面代码示例),当使用静态类型的时候编译器可以提前对代码做出很多针对性的优化,提升性能。6.8万倍的第2个秘密就是这Mojo是一门支持静态类型的编译语言

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

向量化

前面两个秘密其实还谈不上多神秘,很容易理解和想到。我认为接下来谈到的这个才算是Mojo真正厉害的地方。

正如宣传所说,Mojo是面向人工智能的语言,人工智能计算的特点是什么?大量的向量计算。于是Mojo对向量计算进行了针对性的优化,并且这种优化深入到了底层硬件。为此,Mojo内置了SIMD类型。

单指令多数据(SIMD)是一种并行处理技术,内置于许多现代CPU、GPU和定制加速器中。SIMD允许您一次对多个数据执行单个操作。例如,如果您想对数组中的每个元素求平方根,可以使用SIMD来并行化工作。

Mojo中的SIMD类型就是专门负责针对不同的CPU/GPU进行这种优化的,具体是实现细节在这里就不展开了。在原作者测试的机器上,CPU具有512bit长的向量寄存器,这意味着CPU可以一次操作512/64=8个双精度浮点数,理论上可以实现8x的加速,实测结果是实现了6x以上的加速。此外,原作者在SIMD的基础上还进行了进一步针对CPU的指令的优化:现代 x86 系统具有多个融合乘加(FMA)单元,使其能够在每个时钟周期执行多个 FMA。这一优化也将速度再原有基础上提升了一倍多,不过这一技巧很难适用于所有的计算场景,不多讨论。

鉴于Mojo内置了SIMD数据类型,所以实现上面的优化并不算复杂,这一向量化加速的技术还真是非常适合人工智能计算的场景呢。6.8万倍的第3个秘密就是SIMD向量化加速

多线程加速

Python实是单线程的,如果要利用多核CPU的特性还需要一些特殊的处理,很不方便。Mojo是原生支持多线程的,可以很方便利用多核CPU的特性。用多核CPU对比Python的单核CPU,这不是作弊吗?确实,不过在这里我们先不谈公平问题,先看看在Mojo中是如何实现多线程加速的。

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)

  # !!! 重点代码在这里
  with Runtime(num_cores()) as rt:
    parallelize[compute_row](rt, height)

实事求是的说,在Mojo中实现并行确实方便的多啊!无论如何还是要为这一特性点个赞。所以6.8万倍的第4个秘密就是多线程并行加速

原作者测试的机器具有88个CPU,通过这一“作弊行为”,直接将性能在原有的基础上提升了 30 倍,效果那是相当明显😀。但是你可能好奇,为什么没有提升到88倍呢?

负载均衡和数据倾斜

我相信“负责均衡”和“数据倾斜”这两个概念至少有一个你是比较熟悉的,通俗点讲它们都反应了一个问题:分工不均,活都让少数人干了。这就是上面的例子中为什么88核CPU只实现了30倍加速的原因:计算在88个CPU中并不是均匀分布的。那么如何进一步优化呢?

方法其实也并不复杂,我们可以把任务进一步拆分成更小的单元,拆分的越细,平均分配给每个cpu之后越不容易产生“分工不均”的现象(如果你写过Spark,应该听过这个最佳实践:任务/Task的数量最好是Executor/Core数量的2-3倍,而不是等于)。值得庆幸的是,Mojo 包含一个高性能并发运行时,因此我们不必自己创建线程池或进行循环选择和执行。Mojo 的运行时包含高级功能,可以充分利用像这样的多核系统。

# 只需要对原来代码做很少改动
with Runtime(num_cores()) as rt:
    let partition_factor = 16 # Is autotuned.
    parallelize[compute_row](rt, height, partition_factor * num_cores())

6.8万倍的第5个秘密就是负载均衡

总结

以上每一个优化都会使得性能提升几倍到数十倍不等,这些数字相乘之后确实就得到了6.8万这一耸人听闻的数字。一路学习了解下来,我觉得Mojo确实是一门相当不错的编程语言,同时Mojo团队也是很懂营销啊!

关注【黑客悟理】,不错过任何奇奇怪怪的知识

参考资料

  1. https://www.modular.com/blog/how-mojo-gets-a-35-000x-speedup-over-python-part-1
  2. https://www.modular.com/blog/how-mojo-gets-a-35-000x-speedup-over-python-part-2
  3. https://www.modular.com/blog/mojo-a-journey-to-68-000x-speedup-over-python-part-3
  4. https://mojocn.org/

如果你喜欢我的文章,欢迎到我的个人网站关注我,非常感谢!

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 203,362评论 5 477
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,330评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,247评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,560评论 1 273
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,580评论 5 365
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,569评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,929评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,587评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,840评论 1 297
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,596评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,678评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,366评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,945评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,929评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,165评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 43,271评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,403评论 2 342

推荐阅读更多精彩内容