企业官网建设流程全解析

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

三个核心文件的职责边界：

这一模式与仓内既有 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(low=0.0, high=10.0, size=size).astype(dtype) elif dtype in (np.int32, np.int8): return np.random.randint(0, 1000, size=size, dtype=dtype) return None def golden_generate(args): """生成输入数据和 golden 输出""" os.makedirs(args.out_dir, exist_ok=True) dtype = np.dtype(args.dtype) # 1. 为每个 PE 生成独立输入 for pe_id in range(args.n_pes): input_data = gen_random_data((args.M, args.N), dtype=dtype) input_data.tofile(f"{args.out_dir}/input_pe{pe_id}.bin") # 2. 计算 golden 输出（按算子语义实现） # [替换为实际 golden 逻辑] golden = np.zeros((args.M, args.N), dtype=dtype) 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, indent=2) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--n_pes', type=int, required=True) parser.add_argument('--M', type=int, required=True) parser.add_argument('--N', type=int, required=True) parser.add_argument('--dtype', type=str, default='float16') parser.add_argument('--out_dir', type=str, 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_id（input_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 §4）：

1. Rank Pattern（纯通信算子推荐）：为每个 PE 生成含 rank 特征的输入（如pe_id * 1000000 + token * 1000 + h），golden 按通信语义直接构造。输出可直接看出数据来自哪个 PE，适合 exact compare，参考 examples/allgather/；
2. CPU/Numpy Reference（通信 + 浮点计算）：用 numpy 按通信语义计算 golden，浮点累加使用 float32 中间结果；
3. 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(f"PE {pe}: dtype={dtype_str} MaxAE={max_ae:.6e} MaxRE={max_re:.6e} " f"MeanRE={mean_re:.6e} rtol={rtol} atol={atol}") if np.any(np.isnan(actual)) or np.any(np.isinf(actual)): print(f"PE {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, rtol=rtol, atol=atol) print(f"PE {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])} " f"expected={float(golden.astype(np.float64)[idx])} diff={float(abs_diff[idx])}") return ok def main(): parser = argparse.ArgumentParser() parser.add_argument("--data-dir", required=True) parser.add_argument("--output-dir", required=True) parser.add_argument("--rtol", type=float, required=True) parser.add_argument("--atol", type=float, required=True) 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, f"output_pe{pe}.bin"), dtype=dtype, count=expected_elems) golden = np.fromfile(os.path.join(args.data_dir, f"golden_pe{pe}.bin"), dtype=dtype, count=expected_elems) if not check_pe(actual, golden, pe, args.rtol, args.atol, str(dtype)): failed += 1 if failed: print(f"FAIL: {failed}/{n_pes} PE outputs mismatched") sys.exit(1) print(f"PASS: 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，不进入 allclose；
• rtol == 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, dtype=args.dtype) output = np.fromfile(args.output, dtype=args.dtype) if golden.shape != output.shape: print(f"FAIL: 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(f"dtype: {args.dtype}") print(f"MaxAE: {max_ae:.6e}") print(f"MaxRE: {max_re:.6e}") print(f"MeanRE: {mean_re:.6e}") print(f"rtol: {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, rtol=args.rtol, atol=args.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', type=str, required=True) parser.add_argument('--output', type=str, required=True) parser.add_argument('--dtype', type=str, default='float16') parser.add_argument('--rtol', type=float, required=True) parser.add_argument('--atol', type=float, required=True) 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-get，rtol=0）、COMPUTE_INTEGER（整型计算/索引，rtol=0）、COMPUTE_QUANT（量化）、COMPUTE_FLOAT（浮点规约）、COMPUTE_FLOAT_HIGH_PRECISION（高精度浮点，如 fp32 累加的 matmul+reduce）。compute_times定义为"PE 数 × 每 PE 参与计算/累加的元素数"（例如 8 PE 的 allreduce、每 PE 1024 元素 → compute_times = 8192），并以 2048 为界分档取值，例如 fp16 的 COMPUTE_FLOAT 在 compute_times < 2048 时取 2^-8、≥ 2048 时取 2^-7；atol 统一等于 rtol，rtol=0 时即为 bitwise exact。

该标准还定义了双统计判定：逐元素通过率precision_percent必须为 100%，且平均偏置eb = |mean(actual - golden) / max(|golden|)|必须不超过按 dtype 设定的eb_threshold（fp16 为 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/0；fused_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-ops：SHMEM 仓库根 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 (( idx=0; 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=$! ret=0 for pid in "${pids[@]}"; do wait $pid || ret=1 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 的环境函数实现）：

1. 环境必须完整链路，禁止只设build/lib：setup_shmem_runtime_env的顺序是——若ASCEND_HOME_PATH未设置则先 source CANN 的set_env.sh（CANN_SET_ENV）；再source ${SHMEM_REPO}/install/set_env.sh（SHMEM 原生安装环境）；然后把${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 安装环境与其他依赖库。

2. 动态端口分配，禁止写死27010/8899：setup_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 失败。

3. 超时看门狗上限 120 秒：模板用一个后台子进程sleep ${TIMEOUT} && kill "${pids[@]}"实现——所有 PE 进程正常wait完成后主动kill $watchdog拆除；若有 PE 卡死（设备侧挂起、collective 缺人），120 秒后子进程会杀掉全部 PID 并打印[TIMEOUT]，脚本以 FAIL 退出。TIMEOUT环境变量可覆盖但默认不超过 2 分钟，保证测试不被单个挂死进程永久阻塞。

4. 参数化与执行顺序：位置参数支持PE_SIZE（默认 2）与FIRST_NPU（默认 0）；流程固定为"先生成数据 → 再启动进程 → 最后验证"；每个 PE 作为独立后台进程启动，参数依次为PE_SIZE idx IPPORT npu_id DATA_DIR OUTPUT_DIR。

5. --rtol/--atol必须显式填写，不用占位符：交付物中要按 §3.3 的规则填实际数值。

6. 用${PYTHON_CMD:-python3}指定 Python：允许测试环境通过PYTHON_CMD注入特定解释器（如 conda/venv 环境），不硬编码。

7. 交付入口约束（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_ID=127.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 生成与精度验证禁止放进 main.cpp：golden 若由被测程序自己算，验证就失去了独立性；误差打印混在算子进程里会让多 PE 并发输出难以归属。二者必须外置到 Python 脚本。
• Host RMA 不能作为主通信路径：即使是 correctness-first 实现，也必须使用 Device kernel（可以先用单 block 串行的简单 kernel 保证正确性）——被测路径必须与真实使用路径一致，否则测试结论不可迁移。
• 多 PE 启动只归 run.sh：main.cpp一次启动只对应一个 PE，进程编排（端口、NPU 绑定、超时）由外部 launcher 负责，这与 SHMEM 库"每个进程一个对称内存空间"的模型一致。

拆分信号：当 Host 文件超过约 400 行，且大部分内容不是 lifecycle/资源管理/launch 编排时，应拆出独立.cpp/.h。逻辑归属约定如下：

6. cases.json：可选的 case matrix 批量配置

对于需要批量跑多组参数的场景，可用cases.json描述 case matrix，scripts/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/dtype），rtol/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 总数据量 ≥ 256MB，stress 类至少 1 个 repeat 多轮 case 用于验证 signal/state 复用与 epoch/magic 清零逻辑。完整 case matrix 还应覆盖 correctness.md §3 的 smoke、rank pattern、contract、tail/chunk、repeats、gap、visibility 七类，并把未验证的 invariant 显式列出而非静默省略。

7. 小结：一条可复制的测试工程流水线

把模板要素串起来，一个 SHMEM 算子测试工程的端到端流水线是：

1. 从design.md的correctness字段提取 oracle、tolerance、invariants（缺阈值先补设计，不在 checker 里临时硬编码）；
2. gen_data.py固定种子按 PE 生成input_pe{rank}.bin与 golden，落盘config.json；
3. scripts/run.sh调用setup_shmem_runtime_env建完整环境链路、动态分配IPPORT/SHMEM_UID_SESSION_ID，然后并行拉起 n_pes 个main.cpp进程，120 秒看门狗兜底；
4. 全部进程退出后，check_result.py逐 PE 转 float64 比较、打印 MaxAE/MaxRE/MeanRE、按rtol == 0 and atol == 0决定精确/容差模式，以退出码 0/非 0 输出 PASS/FAIL；
5. 用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