企业官网建设流程全解析

Taichi 性能调优完全指南：循环指令、数据布局、TLS/BLS 局部存储优化与离线缓存

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

Taichi 是面向 Python 的"生产力、可移植性与高性能" GPU 编程框架，其编译器会自动并行化最外层 for 循环。本文围绕官方性能调优文档，系统讲解ti.loop_config循环指令的线程级调优、GPU 线程层级背景、稀疏场数据布局、TLS 线程局部存储与 BLS 块局部存储优化原理，以及ti.init()中离线缓存（Offline Cache）的配置策略，帮助开发者在 MPM、流体、渲染等典型负载上榨取最后几个百分点的性能。

For-loop 指令：细粒度控制循环并行度

Taichi 内核（@ti.kernel）会自动将最外层作用域的 for 循环并行化，编译器默认按目标架构自动设定调度参数。但对于追求极限性能的开发者，Taichi 提供了ti.loop_configAPI 来微调循环指令。例如，为循环指定合适的block_dim，在 mpm3d.py 这类 MPM 模拟中可能带来接近 3 倍的性能提升。

ti.loop_config用于设置下一个for 循环的指令，可用指令包括：

• parallelize：设置 CPU 上使用的线程数
• block_dim：设置 GPU 上一个 block 内的线程数
• serialize：若设为True，for 循环将串行执行，且可在循环体内写 break 语句（仅对 range/ndrange for 生效），等价于将parallelize设为 1

从 python/taichi/lang/misc.py 的loop_config实现可以看到，这些指令最终分别映射到底层 AST builder 的parallelize(v)、block_dim(dim)、strictly_serialize()调用；当parallelize == 1时还会自动调用strictly_serialize()强制串行化，这与文档中"serialize 等价于 parallelize=1"的描述完全一致。

@ti.kernel def break_in_serial_for() -> ti.i32: a = 0 ti.loop_config(serialize=True) for i in range(100): # This loop runs serially a += i if i == 10: break return a break_in_serial_for() # returns 55

n = 128 val = ti.field(ti.i32, shape=n) @ti.kernel def fill(): ti.loop_config(parallelize=8, block_dim=16) # If the kernel is run on the CPU backend, 8 threads will be used to run it # If the kernel is run on the CUDA backend, each block will have 16 threads. for i in range(n): val[i] = i

此外，loop_config还支持另外两个可选指令（见源码 docstring）：

• block_dim_adaptive：是否允许后端自适应设置 block_dim，默认开启；注意自适应 block_dim 仅在 CPU 后端生效，其他后端会给出警告
• bit_vectorize：是否对 quant_array 上的 struct for 启用位向量化，例如将 32 个 1-bit 元素一次性拷贝

背景：GPU 的线程层级

为了理解 for 循环如何被并行化，需要先了解当代 GPU 架构上的线程层级（Thread Hierarchy）。从细粒度到粗粒度，计算单元依次为：iteration（迭代）、thread（线程）、block（线程块）、grid（网格）。

• Iteration：for 循环的循环体即一次迭代，每次迭代对应 for 循环中不同的i值
• Thread：迭代被归类为线程。线程是最小的并行单元，一个线程内的所有迭代串行执行。为最大化并行效率，通常采用"一次迭代对应一个线程"的策略
• Block：线程被组织成块（block），块内所有线程并行执行，且可共享块内局部存储（block local storage，即 CUDA 的 shared memory）
• Grid：块被组成网格（grid），网格是从主机端启动的最小单元，网格内所有块并行执行。在 Taichi 中，每一个被并行化的 for 循环都表示为一个 grid

这里采用 CUDA 的术语，其他后端（如 OpenGL、Metal）遵循类似的线程层级结构。

示例：调整 for 循环的块级并行度

开发者可以在 for 循环前通过ti.loop_config调整其属性，指令只作用于紧接着的下一个循环：

@ti.kernel def func(): for i in range(8192): # no decorator, use default settings ... ti.loop_config(block_dim=128) # change the property of next for-loop: for i in range(8192): # will be parallelized with block_dim=128 ... for i in range(8192): # no decorator, use default settings ...

数据布局（Data Layouts）

由于 Taichi 将数据结构与计算分离，开发者可以自由实验不同的数据布局。与其他编程语言一样，选择高效的数据布局可能显著提升性能。关于 Taichi 中高级数据布局（如结构体数组 SoA、数组结构 AoS、稠密/指针等 SNode 层级、布局推导与手动布局）的完整讨论，请参阅 Fields (advanced) 章节。

局部存储优化（Local Storage Optimizations）

Taichi 内置若干利用快速内存（如 CUDA shared memory、L1 cache）的加速优化。简单来说，只要可行，Taichi 就会把对全局内存（慢）的访问替换为对局部内存（快）的访问，并在结束时把局部内存（如 CUDA shared memory）中的数据写回全局内存。这类变换保持原始程序的语义不变。

线程局部存储（TLS）

TLS 主要用于优化并行归约（parallel reduction）。当 Taichi 在@ti.kernel中识别出全局归约模式时，会在代码生成阶段自动应用 TLS 优化，其做法与常见 GPU 归约实现类似。

以下用 CUDA 术语演示一个例子：

x = ti.field(ti.f32, shape=1000000) s = ti.field(ti.f32, shape=()) @ti.kernel def sum(): for i in x: s[None] += x[i] sum()

Taichi 的并行循环在内部基于Grid-Stride Loops实现：每个物理 CUDA 线程可能处理x中的多个元素，也就是说，为sum启动的线程数可以少于x的 shape。

TLS 优化正是利用这一点：不再直接、原子地将x[i]加到全局内存目标s[None]上，而是进入线程时预分配一个线程局部缓冲区，将x的值非原子地累加进该缓冲区，在退出线程前再把缓冲区结果原子地加回s[None]。如果每个线程处理N个元素，原子加的次数就降为原来的1/N。

此外，最后对全局内存s[None]的原子加还利用了 CUDA 的 warp 级内建指令进一步优化，进一步减少所需原子加次数。

当前 TLS 的适用范围：支持add、sub、min、max这四类归约算子，作用于0D标量/向量/矩阵ti.field；尚不支持ti.ndarray。

在 Nvidia GeForce RTX 3090 上对一维 8M float 的 Taichi field 做全局 max 归约的基准对比：

• TLS 关闭：5.2 × 1e3 us
• TLS 开启：5.7 × 1e1 us

TLS 带来了约100 倍加速，且 TLS 归约求和性能与 CUDA 实现相当。

块局部存储（BLS）

适用场景：对于最后一层为denseSNode 的稀疏场（即层级结构匹配ti.root.(sparse SNode)+.dense），Taichi 会将每个dense容器（或denseblock）分配给一个 CUDA 线程块。BLS 优化专门针对这类场生效。

BLS 旨在利用 CUDA shared memory 加速模板计算（stencil computation）。优化流程如下：

1. 用户通过ti.block_local标注想要缓存的场集合
2. 在编译期，Taichi 尝试识别这些标注场相对于denseblock 的访问范围
3. 若识别成功，Taichi 生成代码：先把范围内可访问的数据全部加载到block local缓冲区（CUDA shared memory），再将循环体中对相应槽位的访问全部替换为该缓冲区

下面演示 BLS 的用法，a是一个 block size 为4x4的稀疏场：

a = ti.field(ti.f32) b = ti.field(ti.f32) # `a` has a block size of 4x4 ti.root.pointer(ti.ij, 32).dense(ti.ij, 4).place(a) @ti.kernel def foo(): # Taichi will cache `a` into the CUDA shared memory ti.block_local(a) for i, j in a: print(a[i - 1, j], a[i, j + 2])

每个循环迭代访问的元素相对于自身坐标的偏移分别是[-1, 0]和[0, 2]。因此，对于一个从[M, N]（含）到[M + 4, N + 4]（不含）的完整 block，其相对于该 block 的访问范围是[M - 1, M + 4) x [N, N + 6)（由[M + (-1), M + 4) x [N, N + 4 + 2)推导得出）。全局坐标i, j与缓冲区局部索引的映射关系如下图：

作为用户，你无需关心这些底层细节——Taichi 自动完成所有推断以及全局/块局部映射：预分配一块5x6的 CUDA shared memory 缓冲区，将a的内容预加载进缓冲区，然后把循环体内的所有a（全局内存）访问替换为缓冲区访问。上面的例子没有修改a本身；而如果被块缓存的场发生了写入，Taichi 会生成把缓冲区写回全局内存的代码。

注意：BLS 并非没有代价。BLS 面向的是具有大量重叠全局内存访问的模板计算。如果不符合这一条件，预加载与回写反而可能降低性能。此外，新一代 Nvidia GPU 已缩小了全局内存与 shared memory 在只读访问上的差距，目前我们观察到 BLS 对于缓存原子操作的目标地址更有效。作为经验法则，建议通过基准测试来判断是否启用 BLS。

从 python/taichi/lang/misc.py 的block_local实现可见，它会把场成员标记为SNodeAccessFlag.block_local；且当优化级别为 0 时会强制提升到opt_level = 1以启用 BLS 分析，这也印证了 BLS 需要依赖编译器优化管线。

离线缓存（Offline Cache）

Taichi 内核在第一次被调用时会被隐式编译。为降低后续调用的开销，编译结果会保存在在线的内存缓存中：只要内核未改变，即可立即加载并启动。但应用退出后该缓存即失效，重启程序后 Taichi 必须重新编译所有内核例程并重建在线内存缓存——由于编译开销，Taichi 函数的首次启动通常较慢。

离线缓存功能解决了这一问题：它将编译缓存转储到磁盘供后续运行使用，从而大幅降低重复运行时的首次启动开销。Taichi 默认构建并维护离线缓存，同时提供若干ti.init()选项用于配置其行为：

• offline_cache: bool：启用或禁用离线缓存。默认值：True
• offline_cache_file_path: str：存放离线缓存文件的目录。Windows 上默认'C:\taichi_cache\ticache\'，Unix 系系统默认'~/.cache/taichi/ticache/'。目录会自动创建
• offline_cache_max_size_of_files: int32：缓存文件的最大字节数。默认值：100MB。当缓存文件大小超过该限制时触发清理进程
• offline_cache_cleaning_policy: str：缓存中过时文件的替换策略。可选值：'never'、'version'、'lru'、'fifo'。默认值：'lru'
  • 'never'：从不清理，无论offline_cache_max_size_of_files配置如何都保留所有缓存文件
  • 'version'：仅丢弃与内核函数相关的旧版本缓存文件
  • 'lru'：丢弃最近最少使用的缓存文件
  • 'fifo'：丢弃最早加入的缓存文件

例如，在ti.init()中开启离线缓存并指定缓存目录：

ti.init(arch=ti.gpu, offline_cache=True, offline_cache_file_path='./ticache', offline_cache_max_size_of_files=2 * 1024 * 1024 * 1024, offline_cache_cleaning_policy='lru')

在 python/taichi/lang/misc.py 中可以看到前端默认将cfg.offline_cache = True，与文档所述"默认启用"一致。

验证效果：将示例程序运行两次并观察启动开销，结果如下图所示（蓝色为禁用离线缓存，橙色为首次启用离线缓存，绿色为第二次启用）：

可以看到，MPM99、MPM3D、Cornell Box 等示例在第二次运行（绿色柱）时启动时间均趋近于 0 秒，而禁用缓存时则需要数秒至数十秒的编译时间。

注意：如果你的代码行为异常，可通过设置环境变量TI_OFFLINE_CACHE=0或在ti.init()中设置offline_cache=False来禁用离线缓存，并向 Taichi 仓库提交 issue 反馈。

小结与实践建议

1. 循环级调优：使用ti.loop_config(parallelize=..., block_dim=...)为 CPU 指定线程数、为 GPU 指定每块线程数；需要 break 时用serialize=True。注意指令只作用于紧接着的下一个 for 循环。
2. 数据布局：善用 Taichi 数据与计算分离的特性，参考 Fields (advanced) 选择高效布局。
3. 归约负载：优先使用 0Dti.field上的add/sub/min/max归约，让 TLS 自动生效；ti.ndarray目前不支持 TLS。
4. 模板计算：对符合ti.root.(sparse)+.dense层级且访存重叠度高的 stencil 核函数，用ti.block_local配合基准测试决定是否启用 BLS。
5. 重复启动场景：保持默认的离线缓存开启，必要时通过ti.init()或TI_OFFLINE_CACHE环境变量调整缓存目录、大小上限与清理策略。

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