TVM 测试框架指南：使用 pytest Target 参数化在多个运行时上运行单元测试

• 编译器
• 深度学习
• 模型优化

导读

TVM 是一套面向 CPU、GPU 及各类专用加速器的开源深度学习编译器栈。对于任意受支持的运行时，TVM 都应产生数值正确的结果，因此任何校验数值输出的单元测试都应覆盖所有受支持的运行时。本文围绕 docs/dev/how_to/pytest_target_parametrization.rst 展开，系统讲解tvm.testing提供的 target 参数化辅助机制：如何让一个测试函数自动展开为针对多个目标设备的多个参数化用例、如何在本地/Docker/CI 中运行这些用例，以及底层 pytest 插件与标记（mark）体系的工作原理。读完本文，你将掌握 TVM 官方推荐的单元测试编写范式，并能在自己的开发环境中复现 CI 的测试选择逻辑。

背景：为什么需要 Target 参数化

TVM 的核心承诺是"同一份代码、同一套算子语义，在不同硬件上获得正确的数值结果"。为了守住这一承诺，回归测试必须同时覆盖 llvm（CPU）、cuda、opencl、vulkan、metal、rocm、hexagon 等运行时。如果手写循环逐一执行，会带来两个问题：

1. 静默跳过：通过循环遍历目标时，被config.cmake禁用或缺少对应硬件的运行时会被静默跳过，难以在测试报告中察觉；
2. 失败定位模糊：循环式测试会在第一个失败目标处中断，无法判断错误是发生在某个特定目标上，还是所有目标共有。

因此，TVM 在tvm.testing中提供了基于 pytest 参数化（parametrization）的辅助函数：一个 Python 测试函数可以展开为多个参数化单元测试，每个用例针对单一目标设备，每个目标的通过/失败/跳过分别独立上报。

一个测试要被真正执行，必须同时满足以下全部条件（源码逻辑见 python/tvm/testing/utils.py 与 python/tvm/testing/plugin.py）：

• 该测试所在文件或目录已被传给pytest；
• 函数上应用（显式或通过 target 参数化隐式应用）的 pytest 标记，必须与pytest -m表达式的过滤要求兼容；
• 使用targetfixture 的参数化测试，其目标必须出现在环境变量TVM_TEST_TARGETS中；
• 使用targetfixture 的参数化测试，其构建配置config.cmake必须启用了对应运行时。

单元测试文件的写法

显式参数化：@tvm.testing.parametrize_targets

推荐的多目标测试方法是参数化。对于固定的一组目标，用@tvm.testing.parametrize_targets('target_1', 'target_2', ...)装饰函数，并让函数接收target或dev参数：

# 显式列出要使用的目标 @tvm.testing.parametrize_targets('llvm', 'cuda') def test_function(target, dev): # 测试代码 pass

该函数会为列出的每个目标各运行一次，且每个目标的成败独立上报。如果某个目标因在config.cmake中被禁用、或当前机器缺少相应硬件而无法运行，则该目标对应的用例会显示为skipped，而不是失败。

从源码看，parametrize_targets的实现非常轻量（python/tvm/testing/utils.py#L1287-L1330）：当带参数使用时，它等价于pytest.mark.parametrize("target", list(args), scope="session")，即以session作用域对整个目标列表做参数化。

隐式参数化：接受target/dev参数即可

对于应该在所有目标上运行的测试，装饰器可以省略。任何接受target或dev参数的测试，都会自动按照TVM_TEST_TARGETS中指定的全部目标进行参数化：

# 隐式按 TVM_TEST_TARGETS 环境变量中的全部目标参数化 def test_function(target, dev): # 测试代码 pass

这一"自动参数化"由 pytest 插件在收集阶段完成。plugin.py中的pytest_generate_tests钩子会依次调用三个处理函数（python/tvm/testing/plugin.py#L79-L89）：

• _parametrize_correlated_parameters：处理tvm.testing.parameter/tvm.testing.parameters定义的相关参数；
• _auto_parametrize_target：如果测试函数使用了targetfixture 但没有任何显式parametrize标记，则自动为其添加对target的参数化（python/tvm/testing/plugin.py#L113-L148）；
• _add_target_specific_marks：为每个目标用例补充对应的@tvm.testing.requires_*标记。

其中_auto_parametrize_target在收集时检查metafunc.fixturenames中是否包含"target"，若无显式参数化则通过utils._get_targets()取得目标列表，并为每个目标生成一个pytest.mark.parametrize标记。这也解释了为什么"只要写上target/dev参数就自动多目标"——这是插件层面的约定行为，而非parametrize_targets装饰器本身的功能。

裸装饰器形式：显式强调参数化

@tvm.testing.parametrize_targets也可以不带参数、作为裸装饰器使用，用于在代码中显式标注"本测试按全部目标参数化"。该形式本身没有额外效果（自动参数化已经生效），仅用于提高可读性，并保持向后兼容：

# 显式表明按 TVM_TEST_TARGETS 中的全部目标参数化 @tvm.testing.parametrize_targets def test_function(target, dev): # 测试代码 pass

parametrize_targets对"无参调用"的处理在 python/tvm/testing/utils.py#L1323-L1330：当唯一的参数可调用（即被装饰的函数本身）时，直接原样返回该函数。

排除与已知失败：exclude_targets/known_failing_targets

在大多数应运行于全部目标、但个别目标存在特殊情况的场景中，应使用以下两个装饰器（更详细的使用场景说明见其 docstring）：

# 从参数化目标中排除特定目标 @tvm.testing.exclude_targets("cuda") def test_function(target, dev): # 测试代码 pass # 将特定目标标记为已知失败（xfail） @tvm.testing.known_failing_targets("cuda") def test_function(target, dev): # 测试代码 pass

实现上，这两个装饰器并不直接添加 pytest 标记，而是在被装饰函数对象上记录属性（tvm_excluded_targets/tvm_known_failing_targets，见 python/tvm/testing/utils.py#L1333-L1411）：

• exclude_targets在_auto_parametrize_target中被消费：自动参数化时会把被排除的target_kind从目标列表中过滤掉；
• known_failing_targets在_add_target_specific_marks中被消费：当某个参数化目标的target_kind命中已知失败列表时，为该用例追加pytest.mark.xfail(reason='Known failing test for target "..."')（python/tvm/testing/plugin.py#L186-L194）。

多参数关联参数化

有些场景需要在多个参数上同时参数化，例如某个目标存在多个实现、需要分别测试时。可以显式地对元组参数列表做参数化，此时只有显式列出的目标会运行，但每个目标仍会被自动打上对应的@tvm.testing.requires_RUNTIME标记：

@pytest.mark.parametrize('target,impl', [ ('llvm', cpu_implementation), ('cuda', gpu_implementation_small_batch), ('cuda', gpu_implementation_large_batch), ]) def test_function(target, dev, impl): # 测试代码 pass

_add_target_specific_marks中的update_parametrize_target_arg会遍历所有显式parametrize标记：只要参数名列表包含"target"，就会把该参数化标记中与 target 对应的值替换为pytest.param(*values, marks=requires_*)形式，从而为每个目标用例挂上正确的运行时要求标记（python/tvm/testing/plugin.py#L151-L221）。这里还隐含一个易错点：参数值必须是以列表形式给出的参数集；如果误用元组，plugin.py会抛出带文件名与行号的TypeError提示。

pytest 标记体系

参数化功能构建在 pytest marks 之上。每个测试函数都可以用 pytest 标记附加元数据，其中最常用的标记如下：

使用参数化 target 时，每个测试用例都会自动挂上与其目标对应的@tvm.testing.requires_RUNTIME标记。因此，如果目标在config.cmake中被禁用或缺少硬件，该用例会被明确列为 skipped。

这套标记体系的底层是utils.py中的Feature类（python/tvm/testing/utils.py#L529-L821）。每个特性（feature）可以声明以下检查维度，由Feature.marks()组合成实际生效的 pytest 标记序列：

• cmake_flag：构建时必须在config.cmake中启用的 CMake 开关（如USE_CUDA、USE_LLVM），通过tvm.support.libinfo()读取构建信息，禁用则追加skipif；
• target_kind_enabled：目标种类必须出现在TVM_TEST_TARGETS（或默认目标列表）中；
• target_kind_hardware：必须存在对应设备（检查tvm.device(kind).exist）；
• compile_time_check/run_time_check：编译期（如 nvcc 版本）与运行期（如 GPU 算力）自定义检查，返回字符串时直接作为跳过原因；
• parent_features：特性继承（如 cuDNN 依赖 CUDA），被依赖特性（除target_kind_enabled外）的检查会一并继承。

标记支持三种support_required模式："compile-and-run"（默认，编译与运行条件都要满足）、"compile-only"（仅编译期检查）、"optional"（只打标记、不跳过，用于兼容旧式enabled_targets()风格）。仓库中已预置大量特性，例如requires_llvm、requires_cuda、requires_rocm、requires_vulkan、requires_metal、requires_opencl、requires_hexagon、requires_cudnn、requires_cublas、requires_nvptx、requires_micro等（python/tvm/testing/utils.py#L844-L1045）。plugin.py的pytest_configure会在收集前把这些特性的标记名全部注册到 pytest（python/tvm/testing/plugin.py#L64-L72）。

旧式写法（不推荐）：enabled_targets()

tvm.testing.enabled_targets()会返回当前机器上"已启用且可运行"的全部目标（以(target, device)对的形式），其判定综合了TVM_TEST_TARGETS环境变量、构建配置与物理硬件。现有测试大多仍显式循环该函数返回值，但新测试不应再使用这种风格，原因有二：

1. 循环式测试会静默跳过config.cmake中禁用或无设备可运行的运行时，报告中看不出跳过；
2. 循环会在第一个失败目标处中断，无法判断错误是特定目标独有还是所有目标共有。

# 旧式写法，请勿在新测试中使用 def test_function(): for target, dev in tvm.testing.enabled_targets(): # 测试代码 pass

从实现看，enabled_targets()基于_get_targets()（python/tvm/testing/utils.py#L405-L451）：对每个候选目标计算is_enabled（tvm.runtime.enabled(kind)或libinfo()中的开关）与is_runnable（is_enabled and tvm.device(kind).exist），仅返回is_runnable的目标。若全部目标都不可运行，会回退到仅运行llvm；若连llvm也未启用则抛出TVMError。

本地运行

在${TVM_HOME}目录下直接执行pytest即可运行 Python 单元测试。

环境变量

TVM_TEST_TARGETS：以分号分隔的目标列表，决定参数化测试覆盖哪些目标。未设置时，默认使用tvm.testing.DEFAULT_TEST_TARGETS。

仓库中的默认目标列表（python/tvm/testing/utils.py#L454-L465）：

DEFAULT_TEST_TARGETS = [ "llvm", "cuda", "nvptx", "vulkan -from_device=0", "opencl", "opencl -device=mali,aocl_sw_emu", "opencl -device=intel_graphics", "metal", "rocm", "hexagon", ]

注意两点：

• 目标串可以携带设备相关参数（如opencl -device=intel_graphics、vulkan -from_device=0），目标种类为第一个空格前的部分；
• 如果TVM_TEST_TARGETS中没有任何"既已启用、又有可用设备"的目标，测试会自动回退到仅运行llvm目标。该回退逻辑与对应的警告日志位于_get_targets()中；若连llvm也未启用，则会抛出TVMError（提示尝试设置TVM_TEST_TARGETS为受支持目标）。_tvm_test_targets()（python/tvm/testing/utils.py#L1117-L1124）在解析时会去重并保持用户指定的顺序。

TVM_LIBRARY_PATH：指向libtvm.so库文件的路径。例如可用它指向 debug 构建的库来运行测试。未设置时，会相对于 TVM 源码目录搜索libtvm.so。

命令行参数

• 传入文件或目录路径：只运行该文件/目录中的单元测试。例如在未安装特定前端（frontend）的机器上，可以通过只传tests/python/frontend之外的路径来避免执行该目录下的测试。
• -m参数：仅运行带指定 pytest 标记的测试。最常见的是-m gpu（只运行标记了@pytest.mark.gpu、需要 GPU 的测试）或-m 'not gpu'（只运行不使用 GPU 的测试）。

一个容易踩坑的细节：-m过滤发生在基于TVM_TEST_TARGETS的目标选择之后。即使指定了-m gpu，如果TVM_TEST_TARGETS中没有包含 GPU 目标，也不会运行任何 GPU 测试。这一过滤顺序在 python/tvm/testing/plugin.py 的pytest_sessionfinish中也有体现：当-m表达式过滤后没有收集到任何测试时，pytest 不会以"未收集到测试"的错误状态退出。

在本地 Docker 容器中运行

docker/bash.sh脚本（仓库根目录下的 docker/bash.sh）可以在与 CI 相同的 docker 镜像中运行单元测试：

# 第一个参数指定要使用的镜像，如 ci_gpu docker/bash.sh ci_gpu

允许的镜像名定义在 TVM 源码目录中 Jenkinsfile 的顶部，并映射到 tlcpack 组织发布的 docker 镜像（相关镜像构建文件见 docker/Dockerfile.ci_cpu、docker/Dockerfile.ci_gpu 等）。

• 如果不传额外参数，会进入镜像内的交互式 bash 会话；
• 如果传入脚本作为可选参数，则该脚本会在镜像内执行，例如：

docker/bash.sh ci_gpu tests/scripts/task_python_unittest.sh

需要特别注意的是：docker 镜像包含全部系统依赖，但并未内置build/config.cmake配置文件。TVM 源码目录被挂载为镜像的家目录，因此默认会使用本地已有的 config/build 目录。一个常见做法是分别维护build_local和build_docker两个目录，进入/退出 docker 时通过符号链接将build指向对应目录。

在 CI 中运行

CI 中所有环节都从 Jenkinsfile 中的任务定义出发，包括：使用哪个 docker 镜像、编译期配置是什么、哪些测试进入哪些阶段。相关定义位于 ci/jenkins/generated（生成的 Jenkinsfile，如cpu_jenkinsfile.groovy、gpu_jenkinsfile.groovy）与 ci/jenkins/templates（生成模板，如 ci/jenkins/templates/cpu_jenkinsfile.groovy.j2）。

Docker 镜像

Jenkinsfile 中的每个任务（如BUILD: CPU）都会调用docker/bash.sh。与本地用法一致，docker/bash.sh后面的参数定义了 CI 中使用的镜像。例如 ci/jenkins/templates/utils/Build.groovy.j2 中，Python 单测任务会以指定镜像执行./tests/scripts/task_python_unittest.sh。

编译期配置

docker 镜像中没有内置config.cmake，因此它是每个BUILD任务的第一步工作，通过tests/scripts/task_config_build_*.sh脚本完成（例如task_config_build_gpu.sh、task_config_build_cpu.sh等，见 tests/scripts 目录）。具体使用哪个脚本取决于被测试的构建类型，由 Jenkinsfile 指定。每个BUILD任务最后会把编译好的库打包，供后续测试阶段使用。

测试运行

Jenkinsfile 的Unit Test与Integration Test阶段决定pytest的调用方式。每个任务先解包BUILD阶段编译好的库，然后运行测试脚本（如tests/scripts/task_python_unittest.sh）。这些脚本设置了传给pytest的文件/目录与命令行选项。

CI 脚本对TVM_TEST_TARGETS与-m gpu的组合使用是理解整个体系的最佳范例（见 tests/scripts 目录）：

• tests/scripts/task_python_unittest_gpuonly.sh 中：export TVM_TEST_TARGETS="cuda;opencl;metal;rocm;nvptx;opencl -device=mali,aocl_sw_emu"并设置PYTEST_ADDOPTS="-m gpu"，只运行 GPU 标记的测试；后续还用TVM_TEST_TARGETS="vulkan -from_device=0"单独覆盖 Vulkan 设备；
• tests/scripts/task_python_integration_gpuonly.sh 中：TVM_TEST_TARGETS="cuda;opencl;metal;rocm;nvptx;opencl -device=mali,aocl_sw_emu,adreno"，配合-m gpu；
• tests/scripts/task_python_frontend.sh 使用TVM_TEST_TARGETS="llvm;cuda"，而纯 CPU 的前端任务 tests/scripts/task_python_frontend_cpu.sh 使用TVM_TEST_TARGETS="llvm"；
• tests/scripts/task_python_adreno.sh 针对 Adreno GPU 将TVM_TEST_TARGETS设为"opencl"。

多个 CI 脚本都包含-m gpu选项，用于将测试范围限制为带@pytest.mark.gpu标记的用例。

其他测试工具速览

tvm.testing包（python/tvm/testing/init.py）除 target 参数化外，还提供一组与 pytest 深度集成的辅助设施，写作单元测试时经常会用到：

• tvm.testing.parameter/tvm.testing.parameters：定义参数化 fixture。前者是多个参数全组合；后者是多个参数按组对齐（每组值只运行一次）。二者均以session作用域运行，适合无设置成本的参数（字符串、整数、元组等）；
• tvm.testing.fixture：定义带设置成本的 fixture，可通过cache_return_value=True在测试间缓存返回值（可用环境变量TVM_TEST_DISABLE_CACHE强制关闭缓存）；
• tvm.testing.assert_allclose：带默认atol/rtol的数值比较工具；
• tvm.testing.CompareBeforeAfter：编写 TIR 变换测试的基类，通过定义before/transform/expected成员即可自动生成测试；
• pytest_plugins = ['tvm.testing.plugin']：在conftest.py中声明该行，即可在 TVM 测试目录之外复用这套参数化与标记体系（python/tvm/testing/plugin.py#L19-L33）。

总结

TVM 的 target 参数化机制把"多运行时覆盖"从手动循环变成了 pytest 原生参数化：一个测试函数只需声明target/dev参数，pytest 插件就会自动按TVM_TEST_TARGETS展开用例，并为每个用例挂上对应的requires_*标记，使禁用或缺失硬件的目标被显式跳过。撰写新测试时，优先采用target/dev参数配合exclude_targets、known_failing_targets的表达方式，避免使用enabled_targets()循环的旧式写法；运行测试时，则通过TVM_TEST_TARGETS、-m与config.cmake三者的配合精确控制覆盖范围——这正是 TVM 官方 CI 在 CPU、GPU 与各类专用硬件节点上复用同一套测试套件的核心机制。

• 编译器
• 深度学习
• 模型优化