用 Taichi 加速 Python：提速 100+ 倍！

如果你平时用 Python 做数值计算、图形模拟、粒子系统、或是想跑一些高密度数学运算，你一定遇到过一个问题：Python 太慢了。

慢得让人怀疑人生。

但很多人不知道，其实你可以继续写 Python，但是把关键部分“换个引擎”跑。这里要说的，就是 Taichi —— 一个让你写着 Python，却能跑出 C++、甚至 GPU 速度的计算框架。

Taichi 到底是干嘛的？

一句话：你写 Python，它替你把关键代码 JIT 编译成高性能机器码，还能直接跑在 GPU 上。

比如本来你 Python 要跑 10 秒，那 Taichi 可能让它只要 0.1 秒。

不是“优化算法”，是真正给你“换跑道”。

它适合：

• 大规模循环（如 10^7、10^8 级别）
• 粒子模拟、物理仿真
• 数值方法（Jacobi、FFT、PDE 离散方法等）
• 图形渲染
• 并行计算

写法不难，你只要把一个函数改成：

@ti.kernel
defcompute():
    ...

Taichi 就帮你做剩下的事情。

来个最典型的例子：把纯 Python 的循环炸掉

先看一段纯 Python 的代码，做个简单的向量相加：

import numpy as np
N = 10_000_000
a = np.ones(N, dtype=np.float32)
b = np.ones(N, dtype=np.float32)
c = np.zeros(N, dtype=np.float32)
for i in range(N):
    c[i] = a[i] + b[i]

如果你试一下，会发现跑得非常慢，循环就是 Python 的老大难。

现在用 Taichi 版本：

import taichi as ti
ti.init(arch=ti.gpu)   # 没 GPU 就用 ti.cpu
N = 10_000_000
a = ti.field(dtype=ti.f32, shape=N)
b = ti.field(dtype=ti.f32, shape=N)
c = ti.field(dtype=ti.f32, shape=N)
@ti.kernel
defadd():
for i in range(N):
        c[i] = a[i] + b[i]
add()

这段一看就是 Python，但实际上内部已经变成了 GPU 并行执行的机器码。

速度通常是 Python for 循环的几十到上百倍。

你完全没有写 C++ 或 CUDA，但 Taichi 已经帮你干了这些。

为啥能这么快？

Taichi 的核心优势有三个：

1. JIT 编译（即时编译）

你写的 @ti.kernel 内部不是跑 Python，而是被 Taichi 编译成：

• CPU SIMD 指令
• GPU kernel
• 甚至 Vulkan / CUDA 代码

1. 自动并行

Taichi 的 for 循环不是普通的 Python 循环，而是被拆成 SIMD 或 GPU 线程。

比如：

for i in range(N):

可能会变成 GPU 上 N 个线程同时执行。

1. 结构化字段（SNode）加速内存访问

Taichi 有自己的数据结构系统，可以把数据排布成适合 GPU 的格式，提高带宽利用率。

这对大规模模拟非常重要。

一个更实际例子：2D 粒子模拟（Python 写法 vs Taichi 写法）

纯 Python 版本：

# 完全跑不动
for i in range(n):
for j in range(n):
        v[i, j] += force[i, j] * dt

Taichi 版本：

@ti.kernel
defupdate():
for i, j in v:
        v[i, j] += force[i, j] * dt

这段在 GPU 上几乎是瞬间完成的，而 Python 版甚至有可能直接卡死。

PDE 模拟（比如热传导）

传统 Python 版本：

for _ in range(1000):
for i in range(1, N-1):
for j in range(1, N-1):
            nxt[i, j] = (cur[i-1, j] + cur[i+1, j] +
                         cur[i, j-1] + cur[i, j+1]) * 0.25

Taichi：

@ti.kernel
defdiffuse():
for i, j in cur:
if0 < i < N-1and0 < j < N-1:
            nxt[i, j] = (
                cur[i - 1, j] + cur[i + 1, j] +
                cur[i, j - 1] + cur[i, j + 1]
            ) * 0.25

Taichi 在 GPU 上直接把双层循环并行化，速度会非常恐怖。

小心一个关键点：Taichi 不是万能加速工具

Taichi 特别适合：

• 大规模循环
• 数值模拟
• 图形、粒子、物理
• 并行计算

但是 并不是所有 Python 代码都能加速。

比如：

❌ 网络 IO ❌ 字符串操作 ❌ 高层业务逻辑 ❌ Pandas / Django 这种偏数据/业务类代码

Taichi 是为 计算密集型代码 准备的，而不是业务框架。

什么时候应该考虑上 Taichi？

如果你遇到这些情境：

• “为什么我一个 1000×1000 的二维循环能跑半分钟？”
• “物理模拟怎么这么慢？”
• “我想用 Python 写 CUDA，但我不会 CUDA”
• “我需要 GPU 加速，但不想写 C++”
• “想做个实时粒子动效，但 Python 根本跑不动”

那么 Taichi 就非常适合你。

import taichi as ti
ti.init(arch=ti.gpu)
n = 100000
pos = ti.Vector.field(2, dtype=ti.f32, shape=n)
vel = ti.Vector.field(2, dtype=ti.f32, shape=n)
@ti.kernel
definit():
for i in range(n):
        pos[i] = ti.Vector([ti.random(), ti.random()])
        vel[i] = ti.Vector([0.0, -1.0])
@ti.kernel
defupdate(dt: ti.f32):
for i in range(n):
        pos[i] += vel[i] * dt
if pos[i].y < 0:
            pos[i].y = 1
init()
for _ in range(1000):
    update(0.01)

10 万粒子实时更新，用 Python 是不可能跑得动的，但 Taichi 可以轻松跑在 GPU 上。

如果你本来就在用 Python 做科学计算、图形学、AI 前处理、数值模拟，你会发现 Taichi 是一个非常“狠”的工具：

• 写 Python
• 跑 GPU
• 速度可以比纯 Python 快 100～200 倍

用起来不复杂，性能却很惊人。