PyTorch CPU 量化内核目录解析：AVX/AVX2/AVX-512 多实例编译与运行时调度机制

本篇基于 kernels/README.md 展开，讲清 PyTorch 中 aten/src/ATen/native/quantized/cpu/kernels/ 这个目录存在的目的——同一份源码按 AVX、AVX2、AVX-512 等不同 CPU 向量指令集多次编译、再由运行时调度机制选出最优实现，并结合 QuantizedOpKernels.cpp、DispatchStub.h 与 cmake/Codegen.cmake 的源码，完整呈现"为什么、怎么实现、改代码时必须遵守哪些规则"。读完后你能够理解 PyTorch CPU 端性能优化的"多实例编译 + 调度桩"模式，并在修改该目录代码时避免 ODR 违规与调度错误。

1. 目录定位：为处理器向量能力生成最优代码

README 开篇即说明该目录的核心设计意图：

The files in this directory are compiled multiple times for different CPU vector instruction sets (e.g. AVX, AVX2). The purpose of putting code in this directory is to make sure we can generate the optimal code for a given processor's vector capabilities.

即：目录下的 C++ 文件会被编译器按不同 CPU 向量指令集（SSE、AVX、AVX2、AVX-512 等）重复编译多次，每次编译时编译器只针对当前指令集做向量化优化；最终二进制中同时包含多套实现，运行时再根据实际 CPU 能力选择最合适的一套。这样既能让不支持 AVX-512 的旧 CPU 安全运行，又能让高端 CPU 用上最宽的向量宽度。

README 还指出，很多工作是通过 vec256.h（README 写作 vec256_qint.h，即 256 位量化向量头文件，现已归入 aten/src/ATen/cpu/vec/vec256/ 向量库目录）中的预处理器守护完成的。从当前源码看，这一说法仍然成立：QuantizedOpKernels.cpp 中大量使用 #ifdef CPU_CAPABILITY_AVX2 / #elif defined(CPU_CAPABILITY_AVX512) 来切换 256 位/512 位内联汇编路径。

1.1 构建系统如何"编译多次"

cmake/Codegen.cmake 明确标注了处理"需要多次编译"的源文件列表：

# Handle source files that need to be compiled multiple times for
# different vectorization options
file(GLOB cpu_kernel_cpp_in
  "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cpu/*.cpp"
  "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/quantized/cpu/kernels/*.cpp")

构建脚本随后维护 CPU_CAPABILITY_NAMES（如 DEFAULT、SSE_4_2、AVX2、AVX512 等）与对应的编译标志列表：例如检测到 CXX_AVX512_FOUND 且未通过 USE_CPU_VECTORIZATION=0 关闭时，会定义 HAVE_AVX512_CPU_DEFINITION。之后 process_vec 函数对每个能力变体生成一份独立的目标文件并用该变体的 -mavx2、-mavx512f 等标志编译——这就是"同一 .cpp 编出多个 .o"的落地位置。

针对 QuantizedOpKernels.cpp 在 DEFAULT 变体下还有一条针对性豁免（见 cmake/Codegen.cmake 中 if("${NAME}" STREQUAL "native/quantized/cpu/kernels/QuantizedOpKernels.cpp") 分支），为其追加 -Wno-deprecated-copy 以绕过已知编译器告警。Bazel 构建体系下，build_variables.bzl 也把该文件单独列入了向量多编译源列表，保证两种构建系统行为一致。

2. 调度三件套：DECLARE_DISPATCH / DEFINE_DISPATCH / REGISTER_DISPATCH

README 的第二条规则是：本目录中所有代码都必须通过 DECLARE_DISPATCH、DEFINE_DISPATCH、REGISTER_DISPATCH 这套机制注册，以确保运行时调度到正确的实现。这套机制定义在 aten/src/ATen/native/DispatchStub.h：

• DECLARE_DISPATCH(fn, name)：在头文件中声明一个 DispatchStub<FnPtr, ...> 派生结构体（如 qrelu_stub），其内部持有函数指针表；禁止拷贝/移动，作为进程内单例式的"调度桩"存在；
• DEFINE_DISPATCH(name)：在实现文件中定义该结构体；
• REGISTER_ARCH_DISPATCH(name, arch, fn) / REGISTER_DISPATCH(name, fn)：把某个架构变体下的具体内核函数挂到调度桩的对应槽位上。

在 CPU 路径下（非 CUDA/HIP/MPS），REGISTER_DISPATCH 的行为有一个关键细节（DispatchStub.h 中 #elif defined(CPU_CAPABILITY) 分支）：

#ifdef CPU_CAPABILITY_AVX512
// REGISTER_DISPATCH now dispatches an AVX512 kernel to nullptr
#define REGISTER_DISPATCH(name, fn) \
  REGISTER_ARCH_DISPATCH(name, CPU_CAPABILITY, ((void*)(fn) ? nullptr : nullptr))
#else
#define REGISTER_DISPATCH(name, fn) REGISTER_ARCH_DISPATCH(name, CPU_CAPABILITY, fn)
#endif
#define ALSO_REGISTER_AVX512_DISPATCH(name, fn) REGISTER_ARCH_DISPATCH(name, CPU_CAPABILITY, fn)

含义是：当文件正以 AVX-512 能力编译时，普通 REGISTER_DISPATCH 故意把 AVX-512 槽位注册为 nullptr（回退到更低能力实现），只有显式使用 ALSO_REGISTER_AVX512_DISPATCH 的内核才会把 AVX-512 版本真正挂上去。这样设计允许开发者按性能实测逐个内核决定是否值得启用 512 位路径。

QuantizedOpKernels.cpp 末尾正是这种模式的典型样本：qrelu_stub、qgelu_stub、qadd_stub、qbatch_norm_stub、qmul_tensor_cpu_stub 等几十个量化内核在此集中注册；其中 qmul_tensor_cpu_stub、qadd_tensor_cpu_stub、qbatch_norm_cpu_stub 使用了 ALSO_REGISTER_AVX512_DISPATCH 显式启用 512 位路径，而源码注释明确说明另有一批内核"dispatched to AVX2 because they don't perform as well with AVX512"，体现了"逐内核按性能决定变体"的工程决策。

3. 必须放入匿名命名空间：防 ODR 违规

README 用大写强调的第三条规则：

THE CODE MUST RESIDE IN THE ANONYMOUS NAMESPACE. FAILURE TO ENSURE THIS IS THE CASE CAN LEAD TO HARD-TO-DEBUG ODR VIOLATIONS.

原因可以直接从构建机制推导：同一份 .cpp 会被编译进多个目标文件（每个 CPU 能力一份）。如果其中函数带有外部链接（external linkage），这些符号会在链接阶段重复出现——虽然链接器通常能容忍，但一旦符号被跨翻译单元引用、或不同编译变体间行为出现细微差异，就会产生难以排查的 ODR（One Definition Rule）违规与运行时不可预期行为。将其全部收入匿名命名空间后，每个编译变体的符号都是内部链接，彼此隔离。

实际代码严格执行了这一点：QuantizedOpKernels.cpp 从第 41 行 namespace at::native { 开始即进入第 42 行的 namespace {（匿名命名空间），直到第 4622 行 } // anonymous namespace 才闭合；整个文件还以 // NOLINTBEGIN(*-c-arrays) 开头处理其中的裸数组用法。文件内甚至保留了"醒目警告"：

// ****************** HEY YOU! YES YOU! Read this! ********************
//
// Please read the README.md in this directory before editing this file

4. 预处理器守护：AVX2 与 AVX512 的内联路径切换

README 提到"Much of this is done via preprocessor guards"。以文件内的水平求和辅助函数 hsum/hsum_sq（分别对 uint8_t、int8_t、int32_t 做向量化求和与平方和）为例，可以看到守护的典型写法：

• #ifdef CPU_CAPABILITY_AVX2：使用 _mm256_* 指令族，按 32 字节步进累加（如 _mm256_maddubs_epi16 做无符号 8 位到 16 位的乘加）；
• #elif defined(CPU_CAPABILITY_AVX512)：使用 _mm512_* 指令族，按 64 字节步进累加；
• 两个宏都未定义时：自动落入纯标量回退循环。

hsum_sq 中还有值得注意的工程细节：由于 8 位数据的平方累加存在 int32 溢出风险，代码设置了 overflow_threshold（AVX2 下为 262144，即 2147483647 / (512*512) * 8 的推导结果），分块处理并周期性把向量累加器落回 int64_t 结果，防止大长度输入产生错误结果。这类"同一函数、按指令集生成不同最优实现、再配标量尾处理"的三段式结构，正是该目录代码的普遍形态。

5. 测试要求：每个指令集变体都必须被覆盖

README 最后一条加粗规则是：确保 AVX、AVX2 等不同变体都被测试。CI 中存在刻意禁用 AVX/AVX2 的构建变体（例如 USE_CPU_VECTORIZATION=0），用来验证无向量扩展路径也能工作。

源码里也留有变体差异会暴露真实问题的记录：文件末尾的注释说明部分量化测试在 Windows + AVX-512 组合下出现不稳定失败（引用 GH 56992），因此在 _WIN32 下这些内核改用 REGISTER_DISPATCH 走 AVX2 路径而非 ALSO_REGISTER_AVX512_DISPATCH。这段注释是"多变体必须分别验证"这一规则的现实注脚——不同指令集变体不仅影响性能，还可能暴露数值与稳定性问题。

6. 修改该目录前的检查清单

综合 README 的四条考量与源码证据，改动 aten/src/ATen/native/quantized/cpu/kernels/ 下代码前应确认：

1. 必要性：只把确实受益于"按指令集多编译"的代码放在这里，代码量保持最小（每份代码都会被编译多次，编译时间成倍增加）；
2. 调度注册：新增内核必须走 DECLARE_DISPATCH/DEFINE_DISPATCH/REGISTER_DISPATCH 流程；若需要 AVX-512 加速，显式使用 ALSO_REGISTER_AVX512_DISPATCH 并基于实测性能决策；
3. 命名空间：全部实现置于匿名命名空间内，避免跨编译变体的 ODR 违规；
4. 守护与回退：向量路径用 CPU_CAPABILITY_* 宏守护，并保证标量回退路径正确；
5. 变体测试：在含"无 AVX/无 AVX2"构建变体的 CI 上验证，确认所有路径均可运行。

该目录当前唯一实现文件即 QuantizedOpKernels.cpp，覆盖 qrelu、qgelu、qsigmoid、qhardsigmoid、qclamp、qadd/qmul（含 ReLU 融合模板变体）、qbatch_norm、qmaxpool（NHWC/NTHWC）、量化/反量化（per-tensor 与 per-channel）、masked_fill/index_put 等 CPU 量化内核；上游 CPU 算子层如 cpu/qrelu.cpp、cpu/qgelu.cpp、cpu/qthreshold.cpp 等通过引用这些调度桩（*_stub）完成跨指令集的统一入口，读者可沿此调用链继续深入。