Taichi AOT 单元测试编程指南：从 Python 编译到 C++ 执行的「双段式」测试体系

Taichi AOT 单元测试编程指南：从 Python 编译到 C++ 执行的「双段式」测试体系

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

Taichi 的 Ahead-of-Time（AOT）工作流允许开发者在 Python 侧将 Kernel 与计算图编译并序列化为文件，再在脱离 Python 运行时的 C++ 应用中加载执行。为了保证这条跨语言、跨后端链路的正确性，Taichi 在tests/cpp/目录下建立了一套专门的 AOT 单元测试体系。本文以仓库内的 AOT_TEST_README.md 为核心，结合 cpptests.yaml 配置、conftest.py 运行器实现以及 LLVM / C-API / 图形后端（Vulkan、OpenGL、Metal、DX12）等真实测试用例，完整讲解这套测试的组成结构、配置格式、执行流程与编写方法。读完本文，你将能够：理解 AOT 测试「Python 编译 + C++ 执行」的双段式设计，熟练编写cpptests.yaml配置，并掌握测试选择、环境变量约定与运行命令。

一、AOT 测试的基本组成：为什么需要「双段式」结构

AOT（Ahead-of-Time）意味着编译与运行被刻意分离：编译阶段（Python 侧）产生可序列化的 AOT 模块文件，运行阶段（C++ 侧）读取这些文件并执行。因此，一个完整的 AOT 测试天然由两部分组成（见 AOT_TEST_README.md）：

1. Python 脚本：负责初始化 Taichi 运行时、定义并编译 Kernel（或计算图），将产物序列化写入磁盘文件；
2. C++ 测试：加载上述序列化文件，通过对应的运行时/加载器恢复 Kernel 或计算图，执行并校验计算结果。

两者之间通过测试用例运行器（test case runner）串联：运行器先执行 Python 脚本，再执行对应的 C++ 测试。这一设计与 Taichi AOT 的实际生产流程完全一致——cpp_examples/aot_save.cpp 展示了同样的「保存」思想，而 c_api/tests 下的 C-API 测试则是「加载执行」侧的规模化应用。

从源码结构看，AOT 测试用例按后端组织在 tests/cpp/aot 下：

• llvm/：CPU、CUDA、DX12 等 LLVM 后端测试（如 kernel_aot_test.cpp）；
• vulkan/、opengl/：图形后端测试；
• python_scripts/：所有与 C++ 测试配对的 Python 编译脚本（如 kernel_aot_test1.py）；
• gfx_utils.cpp 与 gfx_utils.h：图形后端共享的工具函数（run_cgraph1、run_kernel_test1、run_mpm88_graph等）。

二、cpptests.yaml配置格式详解

AOT 测试编写者需要将自己的用例登记在 cpptests.yaml 中，该文件是这套测试体系的「注册表」。其条目格式如下（来自 AOT_TEST_README.md）：

- name: Descriptive name for the binary (not affecting actual test execution) binary: path/to/compiled/gtest_binary tests: - test: GTestModule.GTestCaseName script: path/to/python_scripts.py args: --arg1 --arg2 markers: [marker1, marker2] - ...

顶层是一个列表，每个元素对应一组测试，字段说明：

需要特别强调两条规则：

• tests.test是唯一必填的用例级字段；
• tests.script与tests.args共同描述「前置编译脚本」，二者均是可选的——如果某个 gtest 用例本身能在测试内部完成模块构造（例如只测试加载/解析逻辑），就不需要 Python 脚本。

仓库中的真实配置示例

cpptests.yaml 中给出的完整示例与文档完全对应，且规模更大。第一组是 LLVM 与图形后端用例：

- name: C++ Tests binary: ../../build/taichi_cpp_tests tests: - test: LlvmAotTest.CpuKernel script: aot/python_scripts/kernel_aot_test1.py args: --arch=cpu - test: LlvmAotTest.CudaKernel script: aot/python_scripts/kernel_aot_test1.py args: --arch=cuda - test: LlvmAotTest.DX12Kernel script: aot/python_scripts/kernel_aot_test1.py args: --arch=dx12 - test: LlvmCGraph.RunGraphCpu script: aot/python_scripts/graph_aot_test_.py args: --arch=cpu - test: CGraphAotTest.VulkanMpm88 script: aot/python_scripts/mpm88_graph_aot.py args: --arch=vulkan --cgraph - test: GfxAotTest.VulkanDenseField script: aot/python_scripts/dense_field_aot_test_.py args: --arch=vulkan ...

第二组是 C-API 测试，并展示了markers的实际用法：

- name: C-API Tests binary: ../../build/taichi_c_api_tests tests: - test: CapiTest.Float16Fill script: aot/python_scripts/numerical_aot_test_.py args: --arch=vulkan markers: [sm70] - test: CapiTest.Float16Compute script: aot/python_scripts/numerical_aot_test_.py args: --arch=vulkan markers: [sm70] ...

第三组则通过 YAML 锚点复用同一份用例列表，指向静态链接的 C-API 测试二进制（cpptests.yaml）：

- name: C-API Tests (Static binary) binary: ../../build/taichi_static_c_api_tests tests: *c-api-tests

从该文件可以看出几个实战要点：

• binary使用相对路径（如../../build/taichi_cpp_tests），其基准目录是 tests/cpp，因此对应的绝对构建产物在仓库构建目录下；
• 同一个 Python 脚本可以被多个用例复用，仅通过--arch区分后端（如kernel_aot_test1.py同时服务 CPU、CUDA、DX12、Vulkan、OpenGL、Metal）；
• args可以携带附加开关，例如--cgraph表示同时导出计算图（CGraph）产物；
• 同一 YAML 组的用例可以覆盖多个后端，从而在同一个二进制上完成跨架构的 AOT 一致性验证。

三、运行器原理：pytest 如何驱动 C++ 测试

cpptests.yaml本身不会自动执行——它由 conftest.py 以 pytest 插件的形式解析并驱动。这正是文档所说 "The temporary directory where serialized cache file stays will be generated by test case runner" 的落点。

从源码看，运行器的工作流程如下（conftest.py）：

1. pytest_collect_file识别名为cpptests.yaml的文件，将其收集为CPPTestsFile；
2. collect()用yaml.safe_load解析配置，遍历每个组：若binary不存在则跳过整组；对每个用例创建CPPTestItem，并依据markers字段为测试项动态添加 pytest marker（conftest.py）；
3. 对于 YAML 中未显式列出的 gtest 用例，运行器还会调用binary --gtest_list_tests自动枚举二进制中的全部测试，以「默认配置」（无脚本、无参数）补全（conftest.py）——这保证新写的 gtest 用例即使忘记登记，也不会被静默漏跑。

每个用例的实际执行在CPPTestItem.runtest()中完成（conftest.py）：

with tempfile.TemporaryDirectory(prefix="ti-cpp-tests-") as tmpdir: env = os.environ.copy() env.update({ "TI_DEVICE_MEMORY_GB": "0.5", "TI_LIB_DIR": str(ti_lib_dir), "TAICHI_AOT_FOLDER_PATH": tmpdir, }) if self.script: subprocess.call(f"{sys.executable} {self.script} {self.args}", ...) subprocess.call(f"{self.binary} --gtest_filter={self.test}", ...)

可以看到运行器完成了三件关键工作：

• 为每个测试创建独立的临时目录（前缀ti-cpp-tests-）；
• 将临时目录通过环境变量TAICHI_AOT_FOLDER_PATH注入子进程，同时注入TI_LIB_DIR（指向已安装 Taichi 的运行时资源）与TI_DEVICE_MEMORY_GB=0.5（限制设备显存占用）；
• 依次执行「Python 脚本（带args）」→「gtest 二进制（--gtest_filter精确筛选单个用例）」，任何一步返回非零退出码即判失败。

这套由 pytest 统一编排的设计有一个额外收益：C++ 测试可以无缝参与 pytest 的用例选择、marker 过滤与失败报告，实现了 Python/C++ 测试的「一个入口」管理。

四、环境变量约定：TAICHI_AOT_FOLDER_PATH

这是连接「编译侧」与「执行侧」的唯一通道，也是 AOT 测试最重要的约定（AOT_TEST_README.md）：

The temporary directory where serialized cache file stays will be generated by test case runner. Both python program and C++ tests receives this directory path via environment variableTAICHI_AOT_FOLDER_PATH.

即：

• 临时目录由测试运行器创建，Python 编译脚本与 C++ 测试双方都通过环境变量TAICHI_AOT_FOLDER_PATH拿到该路径；
• Python 脚本必须把 AOT 模块保存到该目录；
• C++ 测试必须从该目录加载模块。

Python 侧的典型写法（见 kernel_aot_test1.py）：

assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) m = ti.aot.Module() m.add_kernel(run, template_args={"arr": arr}) m.save(dir_name)

C++ 侧的典型写法（见 kernel_aot_test.cpp）：

LLVM::AotModuleParams aot_params; const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH"); std::stringstream aot_mod_ss; aot_mod_ss << folder_dir; aot_params.module_path = aot_mod_ss.str(); aot_params.executor_ = &exec; aot_params.kernel_launcher = std::make_unique<cpu::KernelLauncher>(cpu::KernelLauncher::Config{&exec}); std::unique_ptr<aot::Module> mod = LLVM::make_aot_module(std::move(aot_params));

由于每个测试的临时目录独立创建，多个 AOT 测试并行运行也不会互相污染缓存文件。测试结束后，临时目录随TemporaryDirectory上下文自动清理，无需手工删除。

五、Python 侧：编写 AOT 编译脚本

5.1 Kernel 编译脚本

以 kernel_aot_test1.py 为最小范例，一个编译脚本需要完成四步：

1. 初始化运行时并指定架构：ti.init(arch=arch)，架构由--arch命令行参数映射而来（cpu→ti.cpu、cuda→ti.cuda、vulkan→ti.vulkan等）；
2. 定义 Kernel：例如一个逐元素写入的run(base, arr, v)Kernel，其中arr是ti.types.ndarray()模板参数，v是ti.types.vector(3, ti.i32)向量参数；
3. 构造 AOT 模块：ti.aot.Module()，通过add_kernel(run, template_args={"arr": arr})登记 Kernel 及其模板实参（ndarray 的形状/类型在这里被固化到产物中）；
4. 保存：m.save(dir_name)写入TAICHI_AOT_FOLDER_PATH指定的目录。

注意if ti.lang.impl.current_cfg().arch != arch: return这行防御逻辑：当请求的架构在当前环境不可用时直接返回，避免编译失败，同时让 C++ 侧（如CudaKernel用例中if (is_cuda_api_available())的守卫）自行决定是否跳过。

5.2 计算图（CGraph）编译脚本

除单个 Kernel 外，AOT 也支持序列化计算图。以 graph_aot_test_.py 为例，其步骤为：

1. 定义多个 Kernel（run0、run1）；
2. 声明图参数：ti.graph.Arg(ti.graph.ArgKind.NDARRAY, "arr0", dtype=ti.i32, ndim=1)与ti.graph.Arg(ti.graph.ArgKind.SCALAR, "base0", dtype=ti.i32)；
3. 通过ti.graph.GraphBuilder()多次dispatch将 Kernel 与图参数绑定（示例中对同一 Kernel 派发三次、使用不同的标量参数）；
4. g_builder.compile()得到run_graph，再mod.add_graph("run_graph", run_graph)与mod.save(tmpdir)完成序列化。

这解释了cpptests.yaml中args: --arch=cpu --cgraph的含义：带--cgraph的用例（如LlvmCGraph.CpuField）在 Python 侧额外导出计算图产物，供 C++ 侧以图形式加载。

六、C++ 侧：加载 AOT 模块并执行

C++ 测试端负责「从TAICHI_AOT_FOLDER_PATH加载模块 → 取回 Kernel → 构造启动上下文 → 执行并断言」。以LlvmAotTest.CpuKernel（kernel_aot_test.cpp）为例，完整链路是：

1. 构造CompileConfig并指定cfg.arch = Arch::x64；
2. 创建LlvmRuntimeExecutor，调用materialize_runtime实例化 LLVM 运行时；
3. 通过allocate_memory_on_device分配设备内存并包装成Ndarray；
4. 构造AotModuleParams（module_path取TAICHI_AOT_FOLDER_PATH，executor_与kernel_launcher指定 CPU 后端启动器），调用LLVM::make_aot_module加载模块；
5. mod->get_kernel("run")按名字取回 Kernel（名字与 Python 侧@ti.kernel def run对应）；
6. 用LaunchContextBuilder依次set_arg（标量）、set_arg_ndarray（ndarray）、set_struct_arg（向量分量）填充参数，然后k_run->launch(builder)；
7. 读取设备内存并EXPECT_EQ逐元素比对结果。

CUDA 用例的差异在于：使用cuda::KernelLauncher，且最终需通过CUDADriver::get_instance().memcpy_device_to_host将结果拷回主机再断言（kernel_aot_test.cpp）。而LlvmAotTest.DX12Kernel目前仅加载模块并EXPECT_TRUE(k_run)校验 Kernel 可获取（kernel_aot_test.cpp），注释中的FIXME表明其执行与结果校验部分仍在演进——这提示测试编写者在面对不完整后端支持时，可先做「加载级」冒烟断言。

图形后端（Vulkan / OpenGL）的加载逻辑集中在 gfx_utils.cpp，并通过 gfx_utils.h 暴露run_kernel_test1、run_dense_field_kernel、run_mpm88_graph等可复用入口，配合device_test.cpp、graph_aot_test.cpp、kernel_aot_test.cpp组织用例。

七、测试选择：markers 与-k选项

tests.markers的作用是用例选择（AOT_TEST_README.md）。它在实现上转化为 pytest 的 marker（conftest.py），因此可以直接用 pytest 的表达式进行筛选。

例如numerical_aot_test_.py的 Vulkan 用例被打上sm70标记（markers: [sm70]，见 cpptests.yaml），表示其数值语义仅对支持 FP16 的 sm70 及以上 GPU 有意义。运行时可结合-m（pytest 标记过滤）或-k（关键字匹配）精确圈定：

python tests/run_tests.py -c -m "sm70" # 只跑带 sm70 标记的 C++ 测试 python tests/run_tests.py -c -k "CudaKernel" # 按名称关键字筛选

run_tests.py的-c/--cpp选项专门用于「只运行 C++ 测试」，其实现是把测试目录指向tests/cpp后交给同一套 pytest 管线（run_tests.py），这也印证了文档中 "-koption ofrun_tests.py" 的说法。markers同时作为 CI 环境下的能力门控：不具备对应硬件/特性的机器可以排除相关 marker，避免误报。

八、运行 AOT 测试：从构建到执行

综合前文，AOT 测试的完整运行前提与命令如下：

1. 构建 gtest 二进制：cpptests.yaml中引用的taichi_cpp_tests、taichi_c_api_tests、taichi_static_c_api_tests需已通过 CMake 构建（对应 cmake/TaichiTests.cmake 与 cmake/TaichiCAPITests.cmake 等目标）；
2. 保证 Python 侧可导入taichi（编译脚本依赖import taichi as ti，且运行器会把已安装的_lib/runtime目录通过TI_LIB_DIR注入子进程）；
3. 执行测试：

# 运行全部 C++ 测试（含 AOT 用例） python tests/run_tests.py -c # 按 gtest 名称关键字筛选 python tests/run_tests.py -c -k "LlvmAotTest" # 仅运行带 sm70 marker 的用例 python tests/run_tests.py -c -m "sm70"

执行时，每个用例都会经历「创建临时目录 → 注入TAICHI_AOT_FOLDER_PATH等环境变量 → 运行 Python 编译脚本 → 运行--gtest_filter精确筛选的 C++ 测试」这一固定流程。若某一步失败，pytest 会报告对应阶段（Python 脚本退出码或 C++ 退出码），便于快速定位是「编译侧」还是「执行侧」的问题。

九、结语

AOT 测试是 Taichi「Python 开发、C++ 部署」工作流的质量防线。通过 AOT_TEST_README.md 描述的这套约定，开发者可以用最少的样板把任意 Kernel 或计算图的「编译 → 序列化 → 加载 → 执行 → 校验」纳入自动化回归：cpptests.yaml负责注册，conftest.py 负责驱动，TAICHI_AOT_FOLDER_PATH负责跨进程传递产物，而markers与-k让用例选择变得灵活可控。对于想要为 Taichi 贡献新 AOT 能力或新后端的开发者而言，参照本文的流程添加一个用例，通常只需：写一个 Python 编译脚本 → 写一个 gtest 用例 → 在cpptests.yaml登记三行配置，即可完成闭环。

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