企业官网建设流程全解析

Taichi Hello World 入门指南：用 Python 编写 GPU 并行计算与 Julia 分形

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi

导读

本文以 Taichi 官方入门文档《Hello, World!》为核心骨架，带领你从零写出并运行第一个 Taichi 程序——Julia 分形动画。你将掌握 Taichi 的三大核心概念：@ti.kernel与@ti.func装饰器如何触发 JIT 即时编译、ti.field数据容器的定义与遍历、以及外层for循环的自动并行化机制，同时理解后端（arch）的选择逻辑与 GUI 渲染流程。读完本文，你将具备独立编写高性能并行计算程序的基础能力。

Taichi：嵌入 Python 的领域专用语言

Taichi 是一种面向高性能并行计算的领域专用语言（DSL），内嵌于 Python 中。当编写计算密集型任务时，只需遵循一套额外的规则，并使用两个装饰器@ti.func与@ti.kernel，即可让 Taichi 接管计算任务，将装饰的函数通过**即时编译（JIT）**编译为机器码。此后对这些函数的调用将运行在多核 CPU 或 GPU 上，相比原生 Python 代码可实现约 50x～100x 的加速。

除了 JIT 之外，Taichi 还提供**预先编译（AOT）**系统，可将代码导出为二进制/着色器文件，脱离 Python 环境运行，参见 Tutorial: Run Taichi programs in C++ application。

前置条件与安装

在开始之前，请确认你的环境满足以下要求：

• Python：3.7 / 3.8 / 3.9 / 3.10（64 位）
• 操作系统：Windows、macOS、Linux（64 位）

Taichi 以 PyPI 包的形式分发，安装只需一条命令：

pip install taichi

Hello, world！Julia 分形程序

一个基础的 Julia 分形（Julia set）动画是理解 Taichi 编程语言核心概念的最佳起点。将下面的代码保存为fractal.py：

import taichi as ti import taichi.math as tm ti.init(arch=ti.gpu) n = 320 pixels = ti.field(dtype=float, shape=(n * 2, n)) @ti.func def complex_sqr(z): # complex square of a 2D vector return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1]) @ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixels c = tm.vec2(-0.8, tm.cos(t) * 0.2) z = tm.vec2(i / n - 1, j / n - 0.5) * 2 iterations = 0 while z.norm() < 20 and iterations < 50: z = complex_sqr(z) + c iterations += 1 pixels[i, j] = 1 - iterations * 0.02 gui = ti.GUI("Julia Set", res=(n * 2, n)) i = 0 while gui.running: paint(i * 0.03) gui.set_image(pixels) gui.show() i += 1

运行方式：如果你不使用 IDE，直接在终端进入脚本所在目录，执行python3 fractal.py（将filename替换为你的脚本名，务必保留.py扩展名）即可看到分形动画。

该示例对应的官方实现位于 python/taichi/examples/simulation/fractal.py，运行ti example -s fractal即可将其保存到当前工作目录，两者逻辑一致，可互为对照。

下面我们逐段剖析这个看似简单的程序。

导入 Taichi 与 math 模块

import taichi as ti import taichi.math as tm

前两行导入了 Taichi 及其math模块。math模块内置了小型向量与矩阵类型，例如 2D 实向量vec2、3×3 实矩阵mat3等。更多用法参见 Math Module。

初始化 Taichi：选择后端

ti.init(arch=ti.gpu)

ti.init()初始化 Taichi 运行环境，并根据传入的可选参数定制运行时。其中arch参数指定执行编译后代码的后端（backend），可为ti.cpu或ti.gpu：

• 指定ti.gpu时，Taichi 会按ti.cuda→ti.vulkan→ti.opengl/ti.Metal的顺序尝试使用 GPU 后端；
• 若 GPU 架构均不可用，则回退使用 CPU 作为后端；
• 也可以直接指定具体后端，例如arch=ti.cuda。若指定的架构不可用，Taichi 会直接抛出错误。

从源码看，ti.init()定义于 python/taichi/lang/misc.py，其完整签名为init(arch=None, default_fp=None, default_ip=None, _test_mode=False, enable_fallback=True, require_version=None, **kwargs)，核心职责是"设置整个程序使用的后端"。所有受支持的架构在 taichi/inc/archs.inc.h 中集中声明，包括x64、arm64、cuda、metal、opengl、dx11、dx12、opencl、amdgpu、vulkan、gles等，其中部分处于 WIP（开发中）或 N/A（不可用）状态。

除了arch，ti.init()还接受其他常用参数（来自 python/taichi/lang/misc.py 的 docstring 整理）：

此外，TI_DEFAULT_FP、TI_DEFAULT_IP等环境变量也可用于配置默认浮点/整型类型，且ti.init()的显式参数优先级更高（见 python/taichi/lang/misc.py）。更完整的全局设置说明见 Global Settings。

定义 Taichi field：数据容器

n = 320 pixels = ti.field(dtype=float, shape=(n * 2, n))

ti.field(dtype, shape)定义一个形状为shape、元素类型为dtype的 Taichi field。Field 是 Taichi 中最基础也最常用的数据结构，可以类比 NumPy 的ndarray或 PyTorch 的tensor，但更加灵活：field 可以是空间稀疏的，也可以方便地切换不同的数据布局。对本例而言，只需知道pixels是一个稠密二维数组即可。

从源码看，ti.field()定义于 python/taichi/lang/impl.py，支持如下参数（摘自其 docstring）：

实现上，ti.field(dtype=float, shape=(640, 320))等价于声明x = ti.field(ti.f32)后通过ti.root.dense(ti.ij, shape=(640, 320)).place(x)放置数据（见 python/taichi/lang/impl.py 中的等价示例）。向量/矩阵类型（如ti.math.vec3）的 dtype 会被转发到Vector.field/Matrix.field处理。

内核（kernel）与 Taichi 函数（func）

@ti.func def complex_sqr(z): # complex square of a 2D vector return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1]) @ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixels c = tm.vec2(-0.8, tm.cos(t) * 0.2) z = tm.vec2(i / n - 1, j / n - 0.5) * 2 iterations = 0 while z.norm() < 20 and iterations < 50: z = complex_sqr(z) + c iterations += 1 pixels[i, j] = 1 - iterations * 0.02

上述代码定义了两个函数，分别用@ti.func与@ti.kernel装饰，前者称为Taichi 函数（Taichi function），后者称为内核（kernel）。二者都不会被 Python 解释器执行，而是由 Taichi 的 JIT 编译器接管，并部署到由ti.init()的arch参数决定的并行多核 CPU 或 GPU 上。

Taichi 函数与内核的主要区别如下：

• 调用入口不同：kernel 是 Taichi 接管执行的入口，可在程序任意位置调用；而 Taichi 函数只能在内核或其他 Taichi 函数内部调用。例如本例中 Taichi 函数complex_sqr是在内核paint内被调用的。
• 类型标注要求不同：kernel 的参数与返回值必须显式标注类型，而 Taichi 函数不需要。例如内核paint的参数t: float标注了类型，而 Taichi 函数complex_sqr的参数z没有。
• 嵌套与递归限制：Taichi 支持嵌套函数，但不支持嵌套内核；同时 Taichi 函数内部不支持递归调用。

对于熟悉 CUDA 的读者：Taichi 的ti.func相当于 CUDA 的__device__，ti.kernel相当于 CUDA 的__global__。 对于熟悉 OpenGL 的读者：ti.func可比作 GLSL 中的普通函数，ti.kernel可比作 compute shader（计算着色器）。

在本例中，complex_sqr计算二维向量的复数平方（实部z0² - z1²、虚部2·z0·z1），而paint对每个像素迭代求解 Julia 集迭代公式z = z² + c，迭代次数决定像素灰度值。

并行 for 循环：Taichi 高性能的关键

@ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixels

Taichi 高性能的关键在于高效的迭代方式——利用并行化循环处理数据。内核最外层作用域中的for循环会被自动并行化，无需手动分配线程、回收资源或管理内存：

• 本例中for循环同时对i、j两个索引迭代，各迭代可并发执行；
• pixelsfield 被当作迭代器，i、j分别为取值0到2*n-1、0到n-1的整数索引；
• (i, j)组合按(0,0), (0,1), …, (0,n-1), (1,0), (1,1), …, (2*n-1, n-1)的顺序被并行遍历。

需要特别注意的是：嵌套在其他结构（如if/else或其他循环）内部的for循环不会被自动并行化，而是按顺序（串行）执行：

@ti.kernel def fill(): total = 0 for i in range(10): # Parallelized for j in range(5): # Serialized in each parallel thread total += i * j if total > 10: for k in range(5): # Not parallelized because it is not at the outermost scope

若想将最外层作用域的for循环改为串行，可使用ti.loop_config(serialize=True)，详见 Serialize a specified parallel for loop。

⚠️ 警告：break语句不支持出现在并行化的循环中：

@ti.kernel def foo(): for i in x: ... break # Error! @ti.kernel def foo(): for i in x: for j in range(10): ... break # OK!

内层串行循环中使用break是允许的。

注意，本例内核paint中使用的while循环位于并行for循环体内，属于每线程内部的串行控制流，因此不违反上述限制。

使用 GUI 显示结果

gui = ti.GUI("Julia Set", res=(n * 2, n)) # Sets the window title and the resolution i = 0 while gui.running: paint(i * 0.03) gui.set_image(pixels) gui.show() i += 1

为了将结果渲染到屏幕上，Taichi 提供了内置的 GUI System。gui.set_image()用于设置窗口内容，gui.show()用于显示更新后的图像。

Taichi 的 GUI 系统使用标准笛卡尔坐标系定义像素坐标，原点位于屏幕左下角：pixels中的(0, 0)元素映射到窗口左下角，(639, 319)元素映射到窗口右上角。从实现角度看，set_image接受numpy.ndarray、MatrixField、Field 或 Texture 等输入（见 python/taichi/ui/canvas.py）。官方 fractal 示例中采用了gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT)的循环方式来响应退出事件（见 python/taichi/examples/simulation/fractal.py），两种写法均可实现动画循环。

关键要点回顾

通过上述示例，你已经掌握了 Taichi 最重要的特性：

• Taichi 在指定后端上编译并执行 Taichi 函数与内核；
• 内核最外层作用域的for循环会被自动并行化；
• Taichi 提供了名为field的灵活数据容器，并可通过索引对其进行迭代。

更多 Taichi 示例

Julia 分形是 Taichi 收录的特色示例之一。要查看 Taichi Gallery 中更多精选示例，运行：

ti gallery

屏幕会弹出一个新的交互式窗口供你选择并运行示例。要查看完整示例列表，运行ti example。以下是一些实用的命令行：

• ti example -p fractal或ti example -P fractal：在终端打印 fractal 示例的源代码；
• ti example -s fractal：将示例保存到当前工作目录。

这些 CLI 命令由 python/taichi/_main.py 中的gallery与example方法实现，其中example会以彩色表格列出全部可用示例并按名称运行或导出。示例源文件位于 python/taichi/examples 目录，涵盖渲染（如fractal3d_ggui.py）、模拟（如mpm88.py、game_of_life.py）、算法（如mgpcg.py）等众多场景。

支持的系统与后端

下表概览了 Taichi 支持的操作系统及与之兼容的后端：

• ✅：支持；
• N/A：不可用。

注：macOS 上没有 CUDA 与 OpenGL 后端，对应的是 Apple Metal 后端；该表与 taichi/inc/archs.inc.h 中声明的架构清单相互印证。对于机器上存在多个可用后端的情况，ti.init(arch=ti.gpu)会按 CUDA → Vulkan → OpenGL/Metal 的优先级自动选择。

下一步

现在你已经准备好编写自己的 Taichi 程序了。以下文档介绍了 Taichi 在典型应用场景中的使用方式：

• Accelerate Python with Taichi
• Conduct Physical Simulation
• Accelerate PyTorch with Taichi

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi