一门号称比Python快68000倍的新型AI编程语言

见贤思齐

作者：chanceModular公司在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用户像从Python2迁移到Python3那样痛苦。幸运的是，虽然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 LTSx86-64 CPU (with SSE4.2 or newer) and a minimum of 8 GiB memoryPython 3.8 - 3.10g++ or clang++ C++ compilerAI助手搭建基础镜像安装模块化CLI，该工具类似于包管理器来安装和更新Mojocurlhttps://get.modular.com|\MODULAR_AUTH=xxxxxxxxxxxxxxx\sh-MODULAR_AUTH可在https://developer.modular.com/download注册后获取安装成功界面如下所示：安装mojosdkmodularinstallmojo安装过程中遇到了如下报错：经过排查后发现是权限问题，解决方法是加参数--cap-add=SYS_PTRACE：dockerrun--cap-add=SYS_PTRACE安装成功界面如下所示：设置环境变量echo'exportMODULAR_HOME="$HOME/.modular"'>>~/.bashrcecho'exportPATH="PATH"'>>~/.bashrcsource~/.bashrc基础命令查看Mojo版本mojo--version更新Mojo版本modularupdatemojo更新modular工具sudoaptupdatesudoaptinstallmodular方式二：ModularPlaygroundModular通过ModularPlayground提供了对Mojo的早期访问，这是一个基于网络的JupyterNotebook环境，可以在上面直接运行Mojo代码，网址是https://playground.modular.com/方式三：腾讯云CloudStudio腾讯云CloudStudio是腾讯云的面向云端开发的IDE产品。内置了Mojo镜像和官方全部Mojo示例https://ide.cloud.tencent.com/登陆后选择Mojo镜像，点击和直接可以编辑、运行，也可以按需提高运行的资源配置，使用示例如下所示：代码运行通过REPL命令行输入mojo回车后开启REPL会话输入代码后，连按两次回车就会开始运行，如下所示：运行Mojo代码文件创建代码文件hello.mojo写入以下代码保存fn main():print("Hello, chance!")AI助手运行代码mojohello.mojo运行结果如下所示：构建可执行的二进制文件构建命令：mojobuildhello.mojo-ohello运行：./hello常用基础语法下面来介绍一些常用的基础语法，总体来说还是比较易用的主函数构建Mojo程序需要一个main()函数作为程序的入口点，例如：fn main():    var x: Int = 1    x += 1    print(x)AI助手如果是构建一个Mojo的API库就不需要main函数引入python模块Mojo还不是python的完整超集，现在还只支持部分的python模块，引入方法如下所示：from python import ythonlet np = ython.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 + yz = 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 + yvar a = 1var b = 2c = add_inout(a, b)print(a)print(b)print(c)AI助手输出为：235AI助手如果希望在函数内改变传参，并且不影响函数外部的变量，可以用owned来修饰，代码示例如下：fn set_fire(owned text: String) -> String:    text += "🔥"    return textfn mojo():    let a: String = "mojo"    let b = set_fire(a)    print(a)    print(b)mojo()AI助手输出为：mojomojo🔥AI助手以上方式传参Mojo会赋值一份a传递到text，类似于c++中的值传递，会多一次拷贝的消耗，如果希望减少拷贝消耗可以在a后面加上^，即调用语句变为letb=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 = yMatrix(list(np.random.rand(M, K)), M, K)    B = yMatrix(list(np.random.rand(K, N)), K, N)    C = yMatrix(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.0018574928418138128GFLOP/smojo普通版本，后面的mojo版本都只是改变了矩阵运算函数，复用benchmarkfn 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        # revent 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 ython")AI助手运行结果为3.0032286709145626GFLOP/s，是python版本的1616倍mojovectorization（向量化计算）加速版本，使用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.56889670260691GFLOP/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.339894628945956GFLOP/s，是python版本的29792倍大模型加速效果测评跑了一下llama2的15M模型对比速度差异，具体数据如下：python版本速度为0.56token/smojo版本速度为322.37token/s由整体实验的加速效果来看官方宣称的68000倍肯定是有些许夸大的，这个68000是对于特定程序在特定环境下的最大加速效果，一般代码优化后是达不到那么大的加速的，但是相比于python来说确实加速了不少，而且mojo也还在起步阶段，如果它真能达到它所畅想的目标，那还是很有前景的。