CANN SHMEM 算子测试工程实践标准目录结构、golden 生成与精度校验脚本模板全解析【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem本文以 CANN SHMEM 仓库中shmem-ops-testcase-gen技能的核心参考文档 test-structure-template.md 为主体系统讲解 SHMEM 多 PE 算子测试脚本的标准目录结构、gen_data.py/check_result.py/scripts/run.sh三件套模板与main.cpp边界规则。读完本文你可以为任意基于 SHMEM 的自定义算子搭建可复现、可验证的多进程测试工程数据按 PE 分文件生成、golden 用 float32 中间结果计算、多 PE 进程带超时看门狗并发启动、精度校验显式传入 rtol/atol 并输出 MaxAE/MaxRE/MeanRE 统计。1. 测试文件的职责划分与标准目录结构SHMEM 算子正确性验证遵循职责分离原则见 correctness.md算子main.cpp只负责单 PE 的 Host 编排、kernel 调用和输出文件写入golden 生成与精度验证全部交给独立 Python 脚本多 PE 测试由scripts/run.sh启动多个独立进程完成。这一原则保证了单进程单 PE 的语义清晰也避免了 golden 逻辑和校验逻辑污染可交付的算子工程。标准目录结构如下op_name/ ├── scripts/ │ ├── gen_data.py # 输入数据生成 golden 计算 │ ├── check_result.py # 精度验证 │ ├── run.sh # 端到端测试入口 │ └── cases.json # 可选case matrix 配置 ├── data/ # gen_data.py 生成的输入和 golden ├── output/ # 算子输出 └── src/ ├── main.cpp # Host 入口 ├── op_name_kernel.cpp └── op_name_kernel.h三个核心文件的职责边界文件职责运行时机scripts/gen_data.py固定种子生成每 PE 独立输入input_pe{rank}.bin按算子语义计算 golden保存config.json算子进程启动前算子二进制由main.cpp编译读取输入 → 拷贝到设备 → 调用 Device kernel → 输出写回output_pe{rank}.bin由 run.sh 并行启动 n_pes 个进程scripts/check_result.py逐 PE 读取输出与 golden计算 MaxAE/MaxRE/MeanRE按 rtol/atol 判定退出码 0/非 0 表示 PASS/FAIL全部进程退出后这一模式与仓内既有 examples 的落地方式一致例如 examples/tp_allreduce_udma/scripts/check_result.py 就按 rank 逐一加载output_{rank}.bin与golden.bin做比较并在首个 mismatch 处抛出带 index、actual、golden 绝对差的断言信息examples/allgather/ 则展示了数据生成scripts/data_gen.py与多 PE 进程启动run.sh分离的做法。2. gen_data.py 模板可复现的输入与 golden 生成2.1 模板代码文档给出的gen_data.py标准模板如下golden_generate中替换为实际 golden 逻辑的部分需按算子语义local compute → communication → finalize实现#!/usr/bin/env python3 import os import json import argparse import numpy as np np.random.seed(42) def gen_random_data(size, dtype): if dtype in (np.float16, np.float32, np.bfloat16): return np.random.uniform(low0.0, high10.0, sizesize).astype(dtype) elif dtype in (np.int32, np.int8): return np.random.randint(0, 1000, sizesize, dtypedtype) return None def golden_generate(args): 生成输入数据和 golden 输出 os.makedirs(args.out_dir, exist_okTrue) dtype np.dtype(args.dtype) # 1. 为每个 PE 生成独立输入 for pe_id in range(args.n_pes): input_data gen_random_data((args.M, args.N), dtypedtype) input_data.tofile(f{args.out_dir}/input_pe{pe_id}.bin) # 2. 计算 golden 输出按算子语义实现 # [替换为实际 golden 逻辑] golden np.zeros((args.M, args.N), dtypedtype) golden.tofile(f{args.out_dir}/golden.bin) # 3. 保存配置 config { n_pes: args.n_pes, M: args.M, N: args.N, dtype: args.dtype, numpy_dtype: str(dtype), expected_elems: args.M * args.N, } with open(f{args.out_dir}/config.json, w) as f: json.dump(config, f, indent2) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(--n_pes, typeint, requiredTrue) parser.add_argument(--M, typeint, requiredTrue) parser.add_argument(--N, typeint, requiredTrue) parser.add_argument(--dtype, typestr, defaultfloat16) parser.add_argument(--out_dir, typestr, default./data) args parser.parse_args() golden_generate(args)2.2 关键规范与设计意图模板配套的关键规范及原因固定np.random.seed(42)保证可复现同一份输入在多轮调试中不变mismatch 才能稳定复现定位dtype 从命令行参数接收不硬编码同一份脚本可服务 fp16/fp32/bf16/int32 多组 case配合cases.json做矩阵化测试输入/golden 按 PE 分文件文件名包含 pe_idinput_pe{rank}.bin每个 PE 输入相互独立golden 按通信语义直接构造时才能定位数据来自哪个 PE通信错位问题一目了然输出目录可配置--out_dir默认./data便于多 case 并行跑时隔离数据目录必须保存config.json包含n_pes、M、N、dtype、numpy_dtype、expected_elems是check_result.py逐 PE 校验时的唯一配置来源避免校验脚本与生成脚本之间靠默认值猜参数浮点累加优先用 float32 中间结果golden 若用 fp16 累加校验器测到的将是 golden 自身的累积误差而非算子误差。golden 构造方法上文档区分了三种策略详见 correctness.md §4Rank Pattern纯通信算子推荐为每个 PE 生成含 rank 特征的输入如pe_id * 1000000 token * 1000 hgolden 按通信语义直接构造。输出可直接看出数据来自哪个 PE适合 exact compare参考 examples/allgather/CPU/Numpy Reference通信 浮点计算用 numpy 按通信语义计算 golden浮点累加使用 float32 中间结果PyTorch Reference通算融合算子模拟完整语义 local compute → communication → finalize计算使用 float32最终按设计的 dtype/cast 规则转换。3. check_result.py 模板两种比较模式与统计输出文档提供了两种校验模式选择依据是每个 PE 是否有独立 golden。3.1 per-PE 模式推荐适用于每个 PE 有独立 golden 的集合通信/分布式算子。模板核心是check_pe函数与主流程#!/usr/bin/env python3 import sys import argparse import json import os import numpy as np def check_pe(actual, golden, pe, rtol, atol, dtype_str): abs_diff np.abs(actual.astype(np.float64) - golden.astype(np.float64)) golden_abs np.abs(golden.astype(np.float64)) eps np.finfo(np.float64).tiny rel_diff abs_diff / np.maximum(golden_abs, eps) max_ae float(np.max(abs_diff)) max_re float(np.max(rel_diff)) mean_re float(np.mean(rel_diff)) print(fPE {pe}: dtype{dtype_str} MaxAE{max_ae:.6e} MaxRE{max_re:.6e} fMeanRE{mean_re:.6e} rtol{rtol} atol{atol}) if np.any(np.isnan(actual)) or np.any(np.isinf(actual)): print(fPE {pe}: FAIL (nan/inf detected)) return False if rtol 0 and atol 0: ok np.array_equal(actual, golden) else: ok np.allclose(actual, golden, rtolrtol, atolatol) print(fPE {pe}: {PASS if ok else FAIL} ({len(actual)} elements)) if not ok: bad np.logical_and(abs_diff atol, rel_diff rtol) if rtol 0 else ~np.equal(actual, golden) idx int(np.argmax(bad)) print(f first mismatch idx{idx} actual{float(actual.astype(np.float64)[idx])} fexpected{float(golden.astype(np.float64)[idx])} diff{float(abs_diff[idx])}) return ok def main(): parser argparse.ArgumentParser() parser.add_argument(--data-dir, requiredTrue) parser.add_argument(--output-dir, requiredTrue) parser.add_argument(--rtol, typefloat, requiredTrue) parser.add_argument(--atol, typefloat, requiredTrue) args parser.parse_args() with open(os.path.join(args.data_dir, config.json)) as f: cfg json.load(f) n_pes int(cfg[n_pes]) expected_elems int(cfg[expected_elems]) dtype np.dtype(cfg.get(numpy_dtype, float16)) failed 0 for pe in range(n_pes): actual np.fromfile(os.path.join(args.output_dir, foutput_pe{pe}.bin), dtypedtype, countexpected_elems) golden np.fromfile(os.path.join(args.data_dir, fgolden_pe{pe}.bin), dtypedtype, countexpected_elems) if not check_pe(actual, golden, pe, args.rtol, args.atol, str(dtype)): failed 1 if failed: print(fFAIL: {failed}/{n_pes} PE outputs mismatched) sys.exit(1) print(fPASS: all {n_pes} PE outputs match) sys.exit(0) if __name__ __main__: main()per-PE 模式的关键细节比较前统一astype(np.float64)避免 fp16/bf16 差值在低精度下溢出或丢失有效位分母用np.maximum(golden_abs, eps)eps np.finfo(np.float64).tiny防止除零先判 nan/inf 再做容差比较——含 nan/inf 的输出直接 FAIL不进入 allclosertol 0 and atol 0时退化为np.array_equal位级精确比较这是 transport/纯搬运算子allgather、put/get 等的强制要求FAIL 时打印首个 mismatch 的 index、actual、expected 与绝对差让排查不必再跑一次退出码约定0 PASS非 0 FAIL使scripts/run.sh可以直接用exit $?传导结果。3.2 单文件模式适用于所有 PE 共享同一 golden 的场景例如 broadcast 后各 PE 输出应一致。模板结构更简单#!/usr/bin/env python3 import sys import argparse import numpy as np def check(args): golden np.fromfile(args.golden, dtypeargs.dtype) output np.fromfile(args.output, dtypeargs.dtype) if golden.shape ! output.shape: print(fFAIL: shape mismatch golden{golden.shape} output{output.shape}) return 1 if np.any(np.isnan(output)) or np.any(np.isinf(output)): print(FAIL: nan/inf detected in output) return 1 abs_diff np.abs(output.astype(np.float64) - golden.astype(np.float64)) golden_abs np.abs(golden.astype(np.float64)) eps np.finfo(np.float64).tiny max_ae np.max(abs_diff) max_re np.max(abs_diff / np.maximum(golden_abs, eps)) mean_re np.mean(abs_diff / np.maximum(golden_abs, eps)) print(fdtype: {args.dtype}) print(fMaxAE: {max_ae:.6e}) print(fMaxRE: {max_re:.6e}) print(fMeanRE: {mean_re:.6e}) print(frtol: {args.rtol}, atol: {args.atol}) if args.rtol 0 and args.atol 0: pass_flag np.array_equal(output, golden) else: pass_flag np.allclose(output, golden, rtolargs.rtol, atolargs.atol) print(PASS if pass_flag else FAIL) return 0 if pass_flag else 1 if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(--golden, typestr, requiredTrue) parser.add_argument(--output, typestr, requiredTrue) parser.add_argument(--dtype, typestr, defaultfloat16) parser.add_argument(--rtol, typefloat, requiredTrue) parser.add_argument(--atol, typefloat, requiredTrue) args parser.parse_args() sys.exit(check(args))单文件模式额外做了shape 一致性检查golden 与 output 形状必须相同这是 per-PE 模式由expected_elems隐式保证的。3.3 容差必须显式传入且按 OpTypes 体系选取模板明确要求--rtol和--atol为required 参数由scripts/run.sh显式传入不依赖脚本默认值每个 PE 必须打印 MaxAE、MaxRE、MeanRE 和 tolerance 阈值。容差值本身不是随意指定的而是按 precision-standard.md 的OpTypes × dtype × compute_times三要素选取。算子先映射到五个分类之一MOVE纯搬运如 broadcast/scatter/allgather/put-getrtol0、COMPUTE_INTEGER整型计算/索引rtol0、COMPUTE_QUANT量化、COMPUTE_FLOAT浮点规约、COMPUTE_FLOAT_HIGH_PRECISION高精度浮点如 fp32 累加的 matmulreduce。compute_times定义为PE 数 × 每 PE 参与计算/累加的元素数例如 8 PE 的 allreduce、每 PE 1024 元素 → compute_times 8192并以 2048 为界分档取值例如 fp16 的 COMPUTE_FLOAT 在 compute_times 2048 时取 2^-8、≥ 2048 时取 2^-7atol 统一等于 rtolrtol0 时即为 bitwise exact。该标准还定义了双统计判定逐元素通过率precision_percent必须为 100%且平均偏置eb |mean(actual - golden) / max(|golden|)|必须不超过按 dtype 设定的eb_thresholdfp16 为 2^-10、bf16 为 2^-7、fp32 为 2^-14对通算融合算子可另启用--torch-output的 reference pass 兜底output 对 golden 的 MARE/MERE/RMSE 与 torch 输出的比值分别在 10/2/2 门限内overall_pass basic_pass OR reference_pass。生成check_result.py时应先按design.md的meta.op_kind查映射表transport → MOVE/0/0fused_compute_comm → COMPUTE_FLOAT 且必须开 reference pass再按算子名表微调。仓内 examples/tp_allreduce_udma/scripts/check_result.py 展示了该思想的真实工程化形态整型用np.array_equal精确比较浮点转 float32 后np.allclose失败时抛出含 rank、index、output、golden、abs_diff 的AssertionError。4. scripts/run.sh 模板环境链路、多 PE 并发与超时看门狗4.1 模板代码#!/bin/bash set -euo pipefail SCRIPT_DIR$(cd $(dirname ${BASH_SOURCE[0]}) /dev/null pwd) OP_DIR$(dirname $SCRIPT_DIR) # custom-opsSHMEM 仓库根 SHMEM_REPO$(cd ${OP_DIR}/../.. pwd) EXEC_BIN${OP_DIR}/build/bin/op_name # 参数 PE_SIZE${1:-2} FIRST_NPU${3:-0} TIMEOUT${TIMEOUT:-120} # 超时上限 120 秒2 分钟 # IPPORT / SHMEM_UID_SESSION_ID 由 setup_shmem_dynamic_endpoints 分配用户 export 时尊重其值 DATA_DIR${OP_DIR}/data OUTPUT_DIR${OP_DIR}/output # 环境MUST 完整链路禁止只设 build/lib # 内联 setup_shmem_runtime_env函数体见 env-setup.snippet.md setup_shmem_runtime_env ${SHMEM_REPO} ${OP_DIR} || exit 1 if [[ ! -x ${EXEC_BIN} ]]; then echo Build first: cmake --build ${OP_DIR}/build (see shmem-ops-compile-debug/references/custom-ops-entrypoints.md §1) 2 exit 1 fi # 1. 生成测试数据 rm -rf ${DATA_DIR} ${OUTPUT_DIR} mkdir -p ${DATA_DIR} ${OUTPUT_DIR} ${PYTHON_CMD:-python3} ${SCRIPT_DIR}/gen_data.py --n_pes ${PE_SIZE} --out_dir ${DATA_DIR} # 2. 启动多 PE 进程 pids() for (( idx0; idx${PE_SIZE}; idx )); do npu_id$(( FIRST_NPU idx )) ${EXEC_BIN} ${PE_SIZE} ${idx} ${IPPORT} ${npu_id} ${DATA_DIR} ${OUTPUT_DIR} pids($!) done # 3. 超时等待上限 ${TIMEOUT} 秒 ( sleep ${TIMEOUT} echo [TIMEOUT] exceeded ${TIMEOUT}s, killing processes 2 kill ${pids[]} 2/dev/null ) watchdog$! ret0 for pid in ${pids[]}; do wait $pid || ret1 done kill $watchdog 2/dev/null wait $watchdog 2/dev/null || true if [[ $ret -ne 0 ]]; then echo [FAIL] one or more PE processes failed or timed out exit 1 fi # 4. 验证结果 ${PYTHON_CMD:-python3} ${SCRIPT_DIR}/check_result.py \ --data-dir ${DATA_DIR} \ --output-dir ${OUTPUT_DIR} \ --rtol RTOL \ --atol ATOL exit $?4.2 关键规范逐条解读模板的关键规范清单及工程原因对照仓内 env-setup.snippet.md 的环境函数实现环境必须完整链路禁止只设build/libsetup_shmem_runtime_env的顺序是——若ASCEND_HOME_PATH未设置则先 source CANN 的set_env.shCANN_SET_ENV再source ${SHMEM_REPO}/install/set_env.shSHMEM 原生安装环境然后把${SHMEM_REPO}/build/lib、可选的${OP_DIR}/build/lib、${ASCEND_HOME_PATH}/lib64依次前置到LD_LIBRARY_PATH最后调用setup_shmem_dynamic_endpoints与warn_shmem_stale_processes。只设build/lib会漏掉 SHMEM 安装环境与其他依赖库。动态端口分配禁止写死27010/8899setup_shmem_dynamic_endpoints在用户未 export 时把IPPORT随机化为tcp://127.0.0.1:(27010 RANDOM % 900)、SHMEM_UID_SESSION_ID随机化为127.0.0.1:(8899 RANDOM % 900)用户已显式 export 时尊重其值。这是为了多轮测试、多套 run 并行如与 Torch 测试同机跑时不撞 SHMEM 会话端口导致 init 失败。超时看门狗上限 120 秒模板用一个后台子进程sleep ${TIMEOUT} kill ${pids[]}实现——所有 PE 进程正常wait完成后主动kill $watchdog拆除若有 PE 卡死设备侧挂起、collective 缺人120 秒后子进程会杀掉全部 PID 并打印[TIMEOUT]脚本以 FAIL 退出。TIMEOUT环境变量可覆盖但默认不超过 2 分钟保证测试不被单个挂死进程永久阻塞。参数化与执行顺序位置参数支持PE_SIZE默认 2与FIRST_NPU默认 0流程固定为先生成数据 → 再启动进程 → 最后验证每个 PE 作为独立后台进程启动参数依次为PE_SIZE idx IPPORT npu_id DATA_DIR OUTPUT_DIR。--rtol/--atol必须显式填写不用占位符交付物中要按 §3.3 的规则填实际数值。用${PYTHON_CMD:-python3}指定 Python允许测试环境通过PYTHON_CMD注入特定解释器如 conda/venv 环境不硬编码。交付入口约束custom-ops 独立工程默认算子内仍生成custom-ops/op/scripts/run.sh但 skill/README/交付文档必须以 custom-ops-entrypoints.md §2 为首选运行入口build 失败提示必须写 §1 的编译方式禁止裸 cmake 提示。对照仓内既有脚本可以看到模板的演进脉络examples/allgather/run.sh 是最早形态——python3 ./scripts/data_gen.py生成 golden、export SHMEM_UID_SESSION_ID127.0.0.1:8899写死会话、循环后台启动各 PE 进程并wait汇总退出码而 examples/tp_allreduce_udma/scripts/run.sh 则已采用set -euo pipefail、参数解析校验is_uint等等更严格的写法。模板相当于把这些实践中的教训写死端口会冲突、无超时会挂死、环境链路不全固化成了 MUST 级规范。5. main.cpp 边界规则Host 入口允许做什么、禁止做什么main.cpp是 Host 入口模板以允许/禁止双列表划定边界允许禁止读取输入文件生成 golden 数据拷贝数据到设备精度验证 / 误差打印调用 Device kernel复杂 route/payload/tiling/packing 逻辑拷贝输出并写入文件fork/spawn 多个 PE 子进程ACL/SHMEM 初始化和清理Host RMA 作为主通信路径调用独立模块的 Host 计划函数在 main.cpp 内实现 Host 计划逻辑三条规则的动机golden 生成与精度验证禁止放进 main.cppgolden 若由被测程序自己算验证就失去了独立性误差打印混在算子进程里会让多 PE 并发输出难以归属。二者必须外置到 Python 脚本。Host RMA 不能作为主通信路径即使是 correctness-first 实现也必须使用 Device kernel可以先用单 block 串行的简单 kernel 保证正确性——被测路径必须与真实使用路径一致否则测试结论不可迁移。多 PE 启动只归 run.shmain.cpp一次启动只对应一个 PE进程编排端口、NPU 绑定、超时由外部 launcher 负责这与 SHMEM 库每个进程一个对称内存空间的模型一致。拆分信号当 Host 文件超过约 400 行且大部分内容不是 lifecycle/资源管理/launch 编排时应拆出独立.cpp/.h。逻辑归属约定如下逻辑类型归属route table / payload 编码解码 / 数据预处理Python测试数据或 独立 C Host 模块运行时layout / offset / shape 派生 / tiling 参数独立 C helper文件读写 / 资源管理可拆到 C helper多 PE 启动scripts/run.sh/ 外部 launcher6. cases.json可选的 case matrix 批量配置对于需要批量跑多组参数的场景可用cases.json描述 case matrixscripts/run.sh循环读取逐条执行[ {case_id: smoke_2pe_fp16, n_pes: 2, M: 128, N: 64, dtype: float16, rtol: 0, atol: 0}, {case_id: medium_8pe_fp16, n_pes: 8, M: 2048, N: 1024, dtype: float16, rtol: 0, atol: 0} ]每个 case 的字段与gen_data.py的命令行参数一一对应n_pes/M/N/dtypertol/atol则直接供check_result.py使用避免逐 case 手工查容差表。case matrix 的规模分档XS/S/M/L、PE 数覆盖至少 2 和 8推荐 4、八类边界条件chunk 不整除、PE 不整除、非 2 幂次维度、单行退化、UB 容量边界、signal/state 复用、2PE最大 shape、8PE最小 shape以及总 case 数 ≥ 20的最低要求定义在配套的 testcase-scale-standard.md 中——例如 L 档要求全 PE 总数据量 ≥ 256MBstress 类至少 1 个 repeat 多轮 case 用于验证 signal/state 复用与 epoch/magic 清零逻辑。完整 case matrix 还应覆盖 correctness.md §3 的 smoke、rank pattern、contract、tail/chunk、repeats、gap、visibility 七类并把未验证的 invariant 显式列出而非静默省略。7. 小结一条可复制的测试工程流水线把模板要素串起来一个 SHMEM 算子测试工程的端到端流水线是从design.md的correctness字段提取 oracle、tolerance、invariants缺阈值先补设计不在 checker 里临时硬编码gen_data.py固定种子按 PE 生成input_pe{rank}.bin与 golden落盘config.jsonscripts/run.sh调用setup_shmem_runtime_env建完整环境链路、动态分配IPPORT/SHMEM_UID_SESSION_ID然后并行拉起 n_pes 个main.cpp进程120 秒看门狗兜底全部进程退出后check_result.py逐 PE 转 float64 比较、打印 MaxAE/MaxRE/MeanRE、按rtol 0 and atol 0决定精确/容差模式以退出码 0/非 0 输出 PASS/FAIL用cases.json按 XS/S/M/L 分档批量扩展覆盖至少 20 个 case 与 8 类边界条件。该模板与 shmem-ops-testcase-gen 技能定义 的五步工作流提取 correctness contract → 生成 case matrix → gen_data.py → check_result.py → run.sh一一对应其反模式清单golden 放进 main.cpp、只有 smoke 没有中等规模、transport 算子用 relaxed tolerance、checker 不打印误差统计等可作为代码评审时的检查项直接引用。【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库基于OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考