# 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：

```python
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 使用相同输入：

```python
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：

```python
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 结束时按后注册先执行的顺序调用 |

例如：

```python
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 不输出该字段，通常不需要设置 |

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

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

## 2. 运行

```bash
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. 执行语义

框架严格分两阶段：

1. 依次完成所有 shape 的正确性测试；失败不会提前终止剩余 shape。框架递归调用 `torch.testing.assert_close` 判断和 Golden 的正确性，固定使用
`rtol=1e-3`、`atol=1e-3`。
2. 只有全部 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`。