Taichi太极图形编程语言实践demo

引言

小马无意中看到一种编程语言太极（Taichi），觉得挺有趣的。太极编程语言可以广泛应用在包括实时物理模拟、数字计算、增强现实、人工智能、视觉和机器人技术、电影和游戏中的视觉效果、通用计算等领域，属于计算机图形学。

还有搞游戏同学提议可以作为Unity中使用Nvidia的Flex插件的平替方案。用于搞游戏看起来贼拉爽，如下：

今天我们就来太极 Hello World一下。

一、太极（Taichi）编程语言是什么

至于语言的定义这里就不赘述了，网上非常多了。这里就是简单地介绍下我们关注的重点。

Taichi 是一种嵌入在 Python 中的领域特定语言。为了使 Taichi 能像 Python 包一样易于使用，基于这个目标我们做了大量的工程工作——使得每个 Python 程序员能够以最低的学习成本编写 Taichi 程序。你甚至可以选择你最喜欢的 Python 包管理系统、Python IDE 以及其他 Python 包和 Taichi 一起结合使用。

Taichi 既能在 CPU，也能在 GPU 上运行。 你只需根据你的硬件平台初始化 Taichi：

# 在 GPU 上运行，自动选择后端
ti.init(arch=ti.gpu)# 在 GPU 上运行， 使用 NVIDIA CUDA 后端
ti.init(arch=ti.cuda)
# 在 GPU 上运行， 使用 OpenGL 后端
ti.init(arch=ti.opengl)
# 在 GPU 上运行， 使用苹果 Metal 后端（仅对 OS X）有效
ti.init(arch=ti.metal)# 在 CPU 上运行 (默认)
ti.init(arch=ti.cpu)

不同操作系统所支持的后端：

平台	CPU	CUDA	OpenGL	Metal
Windows	可用	可用	可用	不可用
Linux	可用	可用	可用	不可用
Mac OS X	可用	不可用	不可用	可用

在参数 arch=ti.gpu 下，Taichi 将首先尝试在 CUDA 上运行。如果你的设备不支持 CUDA，那么 Taichi 将会转到 Metal 或 OpenGL。如果所在平台不支持 GPU 后端（CUDA、Metal 或 OpenGL），Taichi 将默认在 CPU 运行。

当在 Windows 平台 或者 ARM 设备（如 NVIDIA Jetson）上使用 CUDA 后端时， Taichi 会默认分配 1 GB 显存用于张量存储。如需重载显存分配，你可以在初始化的时候通过 ti.init(arch=ti.cuda, device_memory_GB=3.4) 来分配 3.4 GB 显存，或者使用 ti.init(arch=ti.cuda, device_memory_fraction=0.3) 来分配所有可用显存的 30%。
在其他平台上，Taichi 将会使用它的自适应内存分配器来动态分配内存。

我们有找到一个相对友好的教程《Taichi（太极） 0.6.24 编程语言文档》供大家参考，可惜文档中的版本偏低，部分示例代码可能不兼容新版本太极。

仓库：https://github.com/taichi-dev/taichi
仓库中也含了很多示例代码，大家可以直接debug，官方的比市面上的各种文档靠谱。可见目前Taichi的整个生态确实还是有待提升。

小马还找到个线上运行的地址：zoo.taichi。不过小马跑起来遇到远程接口报错不能用，大家不妨也可以试试。

二、实践运行太极demo

2.1 安装运行环境python

上文说到Taichi是一种嵌入在 Python 中的领域特定语言，所以运行太极需要先安装python环境。

如果已经有python环境的同学请绕过下面这个环节，不过需要注意检查版本，因为太极目前仅支持python 3.6/3.7/3.8版本。

从python官方下载python。

我调试用的是window，因此就找W的下载，但是这里再次注意不要直接安装最新版本的python 3.13.7，因为太极不支持。

所以我们应从链接进去选择3.8.10版本来安装。

如果安装成功可以看到如下画面。

当然如果你手快不小心安装了3.13版本，或者之前已经安装了太极支持以外的版本，那也没关系，我们可以更新版本或者再次安装其他版本，安装完我们设置好环境变量，就可以支持新装的版本了。

2.2 安装太极编程语言

准备好python环境后我们安装太极。

安装命令如下：

python -m pip install taichi

上面说到如果python版本不匹配，则会出现下面的提示。

如果版本匹配，如下：

不好意思，可能是遇到了些网络问题，不要慌，我们try again。再来！

看起来成功了。

2.3 运行太极demo

小马原本是计划参照文档示例代码跑一个经典案例。代码如下：

# fractal.py
import taichi as ti
ti.init(arch=ti.gpu)
n = 320
pixels = ti.var(dt=ti.f32, shape=(n * 2, n))
@ti.func
def complex_sqr(z):return ti.Vector([z[0] ** 2 - z[1] ** 2, z[1] * z[0] * 2])
@ti.kernel
def paint(t: ti.f32):for i, j in pixels: # 对于所有像素，并行执行c = ti.Vector([-0.8, ti.sin(t) * 0.2])z = ti.Vector([float(i) / n - 1, float(j) / n - 0.5]) * 2iterations = 0while z.norm() < 20 and iterations < 50:z = complex_sqr(z) + citerations += 1pixels[i, j] = 1 - iterations * 0.02
gui = ti.GUI("Fractal", (n * 2, n))
for i in range(1000000):paint(i * 0.03)gui.set_image(pixels)gui.show()

运行如下：

噢，出错了，跑不动。看起来这是语法版本问题呢。还有一个关于环境的告警，说是需要 添加环境变量TI_ENABLE_CUDA=0以跳过CUDA加载 才能解决。先不管，我们先更新代码解决代码问题。

于是调整代码如下：

import taichi as titi.init(arch=ti.gpu)n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n))@ti.func
def complex_sqr(z):return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])@ti.kernel
def paint(t: float):for i, j in pixels:  # Parallelized over all pixelsc = ti.Vector([-0.8, ti.cos(t) * 0.2])z = ti.Vector([i / n - 1, j / n - 0.5]) * 2iterations = 0while z.norm() < 20 and iterations < 50:z = complex_sqr(z) + citerations += 1pixels[i, j] = 1 - iterations * 0.02gui = ti.GUI("Hello World", res=(n * 2, n))for i in range(1000000):paint(i * 0.03)gui.set_image(pixels)gui.show()

看起来很OK了，我们再运行下。

成功啦！

更多效果案例参看官方这里 https://github.com/taichi-dev/taichi 吧，还是挺好玩的。小马也将继续探索 taichi 在游戏中的应用和大家分享，bye-bye！

---

• 彩蛋时间~ ~