对于AI研究员、机器学习工程师以及软件开发者而言，从代码库中榨取每一丝性能都可能是至关重要的考量。如果您是一位Python用户，想必对Python在这方面的不足有所体会。Python常被认为是一种运行速度较慢的语言，其中一个主要原因便是其全局解释器锁（Global Interpreter Lock, GIL）机制。

尽管Python存在这些固有的特性，但仍有多种方法可以缓解其性能瓶颈，尤其是在使用相对较新的Python版本时。以下是一些常见的优化策略：

• 最新发布的Python版本提供了一种无需使用GIL即可运行代码的方式。
• 可以利用高性能的第三方库，例如NumPy，来执行大量的数值计算。
• Python语言本身现在也内置了许多用于并行和并发处理的方法。

本文将探讨另一种有效的优化途径：在Python代码中调用其他高性能语言，以加速对时间敏感的代码段。具体而言，本文将深入讲解如何从Python中调用Mojo代码，从而显著提升应用程序的执行效率。

或许有读者对Mojo还不太熟悉？以下是Mojo的简要背景介绍。

Mojo是一种相对较新的系统级编程语言，由Modular Inc.公司开发。Modular Inc.是一家专注于AI基础设施的公司，由编译器编写传奇人物克里斯·拉特纳（Chris Lattner，LLVM和Swift的创始人）和前谷歌TPU负责人蒂姆·戴维斯（Tim Davis）于2022年共同创立。Mojo于2023年5月首次公开亮相。

Mojo的诞生源于Python性能不足这一痛点。它通过将Python语法的超集嫁接到基于LLVM/MLIR的编译器流水线之上，直接解决了这一问题。Mojo提供了零成本抽象、静态类型、基于所有权的内存管理、自动向量化以及为CPU和GPU加速器无缝生成代码的能力。

在Mojo发布之初演示的早期基准测试显示，在处理内核密集型工作负载时，Mojo的运行速度比纯Python快了高达35,000倍。这证明了Mojo在原始吞吐量方面可以媲美甚至超越C/CUDA的性能，同时允许开发者继续在熟悉的“Pythonic”编程环境中工作。

然而，推广任何新语言都不可避免地会遇到用户的习惯性阻力。对于习惯Python的开发者而言，直接从Python调用Mojo代码这一新功能无疑是个令人振奋的消息，因为它让开发者可以在不完全转向新语言的情况下，享受到Mojo带来的性能优势。

那么，这是否意味着可以同时获得Python的简洁性和Mojo的卓越性能呢？

为了验证这些说法，本文将编写一些示例代码。首先是纯Python版本，然后是利用NumPy进行优化的版本，最后是一个将部分计算任务卸载到Mojo模块的Python版本。通过对比这些不同版本的运行时间，将揭示Mojo在实际应用中的性能表现。

Mojo能否带来显著的性能提升？请继续阅读以揭晓答案。

环境搭建

在开发环境中，可以使用WSL2 Ubuntu for Windows。最佳实践是为每个项目设置一个独立的开发环境。通常可以选用conda进行环境管理，但鉴于uv包管理器日益普及，本文将尝试使用uv进行设置。安装uv有以下两种方式：

$ curl -LsSf https://astral.sh/uv/install.sh | sh
or...
$ pip install uv

接下来，初始化一个项目：

$ uv init mojo-test 
$ cd mojo-test
$ uv venv
$ source .venv/bin/activate
Initialized project mojo-test at /home/tom/projects/mojo-test

(mojo-test) $ cd mojo-test

(mojo-test) $ ls -al

total 28

drwxr-xr-x  3 tom tom 4096 Jun 27 09:20 .

drwxr-xr-x 15 tom tom 4096 Jun 27 09:20 ..

drwxr-xr-x  7 tom tom 4096 Jun 27 09:20 .git

-rw-r--r--  1 tom tom  109 Jun 27 09:20 .gitignore

-rw-r--r--  1 tom tom    5 Jun 27 09:20 .python-version

-rw-r--r--  1 tom tom    0 Jun 27 09:20 README.md

-rw-r--r--  1 tom tom   87 Jun 27 09:20 main.py

-rw-r--r--  1 tom tom  155 Jun 27 09:20 pyproject.toml

现在，添加所需的外部库：

(mojo-test) $ uv pip install modular numpy matplotlib

Python如何调用Mojo？

假设有一个简单的Mojo函数，它接受一个Python变量作为参数并将其值加2。例如：

# mojo_func.mojo
#
fn add_two(py_obj: PythonObject) raises -> Python
    var n = Int(py_obj)
    return n + 2

当Python尝试加载add_two时，它会寻找一个名为PyInit_add_two()的函数。在PyInit_add_two()中，必须使用PythonModuleBuilder库声明所有可从Python调用的Mojo函数和类型。因此，最终Mojo代码的结构将类似于以下形式：

from python import PythonObject
from python.bindings import PythonModuleBuilder
from os import abort
@export

fn PyInit_mojo_module() -> PythonObject:

try:

var m = PythonModuleBuilder("mojo_func")

m.def_function[add_two]("add_two", docstring="Add 2 to n")

return m.finalize()

except e:

return abort[PythonObject](String("Rrror creating Python Mojo module:", e))
fn add_two(py_obj: PythonObject) raises -> PythonObject:

var n = Int(py_obj)

n + 2

相应的Python代码需要额外的样板代码才能正确运行，如下所示：

import max.mojo.importer
import sys
sys.path.insert(0, "")
import mojo_func
print(mojo_func.add_two(5))
结果应打印 7

代码示例

对于每个示例，本文将展示三种不同版本的代码：纯Python实现、利用NumPy加速的版本，以及将部分计算卸载到Mojo模块的Python版本。

请注意，从Python调用Mojo代码仍处于早期开发阶段。API和使用体验可能会有重大变化。

示例1 — 计算曼德尔布罗特集合

第一个示例将计算并显示曼德尔布罗特集合（Mandelbrot set）。这项计算任务的计算量相当大，纯Python版本完成所需的时间将非常可观。

此示例总共需要四个文件：

1/ mandelbrot_pure_py.py

# mandelbrot_pure_py.py
def compute(width, height, max_iters):
    """使用纯Python生成曼德尔布罗特集合图像。"""
    image = [[0] * width for _ in range(height)]
    for row in range(height):
        for col in range(width):
            c = complex(-2.0 + 3.0 * col / width, -1.5 + 3.0 * row / height)
            z = 0
            n = 0
            while abs(z) <= 2 and n < max_iters:
                z = z*z + c
                n += 1
            image[row][col] = n
    return image

2/ mandelbrot_numpy.py

# mandelbrot_numpy.py
import numpy as np
def compute(width, height, max_iters):

"""使用NumPy进行向量化计算生成曼德尔布罗特集合。"""

x = np.linspace(-2.0, 1.0, width)

y = np.linspace(-1.5, 1.5, height)

c = x[:, np.newaxis] + 1j * y[np.newaxis, :]

z = np.zeros_like(c, dtype=np.complex128)

image = np.zeros(c.shape, dtype=int)
for n in range(max_iters):
    not_diverged = np.abs(z) <= 2
    image[not_diverged] = n
    z[not_diverged] = z[not_diverged]**2 + c[not_diverged]

image[np.abs(z) <= 2] = max_iters
return image.T

3/ mandelbrot_mojo.mojo

# mandelbrot_mojo.mojo 
from python import PythonObject, Python

from python.bindings import PythonModuleBuilder

from os import abort

from complex import ComplexFloat64
这是将在Mojo中快速运行的核心逻辑
fn compute_mandel_pixel(c: ComplexFloat64, max_iters: Int) -> Int:

var z = ComplexFloat64(0, 0)

var n: Int = 0

while n < max_iters:

# abs(z) > 2 与 z.norm() > 4 相同，后者更快

if z.norm() > 4.0:

break

z = z * z + c

n += 1

return n
这是Python将调用的函数
fn mandelbrot_mojo_compute(width_obj: PythonObject, height_obj: PythonObject, max_iters_obj: PythonObject) raises -> PythonObject:
var width = Int(width_obj)
var height = Int(height_obj)
var max_iters = Int(max_iters_obj)

# 将在Mojo中构建一个Python列表来返回结果
var image_list = Python.list()

for row in range(height):
    # 创建一个嵌套列表来表示2D图像
    var row_list = Python.list()
    for col in range(width):
        var c = ComplexFloat64(
            -2.0 + 3.0 * col / width,
            -1.5 + 3.0 * row / height
        )
        var n = compute_mandel_pixel(c, max_iters)
        row_list.append(n)

    image_list.append(row_list)

return image_list

这是将Mojo函数“导出”到Python的特殊函数
@export

fn PyInit_mandelbrot_mojo() -> PythonObject:

try:
    var m = PythonModuleBuilder("mandelbrot_mojo")
    m.def_function[mandelbrot_mojo_compute]("compute", "Generates a Mandelbrot set.")
    return m.finalize()
except e:
    return abort[PythonObject]("error creating mandelbrot_mojo module")

4/ main.py

这个文件将调用其他三个程序，并可以在Jupyter Notebook中绘制曼德尔布罗特图。这里只展示一次绘图结果，读者可以相信它在所有三次代码运行中都正确绘制了。

# main.py (最终版本，包含可视化)
import time

import numpy as np

import sys
import matplotlib.pyplot as plt # 现在，导入pyplot
--- Mojo 设置 ---
try:

import max.mojo.importer

except ImportError:

print("未找到Mojo导入器。请确保MODULAR_HOME和PATH已正确设置。")

sys.exit(1)
sys.path.insert(0, "")
--- 导入我们的模块 ---
import mandelbrot_pure_py

import mandelbrot_numpy

import mandelbrot_mojo
--- 可视化函数 ---
def visualize_mandelbrot(image_data, title="曼德尔布罗特集合"):

"""使用Matplotlib将曼德尔布罗特集合数据显示为图像。"""

print(f"显示图像： {title}")

plt.figure(figsize=(10, 8))

# 'hot', 'inferno' 和 'plasma' 都是很好的颜色映射

plt.imshow(image_data, cmap='hot', interpolation='bicubic')

plt.colorbar(label="迭代次数")

plt.title(title)

plt.xlabel("宽度")

plt.ylabel("高度")

plt.show()
--- 测试运行器 ---
def run_test(name, compute_func, *args):

"""一个辅助函数，用于运行和计时测试。"""

print(f"正在运行 {name} 版本...")

start_time = time.time()

# 计算函数返回图像数据

result_data = compute_func(*args)

duration = time.time() - start_time

print(f"-> {name} 版本耗时: {duration:.4f} 秒")

# 返回数据以便可视化

return result_data
if name == "main":

WIDTH, HEIGHT, MAX_ITERS = 800, 600, 5000
print("开始曼德尔布罗特性能比较...")
print("-" * 40)

# 运行纯Python测试
py_image = run_test("纯Python", mandelbrot_pure_py.compute, WIDTH, HEIGHT, MAX_ITERS)
visualize_mandelbrot(py_image, "纯Python曼德尔布罗特集合")

print("-" * 40)

# 运行NumPy测试
np_image = run_test("NumPy", mandelbrot_numpy.compute, WIDTH, HEIGHT, MAX_ITERS)
# 如果想查看绘图，请取消下面的注释
#visualize_mandelbrot(np_image, "NumPy曼德尔布罗特集合")

print("-" * 40)

# 运行Mojo测试
mojo_list_of_lists = run_test("Mojo", mandelbrot_mojo.compute, WIDTH, HEIGHT, MAX_ITERS)
# 将Mojo的列表转换为NumPy数组以便可视化
mojo_image = np.array(mojo_list_of_lists)
# 如果想查看绘图，请取消下面的注释  
#visualize_mandelbrot(mojo_image, "Mojo曼德尔布罗特集合")

print("-" * 40)
print("比较完成。")

最后，以下是运行结果：

图片来源：作者

可以看到，Mojo的性能表现令人印象深刻。它比纯Python实现快了近20倍，甚至比NumPy代码快了5倍。这无疑为Mojo的强大性能提供了一个有力的开端。

示例2 — 数值积分

在这个示例中，将使用辛普森法则（Simpson’s rule）进行数值积分，以确定sin(X)在0到π区间内的值。辛普森法则是一种计算积分近似值的方法，其定义如下：

∫ ≈ (h/3) * [f(x₀) + 4f(x₁) + 2f(x₂) + 4f(x₃) +… + 2f(xₙ-₂) + 4f(xₙ-₁) + f(xₙ)]

其中：

• h是每个步长的宽度。
• 权重序列为1, 4, 2, 4, 2,…, 4, 1。
• 第一个和最后一个点的权重为1。
• 奇数索引点的权重为4。
• 偶数索引点的权重为2。

需要计算的积分真值为2。下面将测试这些方法在计算精度和速度上的表现。

同样，此示例也需要四个文件。

1/ integration_pure_py.py

# integration_pure_py.py
import math
def compute(start, end, n):

"""使用辛普森法则计算sin(x)的定积分。"""

if n % 2 != 0:

n += 1 # 辛普森法则要求偶数个区间
h = (end - start) / n
integral = math.sin(start) + math.sin(end)

for i in range(1, n, 2):
    integral += 4 * math.sin(start + i * h)

for i in range(2, n, 2):
    integral += 2 * math.sin(start + i * h)

integral *= h / 3
return integral

2/ integration_numpy.py

# integration_numpy.py
import numpy as np
def compute(start, end, n):

"""使用NumPy计算sin(x)的定积分。"""

if n % 2 != 0:

n += 1
x = np.linspace(start, end, n + 1)
y = np.sin(x)

# 应用辛普森法则权重: 1, 4, 2, 4, ..., 2, 4, 1
integral = (y[0] + y[-1] + 4 * np.sum(y[1:-1:2]) + 2 * np.sum(y[2:-1:2]))

h = (end - start) / n
return integral * h / 3 # 返回计算结果

3/ integration_mojo.mojo

# integration_mojo.mojo
from python import PythonObject, Python
from python.bindings import PythonModuleBuilder
from os import abort
from math import sin
注意: 这里使用 'fn' 关键字，因为它与所有版本兼容。
fn compute_integral_mojo(start_obj: PythonObject, end_obj: PythonObject, n_obj: PythonObject) raises -> PythonObject:

# 桥接操作在开始时发生一次。

var start = Float64(start_obj)

var end = Float64(end_obj)

var n = Int(n_obj)
if n % 2 != 0:
    n += 1

var h = (end - start) / n

# 下面的所有计算都使用Mojo原生类型。没有Python互操作。
var integral = sin(start) + sin(end)

# 第一个循环，用于 '4 * f(x)' 项
var i_1: Int = 1
while i_1 < n:
    integral += 4 * sin(start + i_1 * h)
    i_1 += 2

# 第二个循环，用于 '2 * f(x)' 项
var i_2: Int = 2
while i_2 < n:
    integral += 2 * sin(start + i_2 * h)
    i_2 += 2

integral *= h / 3

# 桥接操作在结束时发生一次。
return Python.float(integral)

@export

fn PyInit_integration_mojo() -> PythonObject:

try:

var m = PythonModuleBuilder("integration_mojo")

m.def_functioncompute_integral_mojo

return m.finalize()

except e:

return abortPythonObject

4/ main.py

import time
import sys
import numpy as np
--- Mojo 设置 ---
try:

import max.mojo.importer

except ImportError:

print("未找到Mojo导入器。请确保您的环境已正确设置。")

sys.exit(1)

sys.path.insert(0, "")
--- 导入我们的模块 ---
import integration_pure_py

import integration_numpy

import integration_mojo
--- 测试运行器 ---
def run_test(name, compute_func, *args):

print(f"正在运行 {name} 版本...")

start_time = time.time()

result = compute_func(*args)

duration = time.time() - start_time

print(f"-> {name} 版本耗时: {duration:.4f} 秒")

print(f"   结果: {result}")
--- 主测试执行 ---
if name == "main":

# 使用非常大的步数来突出循环性能

START = 0.0

END = np.pi

NUM_STEPS = 100_000_000 # 1亿步
print(f"计算sin(x)从 {START} 到 {END:.2f} 的积分，步数为 {NUM_STEPS:,}...")
print("-" * 50)

run_test("纯Python", integration_pure_py.compute, START, END, NUM_STEPS)
print("-" * 50)
run_test("NumPy", integration_numpy.compute, START, END, NUM_STEPS)
print("-" * 50)
run_test("Mojo", integration_mojo.compute, START, END, NUM_STEPS)
print("-" * 50)
print("比较完成。")

结果如下：

Calculating integral of sin(x) from 0.0 to 3.14 with 100,000,000 steps...
--------------------------------------------------
Running Pure Python version...
-> Pure Python version took: 4.9484 seconds
   Result: 2.0000000000000346
--------------------------------------------------
Running NumPy version...
-> NumPy version took: 0.7425 seconds
   Result: 1.9999999999999998
--------------------------------------------------
Running Mojo version...
-> Mojo version took: 0.8902 seconds
   Result: 2.0000000000000346
--------------------------------------------------
Comparison complete.

有趣的是，这次NumPy代码略快于Mojo代码，并且其最终结果也更为精确。这凸显了高性能计算中的一个关键概念：向量化与JIT编译循环之间的权衡。

NumPy的优势在于其强大的向量化操作能力。它分配一大块内存，然后调用高度优化的、预编译的C代码，这些代码利用现代CPU特性（如SIMD）同时对数百万个值执行sin()函数。这种“突发处理”模式效率极高。

另一方面，Mojo将简单的while循环JIT编译成高效的机器代码。虽然这避免了NumPy所需的大量初始内存分配，但在本例这种特定场景下，NumPy向量化的原始计算能力使其略胜一筹。

示例3 — Sigmoid函数

Sigmoid函数在人工智能领域是一个重要概念，它是二元分类的基石。

Sigmoid函数，也称为逻辑函数，定义如下：

Sigmoid函数接受任意实数值输入x，并将其平滑地“压缩”到开区间(0,1)内。简而言之，无论传递给Sigmoid函数的输入是什么，它都将返回一个介于0到1之间的值。

例如：

S(-197865) = 0
S(-2) = 0.0009111
S(3) = 0.9525741
S(10776.87) = 1

这使得它非常适合表示概率等概念。

由于Python代码较为简单，可以将其直接包含在基准测试脚本中，因此此示例只需两个文件。

sigmoid_mojo.mojo

from python               import Python, PythonObject
from python.bindings      import PythonModuleBuilder
from os                   import abort
from math                 import exp
from time                 import perf_counter
----------------------------------------------------------------------
快速Mojo例程（内部不包含Python调用）
----------------------------------------------------------------------
fn sigmoid_sum(n: Int) -> (Float64, Float64):

# 确定性填充，一次性设定大小

var data = List[Float64](length = n, fill = 0.0)

for i in range(n):

data[i] = (Float64(i) / Float64(n)) * 10.0 - 5.0   # [-5, +5]
var t0: Float64 = perf_counter()
var total: Float64 = 0.0
for x in data:                       # 单个紧密循环
    total += 1.0 / (1.0 + exp(-x))
var elapsed: Float64 = perf_counter() - t0
return (total, elapsed)

----------------------------------------------------------------------
Python可见的封装器
----------------------------------------------------------------------
fn py_sigmoid_sum(n_obj: PythonObject) raises -> PythonObject:

var n: Int = Int(n_obj)                        # 验证参数

var (tot, secs) = sigmoid_sum(n)
# 最安全的容器: 构建一个Python列表 (自动装箱标量)
var out = Python.list()
out.append(tot)
out.append(secs)
return out                                     # -> PythonObject

----------------------------------------------------------------------
模块初始化器 (名称必须匹配:  PyInit_sigmoid_mojo)
----------------------------------------------------------------------
@export

fn PyInit_sigmoid_mojo() -> PythonObject:

try:

var m = PythonModuleBuilder("sigmoid_mojo")

m.def_functionpy_sigmoid_sum

return m.finalize()

except e:

# 如果有任何异常抛出，给Python一个真实的ImportError

return abortPythonObject

main.py

# bench_sigmoid.py
import time, math, numpy as np
N = 50_000_000
--------------------------- 纯Python -----------------------------------
py_data = [(i / N) * 10.0 - 5.0 for i in range(N)]

t0 = time.perf_counter()

py_total = sum(1 / (1 + math.exp(-x)) for x in py_data)

print(f"纯Python : {time.perf_counter()-t0:6.3f} s  - Σσ={py_total:,.1f}")
--------------------------- NumPy -----------------------------------------
np_data = np.linspace(-5.0, 5.0, N, dtype=np.float64)

t0 = time.perf_counter()

np_total = float(np.sum(1 / (1 + np.exp(-np_data))))

print(f"NumPy       : {time.perf_counter()-t0:6.3f} s  - Σσ={np_total:,.1f}")
--------------------------- Mojo ------------------------------------------
import max.mojo.importer          # 安装 .mojo 导入钩子

import sigmoid_mojo               # 编译并加载共享对象
mj_total, mj_secs = sigmoid_mojo.sigmoid_sum(N)

print(f"Mojo        : {mj_secs:6.3f} s  - Σσ={mj_total:,.1f}")

以下是运行输出：

$ python sigmoid_bench.py
Pure-Python :  1.847 s  - Σσ=24,999,999.5
NumPy       :  0.323 s  - Σσ=25,000,000.0
Mojo        :  0.150 s  - Σσ=24,999,999.5

输出中的Σσ=…显示了所有计算出的Sigmoid值的总和。理论上，当N趋于无穷大时，这个值应该精确等于N除以2。

但从结果可以看出，Mojo实现相较于已经很快的NumPy代码，性能提升超过2倍，并且比基础Python实现快了12倍以上。

Mojo的性能表现确实值得肯定。

总结

本文探讨了从Python中直接调用高性能Mojo代码来加速计算密集型任务的这一激动人心的新能力。Mojo是Modular公司推出的一种相对较新的系统编程语言，它承诺以熟悉的Pythonic语法实现C语言级别的性能，旨在解决Python长期存在的速度限制。

为了验证这一承诺，本文通过实现纯Python、优化后的NumPy以及Python与Mojo混合的方法，对三个计算密集型场景进行了基准测试：曼德尔布罗特集合生成、数值积分和Sigmoid函数计算。

测试结果揭示了在数据可以完全由Mojo原生类型处理的循环密集型算法中，Mojo具有显著的性能优势，能够超越纯Python甚至高度优化的NumPy代码。然而，文章也观察到，对于与NumPy向量化、预编译C函数完美契合的任务，NumPy仍能保持轻微的领先。

这项研究表明，Mojo虽然是Python加速的强大新工具，但要实现最大性能，需要深思熟虑地优化Python和Mojo运行时之间“桥接”的开销。