Unified Ascend operator tests

提供声明式的 Ascend 算子功能与 Torch Event 性能测试。算子文件负责声明 cases、Golden 和自定义执行路径；需要 aclshmem、特殊输入或进程级状态时，还要声明对应配置和生命周期接口。算子文件自身不再实现 warmup、Event、同步、统计或 CSV writer。

1. 算子文件需要定义什么

导出	使用场景	是否必需
`OPERATOR_NAME`	指定终端报告中的算子名称	建议显式声明，默认使用文件名
`DTYPE`	指定框架自动生成输入 Tensor 的 dtype	可选，默认 `torch.bfloat16`
`make_cases(runtime)`	定义全部输入 shape 和csv文件 labels	必需
`golden(...)`	生成正确性基准结果	必需
`custom_op(...)`	不需要预分配资源的直接调用	与 `build_plan` 二选一
`build_plan(context)`	预分配输出、workspace、aclshmem Tensor 或 launcher	与 `custom_op` 二选一
`SHMEM`	声明 MTE/UDMA 和 aclshmem 内存池	使用 aclshmem 时必需
`input_factory(case, runtime)`	构造混合 dtype、标量、mask 或 kwargs	可选，适用于输入不能通过简单 torch.randn() 创建

Cases

统一通过 make_cases(runtime) 构造 cases：

import torch
from operator_api import OperatorCase

OPERATOR_NAME = "qkv_all2all_udma"
DTYPE = torch.bfloat16


def make_cases(runtime):
    cases = []
    for sequence_length in (2048, 4096, 8192):
        for heads_per_rank in range(1, 8):
            global_heads = heads_per_rank * runtime.world_size
            shape = (sequence_length, global_heads, 128)
            cases.append(
                OperatorCase(
                    input_shapes=(shape, shape, shape),
                    labels={
                        "S": sequence_length,
                        "H": heads_per_rank,
                        "D": 128,
                    },
                )
            )
    return cases

input_shapes 和 labels 由用户决定。

runtime 是当前 worker 的运行上下文，包含：

字段	含义
`runtime.rank`	当前进程对应的rank，范围为 `[0, world_size)`
`runtime.world_size`	本次测试使用的总卡数，由 `--num-cards` 决定
`runtime.local_rank`	当前节点上的本地 rank；首版只支持单机，因此与 `rank` 相同
`runtime.device`	当前 worker 已绑定的 NPU device
`runtime.torch`	runner 已加载的 PyTorch 模块，可用于创建输入或查询 dtype
`runtime.dist`	已初始化的 `torch.distributed` 模块

还可以调用 runtime.barrier() 执行多 rank 同步、runtime.synchronize() 等待当前 NPU 完成。

自定义输入

需要自行控制输入创建方式时，实现 input_factory(case, runtime)。例如创建一个 bfloat16 Tensor，并通过 rank 0 广播，保证所有 rank 使用相同输入：

def input_factory(case, runtime):
    tensor = runtime.torch.randn(
        case.input_shapes[0],
        dtype=runtime.torch.bfloat16,
        device=runtime.device,
    )
    runtime.dist.broadcast(tensor, src=0)
    return (tensor,)

返回值是传给 golden(...) 和 custom_op(...) 的位置参数；即使只有一个输入，也要写成 (tensor,)。需要同时传递 kwargs 时可以返回 InputBundle(args=..., kwargs=...)。

aclshmem

使用 MTE 或 UDMA 的算子需要声明 SHMEM，并在 build_plan(context) 中分配 aclshmem Tensor、输出和 workspace：

from operator_api import LaunchPlan, ShmemConfig, ShmemEngine

SHMEM = ShmemConfig(
    engine=ShmemEngine.UDMA,
    local_mem_size=SHMEM_BYTES,
)

SHMEM_BYTES 大小必须覆盖单个 case 同时存活的全部 aclshmem 分配；每个 case 结束后，框架自动释放通过 context.shmem_tensor() 创建的 Tensor，并在全部 case 结束后执行 aclshmem_finalize()。

算子返回值

算子返回值可以是 Tensor，也可以是由 tuple、list、dict 组成的嵌套 Tensor。

build_plan(context)

需要 workspace、输出 buffer、aclshmem Tensor 或提前生成 launcher 时，实现 build_plan(context)。context 是框架创建的 PrepareContext，包含当前 case 的输入、运行环境和自动释放资源：

字段	含义
`context.case`	当前 `OperatorCase`；可读取 `input_shapes` 和 `labels`
`context.inputs`	当前 case 已创建的 `InputBundle`；位置参数在 `inputs.args`，关键字参数在 `inputs.kwargs`
`context.runtime`	当前 worker 的 `RuntimeContext`；字段与前文 runtime 表格一致
`context.label(name)`	读取 `context.case.labels[name]`；不存在时直接报错
`context.empty(shape, dtype=...)`	在当前 NPU device 上创建普通 `torch.empty` Tensor
`context.shmem_tensor(shape, dtype=..., device_id=None)`	创建 aclshmem Tensor，并在当前 case 结束后自动释放；要求已声明 `SHMEM`
`context.defer(callback)`	注册自定义清理函数，在当前 case 结束时按后注册先执行的顺序调用

例如：

from operator_api import LaunchPlan


def build_plan(context):
    (input_tensor,) = context.inputs.args
    rows = int(context.label("ROWS"))

    output = context.empty(
        (rows, *input_tensor.shape[1:]),
        dtype=input_tensor.dtype,
    )
    workspace = context.shmem_tensor(
        (input_tensor.numel(),),
        dtype=input_tensor.dtype,
    )
    launcher = prepare_launcher(
        input_tensor,
        output,
        workspace,
        context.rank,
        context.world_size,
    )
    return LaunchPlan(launch=launcher, outputs=output)

LaunchPlan 不是命令行参数或模块级配置，而是 build_plan(context) 的返回值。只有使用 build_plan 的算子需要创建它；直接实现 custom_op(...) 的简单算子不需要 LaunchPlan。

字段	是否必需	含义
`launch`	必需	无参数可调用对象。框架调用它执行一次自定义算子；通常是提前准备好的 Triton launcher
`outputs`	必需	`launch()` 执行后需要与 Golden 结果比较的 Tensor，或由 tuple、list、dict 组成的嵌套 Tensor
`before_golden`	可选	每次调用 Golden 前执行的无参数函数；用于恢复输入或准备 Golden 状态，不计入 Event 测量区间
`before_custom`	可选	每次调用 `launch` 前执行的无参数函数；用于恢复输出、workspace 或输入状态，不计入 Event 测量区间
`cleanup`	可选	当前 case 结束时执行一次的无参数资源释放函数
`metadata`	可选	预留的算子附加信息；当前终端和性能 CSV 不输出该字段，通常不需要设置

最常见的返回形式只有两个必需字段：

return LaunchPlan(
    launch=launcher,
    outputs=output,
)

2. 运行

python3 perf_unified/launcher.py \
  path/to/operator.py \
  --num-cards 4 \
  --output /tmp/operator_torch_event.csv

默认同时运行正确性和 Torch Event 性能测试。只运行正确性测试时增加 --correctness-only。

以上命令默认使用 0,1,2,3。如需指定其他卡，仍可在命令前设置，例如 ASCEND_RT_VISIBLE_DEVICES=3,4,5,6；--num-cards 不能超过显式列出的卡数。 launcher 不检测卡是否被其他进程占用。

path/to/operator.py 是triton算子文件。 --output 是正确性与性能结果 CSV。

3. 执行语义

框架严格分两阶段：

依次完成所有 shape 的正确性测试；失败不会提前终止剩余 shape。框架递归调用 torch.testing.assert_close 判断和 Golden 的正确性，固定使用 rtol=1e-3、atol=1e-3。
只有全部 shape 都正确，才重新准备输入并开始完整性能测试。

任意 shape 失败时，性能 CSV 只写表头、不写数据行，进程最终返回非零退出码；具体失败 shape 和错误信息由终端输出。

所有 shape 正确时，公共计时器按两个独立阶段测试每个 shape：先执行 5 次 Golden 预热和 50 次 Golden 计时，再执行 5 次 custom/Triton 预热和 50 次 custom/Triton 计时。计时器还负责 Event 外 barrier、NPU 同步、每次迭代跨 rank MAX，以及 min/max/median/average 和 speedup。

4. CSV 输出字段

字段	含义
`OperatorCase.labels` 中的字段	按用户定义的名称展开为独立列，例如 `S`、`H`、`D`
`golden__us`、`custom__us`	两条路径的 min/max/median/average，单位为微秒
`speedup_*`	对应统计值的 `golden / custom`；大于 1 表示自定义算子更快

需要出现在性能表中的维度应由算子通过 labels 明确声明。

5. 示例

QKV 的算子示例位于 perf_unified/examples/qkv_all2all_udma.py。