Taichi LLVM 稀疏运行时深度解析SNode 存储、回收式内存分配器与 GPU 垃圾回收【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 的 LLVM 稀疏运行时sparse runtime是支撑pointer、dynamic等稀疏 SNodeStructural 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 -stdc17 -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.cppstruct 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_tu1对应布尔值。meta指向对应的 meta 结构体node指向 SNode 实例本身。StructMeta定义了所有 SNode 共享的核心属性与接口成员含义snode_idSNode 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 可能返回nullptrrefine_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.denseSNode最简单的连续数组dense是最简单的 SNode 形态它就是一段连续内存中的 cell 数组用 C 类比即std::arrayCell, 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_elementsDense_activate空操作——dense的 cell 永远处于激活状态无需分配动作Dense_is_active恒返回1Dense_lookup_element返回第i个 cell 的地址即node element_size * i是纯算术寻址没有任何间接跳转。内存布局示意如下- node | ------------------------------------------------ | | | | | | cell-0 | cell-1 | cell-2 | cell-3 | | | | | | ------------------------------------------------由于dense不存在未激活单元它的访问开销最低是 Taichi 中默认、最高效的存储形式bitmasked则在dense的基础上附加一个位掩码记录激活状态但存储仍是连续的本文不再展开。4.pointerSNode按需分配 环境值兜底的稀疏存储pointer是稀疏计算中最常用的选择。它只在实际激活 cell 时从内存池动态分配空间去激活时把内存回收进池中复用从而在大规模网格计算中显著节省内存。用 C 类比它可以看作std::arrayCell*, 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.1Pointer_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.dynamicSNode变长向量与分块链表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::vectorint32_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.1Dynamic_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 的地址否则沿链表跳到下一个 chunkchunk_start累加chunk_size。与激活/去激活相关的配套函数还包括Dynamic_activate将node-n原子地更新为max(n, i1)并确保第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_sizesSNode 内存分配器数组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前文提到的各种 chunkcell、dynamic chunk都来自这里free_list空闲 cell 的索引列表每个节点是一个int32_t即list_data_type。分配时优先复用其中的索引不够才向内存分配器申请新空间recycled_list已释放 cell 的索引列表。一次 GC 执行后其中的条目会被转入free_list复用每个节点同样是一个int32_tNodeManager构造时有两个值得注意的参数化逻辑默认每个 chunk 容纳128 * 1024128K个元素同时保证单块不超过 128 MB若chunk_num_elements * element_size 128 MB则不断将chunk_num_elements减半。7.1allocate()先复用后扩容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-getlist_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.ListManagerCPU/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::dequeListManager在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.1reserve_new_element()与touch_chunk()分配新元素的关键方法是reserve_new_element()runtime.cppi32 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.cppPtr 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.cppreserve_new_element()后返回get_element_ptr(i)append(data_ptr)runtime.cppallocate()后std::memcpy拷贝element_size字节clear()runtime.cpp仅把num_elements重置为 0不触碰列表内容——这个惰性清空特性在后续 GC 阶段被巧妙利用。9. 垃圾回收GCGPU 上的三阶段并行回收当某个 offloaded task 可能涉及稀疏 SNode 去激活之后就会触发 GC在 GPU 上以并行方式执行。整个 GC 过程针对每个 SNode 分为三个阶段。9.1 阶段 0gc_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 阶段 1gc_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 阶段 2gc_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-getT(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_listGCoffloaded 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 垃圾回收
Taichi LLVM 稀疏运行时深度解析SNode 存储、回收式内存分配器与 GPU 垃圾回收【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 的 LLVM 稀疏运行时sparse runtime是支撑pointer、dynamic等稀疏 SNodeStructural 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 -stdc17 -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.cppstruct 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_tu1对应布尔值。meta指向对应的 meta 结构体node指向 SNode 实例本身。StructMeta定义了所有 SNode 共享的核心属性与接口成员含义snode_idSNode 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 可能返回nullptrrefine_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.denseSNode最简单的连续数组dense是最简单的 SNode 形态它就是一段连续内存中的 cell 数组用 C 类比即std::arrayCell, 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_elementsDense_activate空操作——dense的 cell 永远处于激活状态无需分配动作Dense_is_active恒返回1Dense_lookup_element返回第i个 cell 的地址即node element_size * i是纯算术寻址没有任何间接跳转。内存布局示意如下- node | ------------------------------------------------ | | | | | | cell-0 | cell-1 | cell-2 | cell-3 | | | | | | ------------------------------------------------由于dense不存在未激活单元它的访问开销最低是 Taichi 中默认、最高效的存储形式bitmasked则在dense的基础上附加一个位掩码记录激活状态但存储仍是连续的本文不再展开。4.pointerSNode按需分配 环境值兜底的稀疏存储pointer是稀疏计算中最常用的选择。它只在实际激活 cell 时从内存池动态分配空间去激活时把内存回收进池中复用从而在大规模网格计算中显著节省内存。用 C 类比它可以看作std::arrayCell*, 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.1Pointer_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.dynamicSNode变长向量与分块链表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::vectorint32_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.1Dynamic_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 的地址否则沿链表跳到下一个 chunkchunk_start累加chunk_size。与激活/去激活相关的配套函数还包括Dynamic_activate将node-n原子地更新为max(n, i1)并确保第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_sizesSNode 内存分配器数组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前文提到的各种 chunkcell、dynamic chunk都来自这里free_list空闲 cell 的索引列表每个节点是一个int32_t即list_data_type。分配时优先复用其中的索引不够才向内存分配器申请新空间recycled_list已释放 cell 的索引列表。一次 GC 执行后其中的条目会被转入free_list复用每个节点同样是一个int32_tNodeManager构造时有两个值得注意的参数化逻辑默认每个 chunk 容纳128 * 1024128K个元素同时保证单块不超过 128 MB若chunk_num_elements * element_size 128 MB则不断将chunk_num_elements减半。7.1allocate()先复用后扩容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-getlist_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.ListManagerCPU/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::dequeListManager在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.1reserve_new_element()与touch_chunk()分配新元素的关键方法是reserve_new_element()runtime.cppi32 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.cppPtr 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.cppreserve_new_element()后返回get_element_ptr(i)append(data_ptr)runtime.cppallocate()后std::memcpy拷贝element_size字节clear()runtime.cpp仅把num_elements重置为 0不触碰列表内容——这个惰性清空特性在后续 GC 阶段被巧妙利用。9. 垃圾回收GCGPU 上的三阶段并行回收当某个 offloaded task 可能涉及稀疏 SNode 去激活之后就会触发 GC在 GPU 上以并行方式执行。整个 GC 过程针对每个 SNode 分为三个阶段。9.1 阶段 0gc_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 阶段 1gc_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 阶段 2gc_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-getT(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_listGCoffloaded 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），仅供参考