Taichi 微基准测试套件benchmarks完全指南安装、运行、结果解析与性能可视化【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi导读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 空格的易读 JSONbokeh可视化工具库。文档中可视化脚本的默认地址为localhost:5006/visualization这正是 Bokeh Server 的默认端口5006说明可视化环节依赖 Bokeh 服务。此外基准测试脚本在 benchmarks/run.py 中导入taichi._lib.core因此你还需要一个可用的 Taichi Python 环境即当前仓库对应的已安装版本。建议在 Taichi 已正确安装的虚拟环境中执行上述命令。三、运行全部基准测试安装完依赖后在benchmarks目录下执行python3 run.py这是官方文档给出的唯一启动入口。根据 benchmarks/run.py 的源码整个运行过程可以拆解为三步记录环境信息BenchmarkInfo通过ti_python_core.get_commit_hash()获取当前 Taichi 的 commit hash并记录带格式的时间戳见 benchmarks/utils.py 的datatime_with_format()输出 ISO 格式时间创建并运行套件BenchmarkSuites实例化benchmark_suites列表中的套件并依次调用run()。当前 benchmarks/run.py 中注册的套件只有一个MicroBenchmark落盘保存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字典将对应项置为Trueget_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.pyFillPlan常量填充含稀疏结构microbenchmarks/fill.pyMathOpsPlan一元数学函数吞吐microbenchmarks/math_opts.pyMatrixOpsPlan矩阵加 / 乘 / 乘加microbenchmarks/matrix_ops.pyMemcpyPlan内存拷贝microbenchmarks/memcpy.pySaxpyPlanSAXPYz 17*x ymicrobenchmarks/saxpy.pyStencil2DPlan2D 模板计算gather / scatter / BLSmicrobenchmarks/stencil2d.py4.2 用例生成维度笛卡尔积每个计划继承BenchmarkPlan通过create_plan(*items)将多个**维度BenchmarkItem**做笛卡尔积生成全部用例benchmarks/microbenchmarks/_plan.pycase_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.pyadd_func(tag_list, func)注册某个实现函数及其标签get_func(tags)返回标签是当前用例标签子集的第一个函数。这正是fill、stencil_2d等计划能够为field、ndarray、sparse分别提供不同 kernel 实现的原因。4.3 公共测试维度由 benchmarks/microbenchmarks/_items.py 定义各计划共享以下维度DataTypedtypei32、i64、f32、f64四种类型。remove_integer()可去掉整型math_ops 只用浮点remove([i64,f64])可限制为i32/f32matrix_ops 如此DataSizedsize按size_bytes (4**i) * 1024i 取 2、4、6、8生成16KB、256KB、4MB、64MB四档数据规模标签由size2tag()benchmarks/microbenchmarks/_utils.py格式化为16KB等Containercontainerfieldti.field与ndarrayti.ndarray两种容器fill 与 stencil_2d 还会额外追加sparse稀疏结构实现为None占位由专门函数处理MetricTypeget_metric计时指标详见 4.5 节。4.4 各计划的特有维度与 kernel 实现SAXPYsaxpy.py维度为Container × DataType × DataSize × MetricType。kernel 为z[i] 17 * x[i] y[i]元素数为dsize / dtype_size / 3三份数组。field走saxpy_field模板参数ndarray走saxpy_arrayti.types.ndarray()参数。Memcpymemcpy.py维度同上。kernel 为dst[I] src[I]ti.grouped遍历元素数为dsize / dtype_size / 2。Fillfill.py维度为Container × DataType × DataSize × MetricTypecontainer额外含sparse。稠密填充 kernel 将每个元素赋值为ti.cast(0.7, dtype)fill_sparse使用ti.root.pointer(...).dense(...)构建稀疏结构先通过ti.activate(block, [i])激活全部块再填充且 repeat 基数固定为 1。AtomicOpsatomic_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会在生成阶段剔除“逻辑原子操作 × 浮点类型”的非法组合。MathOpsmath_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用于填满指令流水线。该测试关注的是一元算子的吞吐上限而非访存带宽。MatrixOpsmatrix_ops.py维度为MatrixOps × BlockMN × ElementNum × DataType(i32/f32) × MetricType。MatrixOps提供mat_addAB、mat_mulAB、mat_mmaABC三种ti.funcBlockMN为 1×1、2×2、3×3、4×4 四种分块每个元素执行 2048 轮 × 4 次矩阵操作。Stencil2Dstencil2d.py维度为Scatter × BloclLocalStorage × Container × DataType × DataSize2D × MetricType是维度最丰富的计划。Scatter区分scattery[Ioffset] x[I]与gethergathery[I] Σ x[Ioffset]BloclLocalStorage控制是否开启块局部存储BLSdsize_2d为 128×128、512×512、2048×2048、8192×8192 四档。sparse分支通过ti.block_local(x/y)、ti.block_dim(64)实现带 BLS 的稀疏模板测试且仅对 16KB64MB 之间的规模生效否则返回None。源码注释还引用了tests/python/bls_test_template.py作为 BLS 用法的参考实现。4.5 计时指标内核耗时与端到端耗时benchmarks/microbenchmarks/_metric.py 定义了两种指标两者都先执行若干次预热compile warmup再正式计时kernel_elapsed_time_ms调用ti.init(kernel_profilerTrue, ...)初始化计时前ti.profiler.clear_kernel_profiler_info()随后用ti.profiler.get_kernel_profiler_total_time()取得纯 kernel 执行时间。它排除了启动与调度开销只统计 GPU/后端内核实际执行耗时end2end_time_ms调用ti.init(kernel_profilerFalse, ...)使用 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 合并为单个 JSONrun.py生成的是分散在多个目录、多个文件中的结果。若需要把全部结果合并为单个 JSON 文件便于脚本分析或归档使用 benchmarks/deserialize.pypython3 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的合并逻辑分两层读取顶层_info.json的suites字段按suite_name → arch → case建立索引并读取每个架构目录下的_info.json与各 case 的 JSON 文件对每个 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/visualization5006 是 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毫秒、指标类型等都足够支撑自定义绘图。七、扩展与自定义从源码出发新增测试计划这套框架的扩展点非常清晰如果你想新增一个微基准测试可以参考既有计划的写法在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()剔除不合理的组合在 benchmarks/microbenchmarks/init.py 的benchmark_plan_list中注册新计划按需在 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 后端合并结果为单 JSONpython3 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。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考