Taichi LLVM 稀疏运行时深度解析：SNode 存储、回收式内存分配器与 GPU 垃圾回收-平芜编程栈

Taichi LLVM 稀疏运行时深度解析：SNode 存储、回收式内存分配器与 GPU 垃圾回收

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

Taichi 的 LLVM 稀疏运行时（sparse runtime）是支撑pointer、dynamic等稀疏 SNode（Structural Node）的核心基础设施。本文以仓库中的设计文档 docs/design/llvm_sparse_runtime.md 为骨架，结合 runtime.cpp 及各个 SNode 头文件的最新源码，逐层拆解 SNode 的统一接口、pointer/dynamic的按需分配实现、NodeManager回收式分配器、ListManager块式存储，以及 GPU 上的三阶段并行垃圾回收（GC）。读完本文，你将能理解 Taichi 稀疏计算（如大规模网格、粒子仿真）背后"按需激活、回收复用、环境值兜底"的完整内存管理链路。

1. 概览：稀疏运行时在 Taichi 中的位置

在 Taichi 中，SNode 是描述数据布局的树形结构节点，共有四种类型：dense、bitmasked、pointer和dynamic。其中dense与bitmasked的单元（cell）在内存中连续排列、始终存在；而dynamic与pointer在空间上天然稀疏——单元不保证连续存储，只有被激活的单元才会占用内存。因此本文后续所说的"稀疏 SNode"专指dynamic和pointer两类。

LLVM 稀疏运行时的全部源码位于 taichi/runtime/llvm/runtime_module 目录，组织方式如下：

每个 SNode 类型对应一个头文件：node_dense.h、node_bitmasked.h、node_pointer.h、node_dynamic.h、node_root.h；
所有 SNode 类型共享一个公共源文件 runtime.cpp，其中定义了运行时核心结构、内存分配器与 GC 逻辑；
并发原语集中在 locked_task.h。

值得注意的一个关键设计是：runtime.cpp并不会被链接进 Taichi 的核心 C++ 库，而是通过 Clang 编译成 LLVM 字节码文件（.bc）。这一点可以在 CMakeLists.txt 中得到印证——其中COMPILE_LLVM_RUNTIME函数对每个后端架构执行clang -c runtime.cpp -emit-llvm ... -D "ARCH_${rtm_arch}"，生成runtime_${rtm_arch}.bc并安装到python/taichi/_lib/runtime目录：

COMMAND ${CLANG_EXECUTABLE} ${CLANG_OSX_FLAGS} -c runtime.cpp -o "runtime_${rtm_arch}.bc" -fno-exceptions -emit-llvm -std=c++17 -D "ARCH_${rtm_arch}" -I ${PROJECT_SOURCE_DIR}

运行时启动时，.bc文件被加载进内存、反序列化为llvm::Module，再与 JIT 编译出的 Taichi kernel 链接。这种"运行时写成 C++、统一编译为字节码"的设计带来两个好处：

一次编写、多后端复用：CPU、CUDA 等所有 LLVM 后端共享同一份运行时代码；
实现语言自由：稀疏运行时可以用具备足够抽象能力的语言（如 C++）编写，而不必直接手写裸 LLVM IR。

2. SNode 统一接口：`StructMeta`与`X_*`系列函数

每个 SNode 类型都带有一个继承自StructMeta的XMeta结构体。当前仓库中StructMeta定义在 runtime.cpp：

struct StructMeta { i32 snode_id; std::size_t element_size; i64 max_num_elements; Ptr (*lookup_element)(Ptr, Ptr, int i); Ptr (*from_parent_element)(Ptr); u1 (*is_active)(Ptr, Ptr, int i); i32 (*get_num_elements)(Ptr, Ptr); void (*refine_coordinates)(PhysicalCoordinates *inp_coord, PhysicalCoordinates *refined_coord, int index); RuntimeContext *context; };

其中Ptr是uint8_t*的别名，i32对应int32_t，u1对应布尔值。meta指向对应的 meta 结构体，node指向 SNode 实例本身。StructMeta定义了所有 SNode 共享的核心属性与接口：

成员	含义
`snode_id`	SNode ID，用于在运行时数组中索引该 SNode 专属资源（如内存分配器）
`element_size`	单个 cell 的字节大小
`max_num_elements`	该 SNode 可容纳的单元容量上限（注意是容量而非当前活跃数）
`X_get_num_elements(meta, node)`	返回该 SNode 可容纳的单元数
`X_activate(meta, node, i)`	激活第`i`个单元
`X_is_active(meta, node, i)`	检查第`i`个单元是否活跃
`X_lookup_element(meta, node, i)`	返回第`i`个单元的指针；对稀疏 SNode 可能返回`nullptr`
`refine_coordinates`	坐标细化（用于多维索引到物理坐标的映射）
`context`	指向运行时上下文`RuntimeContext`

对于稀疏 SNode，还额外实现了一个接口：

X_deactivate(meta, node, i)：将第i个单元去激活，把其内存交还给回收池。

原设计文档也指出，这个额外 API 未来可能被调整，目标是让所有 SNode 共享同一套 API。StructMeta中的函数指针在runtime.cpp中通过STRUCT_FIELD宏暴露给 LLVM 后端访问，LLVMRuntime与NodeManager等结构体同样通过该宏体系被 JIT 后的 kernel 识别，这是"同一份 C++ 运行时既能在 C++ 侧构建、又能在 LLVM IR 中被寻址"的关键机制。

3.`dense`SNode：最简单的连续数组

dense是最简单的 SNode 形态：它就是一段连续内存中的 cell 数组，用 C++ 类比即std::array<Cell, N>。其实现位于 node_dense.h，代码非常直白：

struct DenseMeta : public StructMeta { int morton_dim; }; i32 Dense_get_num_elements(Ptr meta, Ptr node) { return ((StructMeta *)meta)->max_num_elements; } void Dense_activate(Ptr meta, Ptr node, int i) { // Dense elements are always active } u1 Dense_is_active(Ptr meta, Ptr node, int i) { return true; } Ptr Dense_lookup_element(Ptr meta, Ptr node, int i) { return node + ((StructMeta *)meta)->element_size * i; }

四个接口的语义一目了然：

Dense_get_num_elements：直接返回DenseMeta中存的max_num_elements；
Dense_activate：空操作——dense的 cell 永远处于激活状态，无需分配动作；
Dense_is_active：恒返回1；
Dense_lookup_element：返回第i个 cell 的地址，即node + element_size * i，是纯算术寻址，没有任何间接跳转。

内存布局示意如下：

+- node | +------------+------------+------------+------------+ | | | | | | cell-0 | cell-1 | cell-2 | cell-3 | | | | | | +------------+------------+------------+------------+

由于dense不存在"未激活单元"，它的访问开销最低，是 Taichi 中默认、最高效的存储形式；bitmasked则在dense的基础上附加一个位掩码记录激活状态，但存储仍是连续的，本文不再展开。

4.`pointer`SNode：按需分配 + 环境值兜底的稀疏存储

pointer是稀疏计算中最常用的选择。它只在实际激活 cell 时从内存池动态分配空间，去激活时把内存回收进池中复用，从而在大规模网格计算中显著节省内存。用 C++ 类比，它可以看作std::array<Cell*, N>：一个固定大小的指针数组，每个指针要么指向真实的 cell，要么为nullptr。

一个重要的初始化细节：Taichi 在启动时会预分配一块名为ambient_elements的内存，由所有未激活的稀疏 SNode 共享。因此，对未激活稀疏单元的"解引用"会得到ambient_elements中存储的默认值（通常是零）。这一"环境值"设计使得稀疏访问无需显式判空分支即可安全返回默认值。

pointerSNode 的内存布局如下——前半段存放每个 cell 的锁（64 位整数），后半段存放指向 cell 的指针：

+- node | +------------+------------+------------+------------+------------+------------+------------+------------+ | c0-lock | c1-lock | c2-lock | c3-lock | nullptr | *cell-1 | *cell-2 | nullptr | +------------+------------+------------+------------+------------+------------+------------+------------+ | | | +> +------------+ | | | | | cell-2 | | | | | +------------+ | +-------------> +------------+ | | | cell-1 | | | +------------+

4.1`Pointer_activate`：双重检查锁 + warp 内代表线程

以 node_pointer.h 中的Pointer_activate为例，可以看到稀疏运行时基础设施是如何支撑pointerSNode 的：

void Pointer_activate(Ptr meta_, Ptr node, int i) { auto meta = (StructMeta *)meta_; auto num_elements = Pointer_get_num_elements(meta_, node); // 1 volatile Ptr lock = node + 8 * i; volatile Ptr *data_ptr = (Ptr *)(node + 8 * (num_elements + i)); // 2 if (*data_ptr == nullptr) { // 3 // The cuda_ calls will return 0 or do noop on CPUs u32 mask = cuda_active_mask(); if (is_representative(mask, (u64)lock)) { // 4 locked_task( lock, [&] { // 5 auto rt = meta->context->runtime; auto alloc = rt->node_allocators[meta->snode_id]; auto allocated = (u64)alloc->allocate(); atomic_exchange_u64((u64 *)data_ptr, allocated); }, [&]() { return *data_ptr == nullptr; }); } warp_barrier(mask); } }

逐行拆解：

定位锁与数据指针：根据布局，node + 8 * i是第i个 cell 的锁（64 位整数），node + 8 * (num_elements + i)是第i个 cell 的数据指针。这里假设指针宽度为 8 字节。
无锁预检查：直接读data_ptr是否为nullptr。这是经典的双重检查锁定（double-checked locking）模式——先用一次廉价读避免大多数情况下的锁开销。
warp 内选代表线程：如果预检查为真，在 CUDA warp 内选举一个线程去抢锁（is_representative）。这是防止锁竞争的小优化；更重要的是，在没有独立线程调度（independent thread scheduling）的前 Volta 架构上，同一 warp 内多个线程同时争抢同一把锁可能导致死锁，串行化代表线程能规避这一风险。is_representative在 CPU 后端直接返回true，所有cuda_系列调用在 CPU 上返回 0 或空操作。
抢锁执行：胜出的线程通过locked_task（定义在 locked_task.h）获取锁。该原语在 CPU 端等价于互斥锁 + 条件测试；在 CUDA 端则区分计算能力：Volta 之前用"warp 串行化"避免死锁，Volta 及之后因具备独立线程调度而可以直接使用锁，同时通过grid_memfence保证 CUDA 弱内存模型下的可见性。
分配并发布：取回该 SNode 专属的分配器node_allocators[meta->snode_id]，分配一个新的 cell，并把地址原子地写入data_ptr。注意每个 SNode 的 cell 大小不同，因此运行时为每个 SNode 维护独立的内存分配器，由它决定每个 cell 该分配多少空间。

Pointer_deactivate与Pointer_is_active的逻辑与之对称：去激活时在锁内将指针所指向的 cell 交还给分配器（alloc->recycle(data_ptr)）并把指针置空；判断激活只需检查data_ptr != nullptr。Pointer_lookup_element则体现了环境值兜底——当指针为nullptr时返回ambient_elements[snode_id]的地址，从而让"读未激活单元"得到默认零值。

5.`dynamic`SNode：变长向量与分块链表

dynamic在几个方面很特殊：

它必须是1 维的终止 SNode：所谓"终止"指dynamic之后只能place一个 Taichi 字段，不能再挂其他 SNode；
它的轴（axis）必须与其所有前驱节点的轴互不相同。例如：
- dense(ti.ij, (2, 4)).dynamic(ti.k, 8)：合法，dynamic的轴k是唯一的；
- dense(ti.ij, (2, 4)).dynamic(ti.j, 8)：报错，因为dynamic的轴j与dense在轴j上重叠。

逻辑上，dynamic可以看作std::vector<int32_t>，但其物理实现是由 chunk 构成的单链表。dynamic节点的头部结构在 node_dynamic.h 中定义：

struct DynamicNode { i32 lock; i32 n; // 当前元素个数 Ptr ptr; // 指向第一个 chunk };

内存布局示意（n记录当前元素个数，ptr指向首个 chunk，每个 chunk 头部存指向下一个 chunk 的指针，随后是元素数据）：

+- node | +------------+------------+------------+ | lock | n | ptr | +------------+------------+------------+ | +-+------------+ # chunk-0, chunk_start = 0 | >------|--+ +------------+ | | 0 | | +------------+ | | 1 | | +------------+ | | 2 | | +------------+ | | 3 | | +------------+ | | +-+------------+ # chunk-1, chunk_start = 4 | nullptr | +------------+ | 4 | +------------+ | 5 | +------------+ | 6 | +------------+ | 7 | +------------+

5.1`Dynamic_allocate`：原子追加与按需建块

设计文档当时以Dynamic_append为例，而当前仓库已将"追加"重构为Dynamic_allocate（见 node_dynamic.h），它返回指向新元素槽位的指针，并回写元素索引，核心逻辑与Dynamic_append一脉相承：

Ptr Dynamic_allocate(Ptr meta_, Ptr node_, i32 *len) { auto meta = (DynamicMeta *)(meta_); auto node = (DynamicNode *)(node_); auto chunk_size = meta->chunk_size; // 1 auto i = atomic_add_i32(&node->n, 1); *len = i; int chunk_start = 0; auto p_chunk_ptr = &node->ptr; while (true) { // 2 if (*p_chunk_ptr == nullptr) { locked_task(Ptr(&node->lock), [&] { if (*p_chunk_ptr == nullptr) { auto rt = meta->context->runtime; auto alloc = rt->node_allocators[meta->snode_id]; *p_chunk_ptr = alloc->allocate(); } }); } // 3 if (i < chunk_start + chunk_size) { // 3-1 return *p_chunk_ptr + sizeof(Ptr) + (i - chunk_start) * meta->element_size; } // 3-2 p_chunk_ptr = (Ptr *)(*p_chunk_ptr); chunk_start += chunk_size; } }

取索引：用当前长度n作为新元素的下标，atomic_add_i32保证并发追加时每个线程拿到互不相同的索引。
按需建块：chunk_start始终从 0 开始记录当前 chunk 的起始下标，p_chunk_ptr初始指向第一个 chunk 的指针。在while循环里，若当前 chunk 槽位为空，则在锁内从该 SNode 的分配器分配一个新 chunk（双重检查：锁内再次判空，防止锁竞争期间被其他线程抢先分配）。
定位与返回：比较索引i是否落在当前 chunk 范围内。
- 若是，返回该槽位地址。注意偏移量要跳过开头的sizeof(Ptr)字节——它们被保留用于存储下一个 chunk 的地址；
- 否则沿链表跳到下一个 chunk，chunk_start累加chunk_size。

与激活/去激活相关的配套函数还包括：Dynamic_activate（将node->n原子地更新为max(n, i+1)并确保第i个元素所在 chunk 已分配）、Dynamic_deactivate（锁内将n清零并把所有 chunk 逐个recycle回分配器）、Dynamic_is_active（判断i < node->n）以及Dynamic_lookup_element（活跃时沿链表寻址，否则返回环境值）。

6. 运行时核心：`LLVMRuntime`

稀疏运行时的核心数据结构是 runtime.cpp 中的LLVMRuntime，它持有若干关键数据：

所有根 SNode 的信息：roots与root_mem_sizes；
SNode 内存分配器数组node_allocators（类型为NodeManager*，按snode_id索引）；
环境值数组ambient_elements，为每个稀疏 SNode 提供默认值兜底；
ti.random使用的随机状态rand_states；
打印与错误消息缓冲区（error_message_template、error_message_arguments、error_code等）；
内存池与跨后端分配入口allocate_aligned、allocate_from_reserved_memory；
结果缓冲区result_buffer、profile 钩子、线程池等。

其中node_allocators[snode_id]与ambient_elements[snode_id]的初始化在LLVMRuntime的初始化代码中完成（见 runtime.cpp 附近），为每个 SNode 建立专属分配器并预分配共享环境值区域。NodeManager正是稀疏 SNode 的基石，下一节深入分析它的实现。

7.`NodeManager`：回收式内存分配器

每个 SNode 都关联一个专属分配器，统一存放在node_allocators数组中。分配器类型为NodeManager（见 runtime.cpp），内部维护三个ListManager链表：

链表	作用
`data_list`	固定大小内存块列表，每块可存放`chunk_num_elements`个 SNode cell；前文提到的各种 chunk（cell、dynamic chunk）都来自这里
`free_list`	空闲 cell 的索引列表，每个节点是一个`int32_t`（即`list_data_type`）。分配时优先复用其中的索引，不够才向内存分配器申请新空间
`recycled_list`	已释放 cell 的索引列表。一次 GC 执行后，其中的条目会被转入`free_list`复用；每个节点同样是一个`int32_t`

NodeManager构造时有两个值得注意的参数化逻辑：

默认每个 chunk 容纳128 * 1024（128K）个元素；
同时保证单块不超过 128 MB：若chunk_num_elements * element_size > 128 MB，则不断将chunk_num_elements减半。

7.1`allocate()`：先复用，后扩容

NodeManager::allocate()的实现如下：

Ptr allocate() { // 1 int old_cursor = atomic_add_i32(&free_list_used, 1); i32 l; if (old_cursor >= free_list->size()) { // 2 l = data_list->reserve_new_element(); } else { // 3 l = free_list->get<list_data_type>(old_cursor); } // 4 return data_list->get_element_ptr(l); }

读取候选空闲索引：free_list_used是free_list中已消费位置的游标，原子自增后得到本次要考察的索引位置old_cursor。
游标越界则扩容：若old_cursor已超出free_list当前长度，说明空闲索引用尽，调用data_list->reserve_new_element()从data_list分配一块新 chunk。
否则复用：直接取free_list中第old_cursor个索引。
统一返回：无论走哪条分支，索引l都指向data_list中的某个内存槽位，返回该槽位指针。

回收过程recycle(ptr)也很直接：先用data_list->ptr2index(ptr)把指针换算成索引，再把索引append进recycled_list，等待下一次 GC 统一回填到free_list。

8.`ListManager`：CPU/GPU 通用的"无限长"块式列表

ListManager（见 runtime.cpp）的官方注释是一句精炼的描述：

A simple list data structure that is infinitely long. Data are organized in chunks, where each chunk is allocated on demand.

虽然名字带"list"，但把它叫链表其实有些误导——它的行为更像 C++ 的std::deque：ListManager在chunks数组中持有一系列内存块，每块可容纳max_num_elements_per_chunk个元素，块按需分配。

构造函数展示了必要的成员变量，并强制要求每块元素数必须是 2 的幂（is_power_of_two断言），以便用移位代替除法：

ListManager(LLVMRuntime *runtime, std::size_t element_size, std::size_t num_elements_per_chunk) : element_size(element_size), max_num_elements_per_chunk(num_elements_per_chunk), runtime(runtime) { taichi_assert_runtime(runtime, is_power_of_two(max_num_elements_per_chunk), "max_num_elements_per_chunk must be POT."); lock = 0; num_elements = 0; log2chunk_num_elements = taichi::log2int(num_elements_per_chunk); }

8.1`reserve_new_element()`与`touch_chunk()`

分配新元素的关键方法是reserve_new_element()（runtime.cpp）：

i32 reserve_new_element() { auto i = atomic_add_i32(&num_elements, 1); auto chunk_id = i >> log2chunk_num_elements; touch_chunk(chunk_id); return i; }

它原子地自增num_elements得到新元素索引，右移得到所属 chunk ID，然后调用touch_chunk()确保该 chunk 已实际分配。touch_chunk()（runtime.cpp）再次使用双重检查锁：

void ListManager::touch_chunk(int chunk_id) { taichi_assert_runtime(runtime, chunk_id < max_num_chunks, "List manager out of chunks."); if (!chunks[chunk_id]) { locked_task(&lock, [&] { // may have been allocated during lock contention if (!chunks[chunk_id]) { grid_memfence(); auto chunk_ptr = runtime->allocate_aligned( runtime->runtime_memory_chunk, max_num_elements_per_chunk * element_size, 4096, true /*request*/); atomic_exchange_u64((u64 *)&chunks[chunk_id], (u64)chunk_ptr); } }); } }

注意两点：一是"锁内再次判空"应对锁竞争期间其他线程抢先分配的情况；二是底层真正的内存分配由LLVMRuntime::allocate_aligned()完成——它从预分配内存（CUDA/AMDGPU）或请求式分配路径按 4096 字节对齐取块。

8.2 索引寻址、指针反查与增删改

get_element_ptr(i)给出了从索引到地址的 O(1) 映射（runtime.cpp）：

Ptr get_element_ptr(i32 i) { return chunks[i >> log2chunk_num_elements] + // chunk base element_size * (i & ((1 << log2chunk_num_elements) - 1)); // slot within the chunk }

与之相反，ptr2index(ptr)（runtime.cpp）逐个遍历 chunk，判断传入指针落在哪一块的内存范围内，从而反推出索引——这正是NodeManager::recycle把指针转回索引的手段。

有了这些原语，其余方法都变得非常简单：

allocate()（runtime.cpp）：reserve_new_element()后返回get_element_ptr(i)；
append(data_ptr)（runtime.cpp）：allocate()后std::memcpy拷贝element_size字节；
clear()（runtime.cpp）：仅把num_elements重置为 0，不触碰列表内容——这个"惰性清空"特性在后续 GC 阶段被巧妙利用。

9. 垃圾回收（GC）：GPU 上的三阶段并行回收

当某个 offloaded task 可能涉及稀疏 SNode 去激活之后，就会触发 GC（在 GPU 上以并行方式执行）。整个 GC 过程针对每个 SNode 分为三个阶段。

9.1 阶段 0（`gc_parallel_0`）：压实`free_list`

第一阶段（runtime.cpp）把free_list中尚未被使用的索引移动到列表头部。正常情况下，这只需一个for循环拷贝数据即可，但在 GPU 上并行执行且源、目标区间可能重叠时，就需要特别小心——代码区分了"目的区间与源区间不重叠"和"重叠"两种情况：

# 图例："\\\" 表示已被复用的 cell 索引，" "（空白）表示仍然可用的 cell 索引。 # src 与 dst 区间不重叠：src=[6, 8), dst=[0, 2) 0 1 2 3 4 5 6 7 8 0 1 2 +---+---+---+---+---+---+---+---+ +---+---+ |\\\|\\\|\\\|\\\|\\\|\\\| | | ---> | | | +---+---+---+---+---+---+---+---+ +---+---+ src dst # src 与 dst 区间重叠：src=[3, 8), dst=[0, 5) 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 +---+---+---+---+---+---+---+---+ +---+---+---+---+---+ |\\\|\\\|\\\| | | | | | ---> | | | | | | +---+---+---+---+---+---+---+---+ +---+---+---+---+---+ src dst

注意在重叠场景中，cell 索引 3 和 4 同时属于源区间与目的区间。对应的实现gc_parallel_impl_0依据free_list_used * 2 > free_list_size判断走哪条拷贝路径：不重叠时直接平移尾部元素；重叠时只移动不重叠的部分，避免并行写覆盖。

9.2 阶段 1（`gc_parallel_1`）：单线程簿记

第二阶段（runtime.cpp）为free_list和recycled_list做簿记：

void gc_parallel_impl_1(NodeManager *allocator) { auto free_list = allocator->free_list; const i32 num_unused = max_i32(free_list->size() - allocator->free_list_used, 0); free_list->resize(num_unused); allocator->free_list_used = 0; allocator->recycle_list_size_backup = allocator->recycled_list->size(); allocator->recycled_list->clear(); }

这一步必须在单线程上执行，以避免在recycled_list上产生数据竞争：GC 期间recycled_list必须被清空，同时其全部内容要转入free_list。若把这件事放到并行执行的第三阶段，GPU 线程之间将很难协调"清空"与"转移"的事件顺序，因此专门设计了这个串行阶段。它做两件事：

把recycled_list的元素个数备份到recycle_list_size_backup；
清空recycled_list。

由于这两步计算量很小，即使串行执行也不会造成明显的处理时间负担。

9.3 阶段 2（`gc_parallel_2`）：回填`free_list`并并行清零

第三阶段（runtime.cpp）把recycled_list中的索引回填进free_list，同时把所有回收的 cell 内存清零：

void gc_parallel_impl_2(NodeManager *allocator) { auto elements = allocator->recycle_list_size_backup; auto free_list = allocator->free_list; auto recycled_list = allocator->recycled_list; auto data_list = allocator->data_list; auto element_size = allocator->element_size; using T = NodeManager::list_data_type; auto i = block_idx(); while (i < elements) { auto idx = recycled_list->get<T>(i); auto ptr = data_list->get_element_ptr(idx); if (thread_idx() == 0) { free_list->push_back(idx); } // memset auto ptr_stop = ptr + element_size; if ((uint64)ptr % 4 != 0) { auto new_ptr = ptr + 4 - (uint64)ptr % 4; if (thread_idx() == 0) { for (uint8 *p = ptr; p < new_ptr; p++) { *p = 0; } } ptr = new_ptr; } // now ptr is a multiple of 4 ptr += thread_idx() * sizeof(uint32); while (ptr + sizeof(uint32) <= ptr_stop) { *(uint32 *)ptr = 0; ptr += sizeof(uint32) * block_dim(); } while (ptr < ptr_stop) { *ptr = 0; ptr++; } i += grid_dim(); } }

逐点理解这段代码：

读回备份：取阶段 1 备份的recycled_list大小。因为ListManager::clear()并不真正修改内容，所以此时仍可安全读取该列表。
块级并行：i初始化为 CUDA block 索引，每个 CUDA block 协同完成一个 SNode cell 内存区域的清零。
读索引：取出recycled_list中第i个元素，即一个待回收 cell 的索引，再换算成data_list中的指针。
线程 0 回填：block 内第一个线程负责把该索引push_back进free_list。
对齐处理：清零前先处理 4 字节不对齐的头部——若ptr不是 4 的倍数，由线程 0 逐字节把开头几个字节清零，直到对齐。
并行清零：对齐后，每个线程从自己的偏移thread_idx() * sizeof(uint32)开始，以步长sizeof(uint32) * block_dim()跳跃式写 0，覆盖整个 cell；尾部不足 4 字节的残余字节再由单线程逐字节清零。以 4 个线程清零 40 字节 cell 为例，写入分布如下：
```
# A SNode cell of 40 bytes, zero-filled by 4 threads: +----+----+----+----+----+----+----+----+----+----+ | t0 | t1 | t2 | t3 | t0 | t1 | t2 | t3 | t0 | t1 | +----+----+----+----+----+----+----+----+----+----+
```
跨块跳转：整个 block 处理完一个 cell 后，i += grid_dim()跳到下一个待回收 cell。

清零的意义在于：复用的 cell 必须以干净的默认状态重新投入使用，避免旧数据泄漏到新激活的 cell 中。CPU 后端则走gc_serial的串行版本（见 runtime.cpp），逻辑等价：压实free_list、逐个memset清零回收 cell 并回填free_list、清空recycled_list。

10. 总结：稀疏运行时的完整生命周期

把上述模块串起来，一次典型的稀疏 SNode 使用流程是：

初始化：LLVMRuntime为每个 SNode 创建专属NodeManager（含data_list/free_list/recycled_list三个ListManager），并预分配共享的ambient_elements环境值区域；
激活：Pointer_activate/Dynamic_allocate通过双重检查锁 + warp 代表线程，从NodeManager::allocate()获取内存——优先复用free_list中的空闲索引，不足时经ListManager::reserve_new_element()按需分配新 chunk；
访问：lookup_element沿布局寻址；未激活单元读取环境值返回默认零值，保证稀疏访问的确定性；
去激活：指针/节点被置空，recycle()把释放的内存索引追加进recycled_list；
GC：offloaded task 结束后触发三阶段回收——并行压实free_list（阶段 0）→ 单线程备份并清空recycled_list（阶段 1）→ 并行回填free_list并清零回收内存（阶段 2），完成一轮"分配—回收—复用"的闭环。

这套设计让 Taichi 的 LLVM 后端（CPU / CUDA / AMDGPU）在只维护一份 C++ 运行时源码的前提下，实现了对稀疏数据结构的按需分配、内存复用与 GPU 并行垃圾回收，是 Taichi 能够以 Python 语法驾驭大规模稀疏计算的重要底层支撑。

延伸阅读：设计文档全文见 docs/design/llvm_sparse_runtime.md；运行时实现见 runtime.cpp，各 SNode 头文件见 node_dense.h、node_pointer.h、node_dynamic.h、node_bitmasked.h；并发原语见 locked_task.h；字节码编译流程见 CMakeLists.txt。

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

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Taichi LLVM 稀疏运行时深度解析：SNode 存储、回收式内存分配器与 GPU 垃圾回收