Taichi 微基准测试套件（benchmarks）完全指南：安装、运行、结果解析与性能可视化

导读

Taichi 仓库的 benchmarks 目录提供了一套面向多后端（CUDA / Vulkan / OpenGL 等）的官方微基准测试（microbenchmark）框架，覆盖填充（fill）、内存拷贝（memcpy）、SAXPY、原子操作（atomic ops）、数学函数吞吐（math ops）、矩阵运算（matrix ops）与 2D 模板计算（stencil 2D）等典型负载。本文以 benchmarks/README.md 为主线，结合 benchmarks 目录下的源码实现，完整讲解这套基准测试的安装、运行、结果序列化与可视化流程，并深入剖析其测试用例生成机制与计时原理，帮助你在自己的机器上复现结果并读懂 Taichi 的性能数据。

---

一、benchmarks 目录结构与整体工作流

在动手运行之前，先了解 benchmarks 目录的组织方式，这有助于理解后续每一步命令背后发生了什么：

benchmarks/
├── README.md                  # 官方使用说明（本文主体）
├── requirements.txt           # 额外依赖：jsbeautifier、bokeh
├── run.py                     # 入口：运行全部基准测试并保存结果
├── deserialize.py             # 将分散的结果文件合并为单个 results.json
├── suite_microbenchmarks.py   # MicroBenchmark 套件定义与结果落盘逻辑
├── utils.py                   # JSON 美化输出、时间戳等公共工具
└── microbenchmarks/           # 各测试计划与基础设施
    ├── _plan.py               # BenchmarkPlan：用例生成与调度核心
    ├── _items.py              # 测试维度（数据类型、数据规模、容器、算子等）
    ├── _metric.py             # 计时指标（内核耗时 / 端到端耗时）
    ├── _utils.py              # 计时器、标签工具、架构映射、随机填充
    ├── atomic_ops.py          # 原子操作（归约）测试
    ├── fill.py                # 填充测试（含稀疏结构）
    ├── math_opts.py           # 一元数学函数吞吐测试
    ├── matrix_ops.py          # 矩阵加 / 乘 / 乘加测试
    ├── memcpy.py              # 内存拷贝测试
    ├── saxpy.py               # SAXPY 测试
    └── stencil2d.py           # 2D 模板计算（gather / scatter / BLS）

整体工作流为：安装依赖 → 运行 run.py 生成 ./results 目录 → 用 deserialize.py 合并为单个 JSON → 用可视化工具（基于 Bokeh）交互式查看性能数据。下面按这个流程逐一展开。

---

二、环境准备与额外依赖安装

按照 benchmarks/README.md 的说明，运行基准测试需要先安装少量额外依赖：

python3 -m pip install -r requirements.txt

其中 benchmarks/requirements.txt 内容只有两个包：

• jsbeautifier：用于美化 JSON 输出。在 benchmarks/utils.py 的 dump2json() 中，通过 jsbeautifier.beautify(json.dumps(...), options) 将结果格式化为缩进 4 空格的易读 JSON；
• bokeh：可视化工具库。文档中可视化脚本的默认地址为 localhost:5006/visualization，这正是 Bokeh Server 的默认端口（5006），说明可视化环节依赖 Bokeh 服务。

此外，基准测试脚本在 benchmarks/run.py 中导入 taichi._lib.core，因此你还需要一个可用的 Taichi Python 环境（即当前仓库对应的已安装版本）。建议在 Taichi 已正确安装的虚拟环境中执行上述命令。

---

三、运行全部基准测试

安装完依赖后，在 benchmarks 目录下执行：

python3 run.py

这是官方文档给出的唯一启动入口。根据 benchmarks/run.py 的源码，整个运行过程可以拆解为三步：

1. 记录环境信息：BenchmarkInfo 通过 ti_python_core.get_commit_hash() 获取当前 Taichi 的 commit hash，并记录带格式的时间戳（见 benchmarks/utils.py 的 datatime_with_format()，输出 ISO 格式时间）；
2. 创建并运行套件：BenchmarkSuites 实例化 benchmark_suites 列表中的套件并依次调用 run()。当前 benchmarks/run.py 中注册的套件只有一个 MicroBenchmark；
3. 落盘保存：suites.save(benchmark_dir) 将结果写入 os.path.join(os.getcwd(), "results")，最后把包含 commit hash、时间与套件信息的 _info.json 一并写入 results 目录。

注意 benchmarks/run.py 使用 os.getcwd() 拼接结果路径，因此结果文件夹 results 会生成在你执行命令时的当前目录下，而不是固定的 benchmarks/results。若希望在仓库内统一管理，可以先 cd 到目标目录再运行。

3.1 默认启用的后端

需要特别留意 benchmarks/suite_microbenchmarks.py 中 MicroBenchmark.config 的默认配置：

config = {
    "cuda":    {"enable": True},
    "vulkan":  {"enable": False},
    "opengl":  {"enable": False},
}

即默认只对 CUDA 后端执行测试，Vulkan 与 OpenGL 默认关闭。如果需要测试其他后端，需要修改该 config 字典将对应项置为 True（get_benchmark_info() 会据此生成启用的架构列表，run() 也只遍历 enable 为 True 的架构）。底层架构映射见 benchmarks/microbenchmarks/_utils.py 的 get_ti_arch()，它支持的键包括 cuda、vulkan、opengl、metal、x64、cc。换句话说，如果你当前机器没有可用的 CUDA 环境，需要自行将 cuda 关闭并把 vulkan/opengl（或 x64）打开后才能跑通。

3.2 结果落盘结构

运行完成后，results 目录的布局如下（由 benchmarks/suite_microbenchmarks.py 的 save_as_json() 决定）：

results/
├── _info.json                    # 顶层信息：commit_hash、datetime、suites
└── microbenchmarks/              # 套件目录（suite_name）
    └── cuda/                     # 架构目录
        ├── _info.json            # 该架构下各 case 的维度 tags 信息
        ├── atomic_ops.json       # 每个测试计划一个 JSON
        ├── fill.json
        ├── math_ops.json
        ├── matrix_ops.json
        ├── memcpy.json
        ├── saxpy.json
        └── stencil_2d.json

其中 run.py 最后写入的顶层 [results/_info.json] 包含 suites 字段，其结构为 {suite_name: {archs: [...]}}，供后续 deserialize.py 索引使用。

---

四、源码级剖析：微基准测试套件的内部机制

理解这套基准测试的输出格式，关键在于掌握 benchmarks/microbenchmarks/_plan.py 的用例生成机制与各测试计划的维度组合。

4.1 七个内置测试计划

benchmarks/microbenchmarks/init.py 中定义了 benchmark_plan_list，共 7 个测试计划：

计划类	测试主题	定义文件
AtomicOpsPlan	原子操作归约	microbenchmarks/atomic_ops.py
FillPlan	常量填充（含稀疏结构）	microbenchmarks/fill.py
MathOpsPlan	一元数学函数吞吐	microbenchmarks/math_opts.py
MatrixOpsPlan	矩阵加 / 乘 / 乘加	microbenchmarks/matrix_ops.py
MemcpyPlan	内存拷贝	microbenchmarks/memcpy.py
SaxpyPlan	SAXPY（z = 17*x + y）	microbenchmarks/saxpy.py
Stencil2DPlan	2D 模板计算（gather / scatter / BLS）	microbenchmarks/stencil2d.py

4.2 用例生成：维度笛卡尔积

每个计划继承 BenchmarkPlan，通过 create_plan(*items) 将多个**维度（BenchmarkItem）**做笛卡尔积生成全部用例（benchmarks/microbenchmarks/_plan.py）：

case_list = list(itertools.product(*items_list))

每个用例用 tags2name()（benchmarks/microbenchmarks/_utils.py）将标签列表用下划线拼接成唯一名称，例如 saxpy_field_f32_16KB_end2end_time_ms。所有用例存放在 self.plan 字典中，result 字段初始为 None，运行后被填充为实际测得的毫秒数。

Funcs 类实现了一套基于标签子集匹配的函数分发机制（benchmarks/microbenchmarks/_plan.py）：add_func(tag_list, func) 注册某个实现函数及其标签，get_func(tags) 返回标签是当前用例标签子集的第一个函数。这正是 fill、stencil_2d 等计划能够为 field、ndarray、sparse 分别提供不同 kernel 实现的原因。

4.3 公共测试维度

由 benchmarks/microbenchmarks/_items.py 定义，各计划共享以下维度：

• DataType（dtype）：i32、i64、f32、f64 四种类型。remove_integer() 可去掉整型（math_ops 只用浮点）；remove(["i64","f64"]) 可限制为 i32/f32（matrix_ops 如此）；
• DataSize（dsize）：按 size_bytes = (4**i) * 1024（i 取 2、4、6、8）生成 16KB、256KB、4MB、64MB 四档数据规模，标签由 size2tag()（benchmarks/microbenchmarks/_utils.py）格式化为 16KB 等；
• Container（container）：field（ti.field）与 ndarray（ti.ndarray）两种容器，fill 与 stencil_2d 还会额外追加 sparse（稀疏结构，实现为 None 占位，由专门函数处理）；
• MetricType（get_metric）：计时指标，详见 4.5 节。

4.4 各计划的特有维度与 kernel 实现

SAXPY（saxpy.py）：维度为 Container × DataType × DataSize × MetricType。kernel 为 z[i] = 17 * x[i] + y[i]，元素数为 dsize / dtype_size / 3（三份数组）。field 走 saxpy_field（模板参数），ndarray 走 saxpy_array（ti.types.ndarray() 参数）。

Memcpy（memcpy.py）：维度同上。kernel 为 dst[I] = src[I]（ti.grouped 遍历），元素数为 dsize / dtype_size / 2。

Fill（fill.py）：维度为 Container × DataType × DataSize × MetricType，container 额外含 sparse。稠密填充 kernel 将每个元素赋值为 ti.cast(0.7, dtype)；fill_sparse 使用 ti.root.pointer(...).dense(...) 构建稀疏结构，先通过 ti.activate(block, [i]) 激活全部块再填充，且 repeat 基数固定为 1。

AtomicOps（atomic_ops.py）：维度为 AtomicOps × Container × DataType × DataSize × MetricType。AtomicOps 完整包含 atomic_add / sub / and / or / xor / max / min 七种原子操作，但 AtomicOpsPlan 构造时 remove(["atomic_sub","atomic_and","atomic_xor","atomic_max"])，实际测试 atomic_add、atomic_or、atomic_min 三种。测试形态为原子归约：atomic_op(y[None], x[i])。由于 and/or/xor 是逻辑操作、只支持整型，_remove_conflict_items()（benchmarks/microbenchmarks/_plan.py）会在生成阶段剔除“逻辑原子操作 × 浮点类型”的非法组合。

MathOps（math_opts.py）：维度为 MathOps × DataType(浮点) × ElementNum × ForLoopCycle × MetricType。MathOps 覆盖 14 个一元函数：sin、cos、tan、asin、acos、tanh、sqrt、rsqrt、exp、log、round、floor、ceil、abs。ForLoopCycle 生成 8/16/32/64/128/256 六档内层循环次数；每个元素是一个 16 分量向量（local_data_num = 16），用于填满指令流水线。该测试关注的是一元算子的吞吐上限而非访存带宽。

MatrixOps（matrix_ops.py）：维度为 MatrixOps × BlockMN × ElementNum × DataType(i32/f32) × MetricType。MatrixOps 提供 mat_add（A+B）、mat_mul（A@B）、mat_mma（A@B+C）三种 @ti.func；BlockMN 为 1×1、2×2、3×3、4×4 四种分块；每个元素执行 2048 轮 × 4 次矩阵操作。

Stencil2D（stencil2d.py）：维度为 Scatter × BloclLocalStorage × Container × DataType × DataSize2D × MetricType，是维度最丰富的计划。Scatter 区分 scatter（y[I+offset] += x[I]）与 gether（gather，y[I] = Σ x[I+offset]）；BloclLocalStorage 控制是否开启块局部存储（BLS），dsize_2d 为 128×128、512×512、2048×2048、8192×8192 四档。sparse 分支通过 ti.block_local(x/y)、ti.block_dim(64) 实现带 BLS 的稀疏模板测试，且仅对 16KB～64MB 之间的规模生效（否则返回 None）。源码注释还引用了 tests/python/bls_test_template.py 作为 BLS 用法的参考实现。

4.5 计时指标：内核耗时与端到端耗时

benchmarks/microbenchmarks/_metric.py 定义了两种指标，两者都先执行若干次预热（compile & warmup），再正式计时：

• kernel_elapsed_time_ms：调用 ti.init(kernel_profiler=True, ...) 初始化，计时前 ti.profiler.clear_kernel_profiler_info()，随后用 ti.profiler.get_kernel_profiler_total_time() 取得纯 kernel 执行时间。它排除了启动与调度开销，只统计 GPU/后端内核实际执行耗时；
• end2end_time_ms：调用 ti.init(kernel_profiler=False, ...)，使用 benchmarks/microbenchmarks/_utils.py 的 End2EndTimer，基于 time.perf_counter() 并在 tick/tock 前后各调用一次 ti.sync()，统计包含启动开销在内的完整调用耗时。

两者最终都换算为单次平均毫秒数：total_time * 1000 / repeat。

另外，benchmarks/microbenchmarks/_utils.py 的 scaled_repeat_times() 会按环境放大重复次数以稳定测量：

if (arch == "cuda") | (arch == "vulkan") | (arch == "opengl"):
    repeat *= 10        # GPU 后端重复 10 倍
if datasize <= 4 * 1024 * 1024:
    repeat *= 10        # 数据 ≤ 4MB 时再重复 10 倍

所有计划的 basic_repeat_times 默认为 10（稀疏类用例固定为 1），小规模数据在 GPU 上最多会重复 1000 次，以降低计时噪声。

---

五、结果序列化：deserialize.py 合并为单个 JSON

run.py 生成的是分散在多个目录、多个文件中的结果。若需要把全部结果合并为单个 JSON 文件便于脚本分析或归档，使用 benchmarks/deserialize.py：

python3 deserialize.py

默认将合并结果写入 ./results/results.json。也可以显式指定输入结果目录与输出路径：

python3 deserialize.py --folder PATH_OF_RESULTS_FOLDER --output_path PATH_YOU_WIHS_TO_STORE

两个参数（benchmarks/deserialize.py）说明如下：

• -f, --folder：结果文件夹路径，默认 ./results；
• -o, --output_path：输出目录（最终生成 <output_path>/results.json），默认 ./results。

从源码看，ResultsBuilder 的合并逻辑分两层：

1. 读取顶层 _info.json 的 suites 字段，按 suite_name → arch → case 建立索引，并读取每个架构目录下的 _info.json 与各 case 的 JSON 文件；
2. 对每个 case 的结果，去掉首位的 case 名称标签（data["tags"] = data["tags"][1:]），并剔除 result 为 None 的条目（benchmarks/deserialize.py）——这正是稀疏模板等场景下因规模不适用而跳过用例的占位结果。

合并完成后，脚本还会调用 print_info() 打印一份去掉 results 字段的概要信息，方便你快速核对本次测试覆盖了哪些架构与用例。输出 JSON 统一使用 benchmarks/utils.py 的 dump2json() 做美化格式化。

---

六、性能可视化：交互式查看基准结果

拿到合并或分散的结果后，官方文档提供了一个基于 Bokeh 的可视化工具，用于交互式剖析性能问题：

python3 visualization.py

默认读取 ./results 目录，也可指定结果文件路径：

python3 visualization.py --folder PATH_OF_RESULTS_FOLDER

默认服务地址为 localhost:5006/visualization（5006 是 Bokeh Server 的默认端口，这也是 requirements.txt 需要安装 bokeh 的原因）。如需远程访问，可显式指定监听地址与端口：

python3 visualization.py --host YOUR_IP_ADDRESS --port PORT_YOU_WISH_TO_USE

需要说明的是：当前仓库快照的 benchmarks/ 目录中并未包含 visualization.py 脚本本体（文档保留了对该工具的使用说明），因此若你的检出中缺少该文件，可以：

• 先依赖 deserialize.py 生成的 results.json 做离线分析；
• 或在本地基于 Bokeh 自行实现同接口的可视化脚本（默认 --folder ./results、--host localhost、--port 5006）。

无论如何，results 与 results.json 中的结构化数据（每项包含 tags、result（毫秒）、指标类型等）都足够支撑自定义绘图。

---

七、扩展与自定义：从源码出发新增测试计划

这套框架的扩展点非常清晰，如果你想新增一个微基准测试，可以参考既有计划的写法：

1. 在 benchmarks/microbenchmarks/ 下新建模块，定义继承 BenchmarkPlan 的计划类，在 __init__ 中：
   • 调用 super().__init__("your_plan_name", arch, basic_repeat_times=...)；
   • 用 create_plan(...) 组合需要的维度（可复用 Container、DataType、DataSize、MetricType，或自定义 BenchmarkItem）；
   • 用 add_func(tag_list, func) 注册实现函数，必要时用 remove_cases_with_tags() 剔除不合理的组合；
2. 在 benchmarks/microbenchmarks/init.py 的 benchmark_plan_list 中注册新计划；
3. 按需在 benchmarks/suite_microbenchmarks.py 的 MicroBenchmark.config 中调整启用的后端。

实现函数遵循统一签名：def func(arch, repeat, ..., get_metric)，最终调用 get_metric(repeat, kernel, *args) 返回毫秒耗时（参考 benchmarks/microbenchmarks/saxpy.py）。

---

八、快速参考：命令速查表

目的	命令	说明
安装额外依赖	python3 -m pip install -r requirements.txt	安装 jsbeautifier 与 bokeh
运行全部基准	python3 run.py	生成 ./results 目录（默认仅 CUDA 后端）
合并结果为单 JSON	python3 deserialize.py	输出 ./results/results.json
指定输入输出	python3 deserialize.py --folder DIR --output_path OUT_DIR	自定义合并路径
启动可视化	python3 visualization.py	默认 localhost:5006/visualization
远程可视化	python3 visualization.py --host IP --port PORT	开放远程访问

运行前请确认：已安装与当前仓库匹配的 Taichi Python 包；benchmarks 目录是执行命令的工作目录（或注意 results 会写入当前工作目录）；本机具备默认启用的 CUDA 后端，否则需按第三节调整 MicroBenchmark.config。