Taichi 内部设计详解IR、SNode 数据结构与虚拟/物理索引机制【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi本文基于 Taichi 官方内部设计文档 docs/lang/articles/internals/internal.md 展开系统讲解 Taichi 运行时与编译器内部的四大核心主题中间表示IR的设计原则、基于 SNode 树的数据结构组织、稀疏数据结构的 List Generation列表生成并行遍历策略以及虚拟索引与物理索引的映射关系。读完本文你将能够看懂print_ir输出的内核 IR理解ti.root.dense/pointer/bitmasked背后的容器container、单元格cell与组件component语义并掌握用physical_index_position分析字段内存布局的方法。中间表示Intermediate Representation, IRTaichi 的计算 IR 是内核从 Python AST 到后端机器码之间的核心桥梁其设计具有以下四个显著特征静态单赋值Static-single assignment, SSA每个变量只被赋值一次便于数据流分析与优化层次化结构Hierarchical与 LLVM 那种控制流图 基本块的扁平结构不同Taichi 的 IR 是嵌套的、层次化的天然贴合 Python 前端for/if等块结构可微分DifferentiableIR 保留了足够的算子信息支持自动微分在 IR 层面求导静态强类型Statically and strongly typedIR 中每条语句的变量都带有明确的类型标注例如i32 $1。用 print_ir 观察一个内核的 IR只需在ti.init()中传入print_irTrue即可打印所有已实例化内核的 IR。例如下面这个简单内核# show_ir.py import taichi as ti ti.init(print_irTrue) ti.kernel def foo(): for i in range(10): if i 4: print(i) foo()它可能被编译为如下 IRkernel { $0 offloaded range_for(0, 10) grid_dim0 block_dim32 body { i32 $1 loop $0 index 0 i32 $2 const [4] i32 $3 cmp_lt $1 $2 i32 $4 const [1] i32 $5 bit_and $3 $4 $6 : if $5 { print $1, \n } } }观察这份 IR 可以看到几个关键信息$0 offloaded range_for(0, 10)表示一个被卸载offloaded到设备的range_for任务block_dim32指定了设备端的线程块大小循环变量$1 loop $0 index 0是 SSA 形式的一次性定义cmp_lt、bit_and等操作均为强类型i32语句if $5 { ... }是层次化嵌套的块结构体现了 Taichi IR 与 LLVM 基本块风格的关键差异。提示ti.init(print_irTrue)会打印所有实例化内核的 IR若只想查看单个内核可在运行该内核前后分别开启/关闭该开关。关于 Taichi JIT 编译系统的更多细节从 Python 源码到设备代码的完整流水线可以继续阅读 docs/lang/articles/internals/compilation.mdLife of a Taichi kernel。数据结构组织SNode 树的三个核心概念Taichi 数据结构的内部组织基于Structural Node结构节点简称SNode读音 /snōd/树系统。SNode 系统对新人开发者来说最容易混淆的地方在于必须严格区分三个概念——SNode容器container、SNode单元格cell和 SNode组件component。容器、单元格与组件的关系容器container可以包含多个单元格cell单元格的数量建议取 2 的幂power of two。例如S ti.root.dense(ti.i, 128)创建了一个 SNodeS每个S容器包含 128 个S单元格。单元格cell可以包含多个组件component。例如P S.dense(ti.i, 4); Q S.dense(ti.i, 4)向每个S单元格中插入两个组件一个P容器和一个Q容器。注意每个 SNode组件本身又是一个更低层级 SNode 的容器。因此Taichi 中无论是稠密还是稀疏的层次化数据结构本质上都是一棵容器层与单元格层交错的树。唯一例外是placeSNodeplace的容器不再拥有单元格而是直接存放数值。一个完整的 SNode 树示例考虑如下示例源自内部文档原始文件名为misc/listgen_demo.pyx ti.field(ti.i32) y ti.field(ti.i32) z ti.field(ti.i32) S0 ti.root S1 S0.pointer(ti.i, 4) S2 S1.dense(ti.i, 2) S2.place(x, y) # S3: x; S4: y S5 S1.dense(ti.i, 2) S5.place(z) # S6: z该数据结构的层次关系可展开为整个数据结构是一个S0root容器包含1 个S0root单元格它只有一个组件即一个S1pointer容器包含4 个S1pointer单元格每个单元格有两个组件即一个S2dense容器包含2 个S2dense单元格每个单元格有两个组件即一个S3place_x容器直接存放x: ti.i32值一个S4place_y容器直接存放y: ti.i32值一个S5dense容器包含2 个S5dense单元格每个单元格有一个组件即一个S6place容器直接存放z: ti.i32值内部文档原图data_structure_organization.png展示了上述层次及各级容器/单元格的index编号S0root容器与单元格没有 index。由于该图托管于外部站点本文以文字层级树代替图示含义完全等价。数量盘点容器与单元格按上述结构总结我们将拥有以下容器1 个S0root容器1 个S1pointer容器4 个S2dense容器4 个S5dense容器8 个S3place_x容器每个直接包含一个i32值8 个S4place_y容器每个直接包含一个i32值8 个S6place_z容器每个直接包含一个i32值以及以下单元格1 个S0root单元格4 个S1pointer单元格8 个S2dense单元格8 个S5dense单元格再次强调S3place_x、S4place_y和S6place_z容器没有对应的单元格。源码中的印证在taichi/ir/snode.h中SNode类直接体现了上述语义std::vectorstd::unique_ptrSNode ch表示子节点组件int64 num_cells_per_container{1}注释明确指出每个容器的单元格数量并在其上方引用了本文所依据的内部文档术语container/cellAxisExtractor extractors[taichi_max_num_indices]记录了每个轴从 root 起的元素个数、形状与累计形状见taichi/ir/snode.h中AxisExtractor结构体的num_elements_from_root、shape、acc_shape、active字段int depth{0}、int id{0}等字段用于标识节点在树中的位置。在各后端的 struct 编译器中每个 SNode 都有container类型与cell类型两种结构体。单元格cell永远不会暴露给最终用户List generation 生成的也是 SNode 容器container的列表而非单元格列表。术语说明Taichi 正在逐步移除children、instances、elements等歧义术语统一改用标准术语container容器、cell单元格、component组件。List generation稀疏结构上的并行遍历策略Taichi 的 struct-for 会并行遍历稀疏数据结构的所有活跃元素。在稀疏数据结构上实现负载均衡很困难如果朴素地把一棵不规则的树切成若干块很容易导致各分片叶子元素数量相差悬殊造成严重的负载不均。核心策略逐层生成活跃容器列表Taichi 的策略是逐层layer by layer生成活跃 SNode 容器的列表。List generation 计算与普通计算内核运行在同一个设备上具体取决于用户在ti.init()中传入的arch参数。List generation 将数据结构的叶子元素压平flatten成一个一维列表从而规避了不完整树的不规则性随后只需对这个一维列表执行一次常规的并行 for即可负载均衡问题随之解决。实例演示# misc/listgen_demo.py import taichi as ti ti.init(print_irTrue) x ti.field(ti.i32) S0 ti.root S1 S0.dense(ti.i, 4) S2 S1.bitmasked(ti.i, 4) S2.place(x) ti.kernel def func(): for i in x: print(i) func()该内核会生成如下 IR$0 offloaded clear_list S1dense $1 offloaded listgen S0root-S1dense $2 offloaded clear_list S2bitmasked $3 offloaded listgen S1dense-S2bitmasked $4 offloaded struct_for(S2bitmasked) block_dim0 { i32 x1 $5 loop index 0 print i, $5 }注意func触发了两次 list generation任务$0与$1基于唯一一个S0root容器的列表生成唯一一个S1dense容器的列表任务$2与$3基于S1dense容器的列表生成S2bitmasked容器的列表。其中S0rootSNode 的列表永远只有一个容器因此从不清理或重新生成S1dense的列表虽然也始终只有一个容器但为了流程统一仍会重新生成S2bitmasked的列表则有 4 个容器。为什么不为 place 叶子节点生成列表place叶子节点的列表永远不会被生成。相反struct-for 直接遍历其父节点的列表并在遍历每个父节点时即时on-the-fly枚举其中的place节点而不实际生成列表。该设计的动机是摊销 list generation 的开销如果为每个叶子place元素都生成一个列表元素代价会非常高甚至远超叶子元素上的实际计算量。因此只生成倒数第二层父节点的列表将列表生成成本分摊到该层节点下的多个子元素上。在上述示例中虽然共有 16 个x实例但我们只生成了 4 个S2bitmasked节点外加 1 个S1dense节点的列表。Statistics程序执行期间的内部事件统计在某些场景下收集 Taichi 程序执行期间内部事件的量化信息很有帮助Statistics类就是为此设计的。在 C 侧的使用方式如下#include taichi/util/statistics.h // 给计数器 codegen_offloaded_tasks 累加 1.0 taichi::stat.add(codegen_offloaded_tasks); // 把 ir 中的语句数量累加到计数器 codegen_statements taichi::stat.add(codegen_statements, irpass::analysis::count_statements(this-ir));注意统计的 key 是std::stringvalue 是double。在 Python 侧可以通过以下方式打印全部统计信息ti.core.print_stat()该机制常用于编译器开发调试例如统计代码生成阶段生成的 offload 任务数、IR 语句数等帮助定位 IR 变换与代码生成阶段的性能或行为问题。为什么选择 Python 前端将 Taichi 嵌入 Python 具有以下优势易于学习Taichi 的语法与 Python 非常接近易于运行无需提前编译no ahead-of-time compilation复用现有 Python 生态IDEPython IDE 对 Taichi 基本开箱即用支持语法高亮、语法检查和自动补全包管理器pip开发好的 Taichi 应用可以轻松提交到 PyPI其他人用pip即可安装现有包与matplotlib、numpy等其他 Python 组件交互非常简单。AST 工具链Python 内置的 AST 操作工具允许我们灵活地操作和分析 Python AST——只要内核函数体能被 Python 解析器解析即可。当然这一设计也有缺点内核必须能被 Python 解析器解析这意味着 Taichi 语法不能超出 Python 语法范畴。例如访问 Taichi 字段元素时始终需要索引即使是 0D 字段也要用x[None] 123来赋值因为x 123在 Python 语法中是把x本身而非其内部值赋为常量123。为了保持 Python 作用域与 Taichi 作用域的代码一致性只能采用更啰嗦的x[None] 123写法。Python 本身性能较低当用纯 Python 脚本初始化大型 Taichi 字段时可能带来性能问题——大型字段应当用 Taichi 内核来初始化。虚拟索引virtual indices与物理索引physical indices在 Taichi 中虚拟索引用于定位字段中的元素物理索引用于指定内存中的数据布局layout。两者是不同层面的概念。概念与示例在a[i, j, k]中i、j、k是虚拟索引在for i, j in x:中i和j是虚拟索引ti.i, ti.j, ti.k, ti.l, ...是物理索引在 struct-for 语句中LoopIndexStmt::index是物理索引。每个SNode的虚拟索引与物理索引之间的映射关系存储在SNode::physical_index_position中。即physical_index_position[i]回答的问题是第 i 个虚拟索引对应哪个物理索引每个SNode可以拥有不同的虚拟到物理映射physical_index_position[i] -1表示第i个虚拟索引在该SNode中不对应任何物理索引。不同布局下的映射差异常见的稠密字段如a ti.field(ti.i32, shape(128, 256, 512))的 SNode 拥有**平凡trivial**的虚拟到物理映射例如physical_index_position[i] i。而更复杂的数据布局例如列主序column-major的二维字段则可能产生physical_index_position[0] 1且physical_index_position[1] 0的映射。以下代码来自内部文档演示了如何通过内部 API 验证映射a ti.field(ti.f32, shape(128, 32, 8)) b ti.field(ti.f32) ti.root.dense(ti.j, 32).dense(ti.i, 16).place(b) ti.lang.impl.get_runtime().materialize() # 这是供开发者使用的内部 API不保证对用户稳定 mapping_a a.snode().physical_index_position() assert mapping_a {0: 0, 1: 1, 2: 2} mapping_b b.snode().physical_index_position() assert mapping_b {0: 1, 1: 0} # 注意 b 是列主序的 # 暴露给用户的虚拟第 0 维索引在内存布局中位于第 1 维。上面第二个断言说明字段b先按ti.j划分 32、再按ti.i划分 16导致用户眼中的第 0 个虚拟索引i对应物理索引ti.j值为 1第 1 个虚拟索引j对应物理索引ti.i值为 0从而实现了列主序的内存布局。索引数量上限Taichi 最多支持12个虚拟索引与物理索引。该上限在源码中定义为constexpr int taichi_max_num_indices 12见 taichi/inc/constants.h并在以下位置被广泛使用taichi/ir/snode.h 中AxisExtractor extractors[taichi_max_num_indices]与int physical_index_position[taichi_max_num_indices]数组的定义Axis构造器会校验维度值必须满足0 value taichi_max_num_indices超出即报错Too many dimensionstaichi/ir/snode.cpp 中SNode::create_node在创建子节点时根据传入轴写入physical_index_position、激活对应AxisExtractor并在该轴已激活但尺寸非 2 的幂时给出性能警告建议设为 2 的幂映射数组随后会被排序见 taichi/ir/snode.cppPython 侧通过 taichi/python/export_lang.cpp 将physical_index_position导出为 Python 可调用的接口即上文示例中的a.snode().physical_index_position()并提供了get_max_num_indices查询接口taichi/python/export_lang.cpp。这些源码实现印证了文档所述每个 SNode 独立维护自己的虚拟到物理映射映射由创建节点时的轴声明dense(ti.i, ...)/dense(ti.j, ...)顺序决定从而让用户能够以字段逻辑维度编程同时自由控制底层内存布局——这正是 Taichi 兼顾生产力与性能的关键设计之一。小结本文围绕 docs/lang/articles/internals/internal.md 这条主线梳理了 Taichi 内部设计最核心的四个主题具备 SSA、层次化、可微分、静态强类型特性的计算 IR以容器/单元格/组件三层语义组织的 SNode 树通过逐层 list generation 将不规则稀疏树压平为一维列表从而并行化的 struct-for 机制以及虚拟索引与物理索引之间的映射规则与 12 维上限。配合print_irTrue、ti.core.print_stat()等调试手段以及physical_index_position()等内部接口开发者可以深入观察内核编译产物与字段内存布局为进一步阅读编译器实现如 docs/lang/articles/internals/compilation.md或参与 Taichi 内核开发打下基础。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Taichi 内部设计详解：IR、SNode 数据结构与虚拟/物理索引机制
Taichi 内部设计详解IR、SNode 数据结构与虚拟/物理索引机制【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi本文基于 Taichi 官方内部设计文档 docs/lang/articles/internals/internal.md 展开系统讲解 Taichi 运行时与编译器内部的四大核心主题中间表示IR的设计原则、基于 SNode 树的数据结构组织、稀疏数据结构的 List Generation列表生成并行遍历策略以及虚拟索引与物理索引的映射关系。读完本文你将能够看懂print_ir输出的内核 IR理解ti.root.dense/pointer/bitmasked背后的容器container、单元格cell与组件component语义并掌握用physical_index_position分析字段内存布局的方法。中间表示Intermediate Representation, IRTaichi 的计算 IR 是内核从 Python AST 到后端机器码之间的核心桥梁其设计具有以下四个显著特征静态单赋值Static-single assignment, SSA每个变量只被赋值一次便于数据流分析与优化层次化结构Hierarchical与 LLVM 那种控制流图 基本块的扁平结构不同Taichi 的 IR 是嵌套的、层次化的天然贴合 Python 前端for/if等块结构可微分DifferentiableIR 保留了足够的算子信息支持自动微分在 IR 层面求导静态强类型Statically and strongly typedIR 中每条语句的变量都带有明确的类型标注例如i32 $1。用 print_ir 观察一个内核的 IR只需在ti.init()中传入print_irTrue即可打印所有已实例化内核的 IR。例如下面这个简单内核# show_ir.py import taichi as ti ti.init(print_irTrue) ti.kernel def foo(): for i in range(10): if i 4: print(i) foo()它可能被编译为如下 IRkernel { $0 offloaded range_for(0, 10) grid_dim0 block_dim32 body { i32 $1 loop $0 index 0 i32 $2 const [4] i32 $3 cmp_lt $1 $2 i32 $4 const [1] i32 $5 bit_and $3 $4 $6 : if $5 { print $1, \n } } }观察这份 IR 可以看到几个关键信息$0 offloaded range_for(0, 10)表示一个被卸载offloaded到设备的range_for任务block_dim32指定了设备端的线程块大小循环变量$1 loop $0 index 0是 SSA 形式的一次性定义cmp_lt、bit_and等操作均为强类型i32语句if $5 { ... }是层次化嵌套的块结构体现了 Taichi IR 与 LLVM 基本块风格的关键差异。提示ti.init(print_irTrue)会打印所有实例化内核的 IR若只想查看单个内核可在运行该内核前后分别开启/关闭该开关。关于 Taichi JIT 编译系统的更多细节从 Python 源码到设备代码的完整流水线可以继续阅读 docs/lang/articles/internals/compilation.mdLife of a Taichi kernel。数据结构组织SNode 树的三个核心概念Taichi 数据结构的内部组织基于Structural Node结构节点简称SNode读音 /snōd/树系统。SNode 系统对新人开发者来说最容易混淆的地方在于必须严格区分三个概念——SNode容器container、SNode单元格cell和 SNode组件component。容器、单元格与组件的关系容器container可以包含多个单元格cell单元格的数量建议取 2 的幂power of two。例如S ti.root.dense(ti.i, 128)创建了一个 SNodeS每个S容器包含 128 个S单元格。单元格cell可以包含多个组件component。例如P S.dense(ti.i, 4); Q S.dense(ti.i, 4)向每个S单元格中插入两个组件一个P容器和一个Q容器。注意每个 SNode组件本身又是一个更低层级 SNode 的容器。因此Taichi 中无论是稠密还是稀疏的层次化数据结构本质上都是一棵容器层与单元格层交错的树。唯一例外是placeSNodeplace的容器不再拥有单元格而是直接存放数值。一个完整的 SNode 树示例考虑如下示例源自内部文档原始文件名为misc/listgen_demo.pyx ti.field(ti.i32) y ti.field(ti.i32) z ti.field(ti.i32) S0 ti.root S1 S0.pointer(ti.i, 4) S2 S1.dense(ti.i, 2) S2.place(x, y) # S3: x; S4: y S5 S1.dense(ti.i, 2) S5.place(z) # S6: z该数据结构的层次关系可展开为整个数据结构是一个S0root容器包含1 个S0root单元格它只有一个组件即一个S1pointer容器包含4 个S1pointer单元格每个单元格有两个组件即一个S2dense容器包含2 个S2dense单元格每个单元格有两个组件即一个S3place_x容器直接存放x: ti.i32值一个S4place_y容器直接存放y: ti.i32值一个S5dense容器包含2 个S5dense单元格每个单元格有一个组件即一个S6place容器直接存放z: ti.i32值内部文档原图data_structure_organization.png展示了上述层次及各级容器/单元格的index编号S0root容器与单元格没有 index。由于该图托管于外部站点本文以文字层级树代替图示含义完全等价。数量盘点容器与单元格按上述结构总结我们将拥有以下容器1 个S0root容器1 个S1pointer容器4 个S2dense容器4 个S5dense容器8 个S3place_x容器每个直接包含一个i32值8 个S4place_y容器每个直接包含一个i32值8 个S6place_z容器每个直接包含一个i32值以及以下单元格1 个S0root单元格4 个S1pointer单元格8 个S2dense单元格8 个S5dense单元格再次强调S3place_x、S4place_y和S6place_z容器没有对应的单元格。源码中的印证在taichi/ir/snode.h中SNode类直接体现了上述语义std::vectorstd::unique_ptrSNode ch表示子节点组件int64 num_cells_per_container{1}注释明确指出每个容器的单元格数量并在其上方引用了本文所依据的内部文档术语container/cellAxisExtractor extractors[taichi_max_num_indices]记录了每个轴从 root 起的元素个数、形状与累计形状见taichi/ir/snode.h中AxisExtractor结构体的num_elements_from_root、shape、acc_shape、active字段int depth{0}、int id{0}等字段用于标识节点在树中的位置。在各后端的 struct 编译器中每个 SNode 都有container类型与cell类型两种结构体。单元格cell永远不会暴露给最终用户List generation 生成的也是 SNode 容器container的列表而非单元格列表。术语说明Taichi 正在逐步移除children、instances、elements等歧义术语统一改用标准术语container容器、cell单元格、component组件。List generation稀疏结构上的并行遍历策略Taichi 的 struct-for 会并行遍历稀疏数据结构的所有活跃元素。在稀疏数据结构上实现负载均衡很困难如果朴素地把一棵不规则的树切成若干块很容易导致各分片叶子元素数量相差悬殊造成严重的负载不均。核心策略逐层生成活跃容器列表Taichi 的策略是逐层layer by layer生成活跃 SNode 容器的列表。List generation 计算与普通计算内核运行在同一个设备上具体取决于用户在ti.init()中传入的arch参数。List generation 将数据结构的叶子元素压平flatten成一个一维列表从而规避了不完整树的不规则性随后只需对这个一维列表执行一次常规的并行 for即可负载均衡问题随之解决。实例演示# misc/listgen_demo.py import taichi as ti ti.init(print_irTrue) x ti.field(ti.i32) S0 ti.root S1 S0.dense(ti.i, 4) S2 S1.bitmasked(ti.i, 4) S2.place(x) ti.kernel def func(): for i in x: print(i) func()该内核会生成如下 IR$0 offloaded clear_list S1dense $1 offloaded listgen S0root-S1dense $2 offloaded clear_list S2bitmasked $3 offloaded listgen S1dense-S2bitmasked $4 offloaded struct_for(S2bitmasked) block_dim0 { i32 x1 $5 loop index 0 print i, $5 }注意func触发了两次 list generation任务$0与$1基于唯一一个S0root容器的列表生成唯一一个S1dense容器的列表任务$2与$3基于S1dense容器的列表生成S2bitmasked容器的列表。其中S0rootSNode 的列表永远只有一个容器因此从不清理或重新生成S1dense的列表虽然也始终只有一个容器但为了流程统一仍会重新生成S2bitmasked的列表则有 4 个容器。为什么不为 place 叶子节点生成列表place叶子节点的列表永远不会被生成。相反struct-for 直接遍历其父节点的列表并在遍历每个父节点时即时on-the-fly枚举其中的place节点而不实际生成列表。该设计的动机是摊销 list generation 的开销如果为每个叶子place元素都生成一个列表元素代价会非常高甚至远超叶子元素上的实际计算量。因此只生成倒数第二层父节点的列表将列表生成成本分摊到该层节点下的多个子元素上。在上述示例中虽然共有 16 个x实例但我们只生成了 4 个S2bitmasked节点外加 1 个S1dense节点的列表。Statistics程序执行期间的内部事件统计在某些场景下收集 Taichi 程序执行期间内部事件的量化信息很有帮助Statistics类就是为此设计的。在 C 侧的使用方式如下#include taichi/util/statistics.h // 给计数器 codegen_offloaded_tasks 累加 1.0 taichi::stat.add(codegen_offloaded_tasks); // 把 ir 中的语句数量累加到计数器 codegen_statements taichi::stat.add(codegen_statements, irpass::analysis::count_statements(this-ir));注意统计的 key 是std::stringvalue 是double。在 Python 侧可以通过以下方式打印全部统计信息ti.core.print_stat()该机制常用于编译器开发调试例如统计代码生成阶段生成的 offload 任务数、IR 语句数等帮助定位 IR 变换与代码生成阶段的性能或行为问题。为什么选择 Python 前端将 Taichi 嵌入 Python 具有以下优势易于学习Taichi 的语法与 Python 非常接近易于运行无需提前编译no ahead-of-time compilation复用现有 Python 生态IDEPython IDE 对 Taichi 基本开箱即用支持语法高亮、语法检查和自动补全包管理器pip开发好的 Taichi 应用可以轻松提交到 PyPI其他人用pip即可安装现有包与matplotlib、numpy等其他 Python 组件交互非常简单。AST 工具链Python 内置的 AST 操作工具允许我们灵活地操作和分析 Python AST——只要内核函数体能被 Python 解析器解析即可。当然这一设计也有缺点内核必须能被 Python 解析器解析这意味着 Taichi 语法不能超出 Python 语法范畴。例如访问 Taichi 字段元素时始终需要索引即使是 0D 字段也要用x[None] 123来赋值因为x 123在 Python 语法中是把x本身而非其内部值赋为常量123。为了保持 Python 作用域与 Taichi 作用域的代码一致性只能采用更啰嗦的x[None] 123写法。Python 本身性能较低当用纯 Python 脚本初始化大型 Taichi 字段时可能带来性能问题——大型字段应当用 Taichi 内核来初始化。虚拟索引virtual indices与物理索引physical indices在 Taichi 中虚拟索引用于定位字段中的元素物理索引用于指定内存中的数据布局layout。两者是不同层面的概念。概念与示例在a[i, j, k]中i、j、k是虚拟索引在for i, j in x:中i和j是虚拟索引ti.i, ti.j, ti.k, ti.l, ...是物理索引在 struct-for 语句中LoopIndexStmt::index是物理索引。每个SNode的虚拟索引与物理索引之间的映射关系存储在SNode::physical_index_position中。即physical_index_position[i]回答的问题是第 i 个虚拟索引对应哪个物理索引每个SNode可以拥有不同的虚拟到物理映射physical_index_position[i] -1表示第i个虚拟索引在该SNode中不对应任何物理索引。不同布局下的映射差异常见的稠密字段如a ti.field(ti.i32, shape(128, 256, 512))的 SNode 拥有**平凡trivial**的虚拟到物理映射例如physical_index_position[i] i。而更复杂的数据布局例如列主序column-major的二维字段则可能产生physical_index_position[0] 1且physical_index_position[1] 0的映射。以下代码来自内部文档演示了如何通过内部 API 验证映射a ti.field(ti.f32, shape(128, 32, 8)) b ti.field(ti.f32) ti.root.dense(ti.j, 32).dense(ti.i, 16).place(b) ti.lang.impl.get_runtime().materialize() # 这是供开发者使用的内部 API不保证对用户稳定 mapping_a a.snode().physical_index_position() assert mapping_a {0: 0, 1: 1, 2: 2} mapping_b b.snode().physical_index_position() assert mapping_b {0: 1, 1: 0} # 注意 b 是列主序的 # 暴露给用户的虚拟第 0 维索引在内存布局中位于第 1 维。上面第二个断言说明字段b先按ti.j划分 32、再按ti.i划分 16导致用户眼中的第 0 个虚拟索引i对应物理索引ti.j值为 1第 1 个虚拟索引j对应物理索引ti.i值为 0从而实现了列主序的内存布局。索引数量上限Taichi 最多支持12个虚拟索引与物理索引。该上限在源码中定义为constexpr int taichi_max_num_indices 12见 taichi/inc/constants.h并在以下位置被广泛使用taichi/ir/snode.h 中AxisExtractor extractors[taichi_max_num_indices]与int physical_index_position[taichi_max_num_indices]数组的定义Axis构造器会校验维度值必须满足0 value taichi_max_num_indices超出即报错Too many dimensionstaichi/ir/snode.cpp 中SNode::create_node在创建子节点时根据传入轴写入physical_index_position、激活对应AxisExtractor并在该轴已激活但尺寸非 2 的幂时给出性能警告建议设为 2 的幂映射数组随后会被排序见 taichi/ir/snode.cppPython 侧通过 taichi/python/export_lang.cpp 将physical_index_position导出为 Python 可调用的接口即上文示例中的a.snode().physical_index_position()并提供了get_max_num_indices查询接口taichi/python/export_lang.cpp。这些源码实现印证了文档所述每个 SNode 独立维护自己的虚拟到物理映射映射由创建节点时的轴声明dense(ti.i, ...)/dense(ti.j, ...)顺序决定从而让用户能够以字段逻辑维度编程同时自由控制底层内存布局——这正是 Taichi 兼顾生产力与性能的关键设计之一。小结本文围绕 docs/lang/articles/internals/internal.md 这条主线梳理了 Taichi 内部设计最核心的四个主题具备 SSA、层次化、可微分、静态强类型特性的计算 IR以容器/单元格/组件三层语义组织的 SNode 树通过逐层 list generation 将不规则稀疏树压平为一维列表从而并行化的 struct-for 机制以及虚拟索引与物理索引之间的映射规则与 12 维上限。配合print_irTrue、ti.core.print_stat()等调试手段以及physical_index_position()等内部接口开发者可以深入观察内核编译产物与字段内存布局为进一步阅读编译器实现如 docs/lang/articles/internals/compilation.md或参与 Taichi 内核开发打下基础。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考