# Unified Ascend operator tests

Provides declarative Ascend operator functionality and Torch Event performance testing. Operator files are responsible for declaring cases, Golden reference, and custom execution paths; when aclshmem, special inputs, or process-level state are needed, corresponding configuration and lifecycle interfaces must also be declared. The operator file itself no longer implements warmup, Event, synchronization, statistics, or CSV writer.

## 1. What the operator file needs to define

| Export | Use Case | Required |
| --- | --- | --- |
| `OPERATOR_NAME` | Specifies the operator name in terminal reports | Recommended to explicitly declare, defaults to filename |
| `DTYPE` | Specifies the dtype for framework-generated input Tensors | Optional, defaults to `torch.bfloat16` |
| `make_cases(runtime)` | Defines all input shapes and csv file labels | Required |
| `golden(...)` | Generates correctness baseline results | Required |
| `custom_op(...)` | Direct invocation without pre-allocated resources | Choose one between this and `build_plan` |
| `build_plan(context)` | Pre-allocate output, workspace, aclshmem Tensor or launcher | Choose one between this and `custom_op` |
| `SHMEM` | Declares MTE/UDMA and aclshmem memory pool | Required when using aclshmem |
| `input_factory(case, runtime)` | Constructs mixed dtype, scalars, masks, or kwargs | Optional, suitable when inputs cannot be created through simple torch.randn() |

### Cases

Uniformly construct cases through `make_cases(runtime)`:

```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` and labels are decided by the user.

`runtime` is the current worker's runtime context, containing:

| Field | Meaning |
| --- | --- |
| `runtime.rank` | Current process's rank, range `[0, world_size)` |
| `runtime.world_size` | Total number of cards used for this test, determined by `--num-cards` |
| `runtime.local_rank` | Local rank on current node; first version only supports single machine, so same as `rank` |
| `runtime.device` | NPU device bound to current worker |
| `runtime.torch` | PyTorch module loaded by runner, can be used to create inputs or query dtype |
| `runtime.dist` | Initialized `torch.distributed` module |

You can also call `runtime.barrier()` to perform multi-rank synchronization, `runtime.synchronize()` to wait for current NPU to complete.

### Custom inputs

When you need to control input creation yourself, implement `input_factory(case, runtime)`. For example, create a bfloat16 Tensor and broadcast from rank 0 to ensure all ranks use the same input:

```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,)
```

The return value is the positional arguments passed to `golden(...)` and `custom_op(...)`; even with only one input, it must be written as `(tensor,)`. When you need to pass kwargs simultaneously, you can return `InputBundle(args=..., kwargs=...)`.

### aclshmem

Operators using MTE or UDMA need to declare `SHMEM` and allocate aclshmem Tensor, output, and workspace in `build_plan(context)`:

```python
from operator_api import LaunchPlan, ShmemConfig, ShmemEngine

SHMEM = ShmemConfig(
    engine=ShmemEngine.UDMA,
    local_mem_size=SHMEM_BYTES,
)
```
`SHMEM_BYTES` size must cover all aclshmem allocations alive for a single case simultaneously; after each case completes, the framework automatically releases Tensors created through `context.shmem_tensor()` and executes `aclshmem_finalize()` after all cases complete.

### Operator return values

Operator return values can be a Tensor, or nested Tensors composed of tuple, list, dict.

### build_plan(context)

When workspace, output buffer, aclshmem Tensor, or pre-generated launcher is needed, implement `build_plan(context)`. `context` is the `PrepareContext` created by the framework, containing the current case's inputs, runtime environment, and auto-released resources:

| Field | Meaning |
| --- | --- |
| `context.case` | Current `OperatorCase`; can read `input_shapes` and `labels` |
| `context.inputs` | Current case's created `InputBundle`; positional arguments in `inputs.args`, keyword arguments in `inputs.kwargs` |
| `context.runtime` | Current worker's `RuntimeContext`; fields consistent with previous runtime table |
| `context.label(name)` | Read `context.case.labels[name]`; raises error if doesn't exist |
| `context.empty(shape, dtype=...)` | Create normal `torch.empty` Tensor on current NPU device |
| `context.shmem_tensor(shape, dtype=..., device_id=None)` | Create aclshmem Tensor and automatically release after current case completes; requires `SHMEM` declared |
| `context.defer(callback)` | Register custom cleanup function, called when current case ends in reverse registration order |

For example:

```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` is not a command-line argument or module-level configuration, but the return value of `build_plan(context)`. Only operators using `build_plan` need to create it; simple operators directly implementing `custom_op(...)` don't need `LaunchPlan`.

| Field | Required | Meaning |
| --- | --- | --- |
| `launch` | Required | Callable with no parameters. Framework calls it to execute custom operator once; usually a pre-prepared Triton launcher |
| `outputs` | Required | Tensor to compare with Golden result after `launch()` executes, or nested Tensors composed of tuple, list, dict |
| `before_golden` | Optional | No-parameter function executed before each Golden call; used to restore inputs or prepare Golden state, not included in Event measurement interval |
| `before_custom` | Optional | No-parameter function executed before each `launch` call; used to restore output, workspace, or input state, not included in Event measurement interval |
| `cleanup` | Optional | No-parameter resource release function executed once when current case ends |
| `metadata` | Optional | Reserved operator additional information; currently terminal and performance CSV don't output this field, usually not needed |

Most common return form has only two required fields:

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

## 2. Running

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

By default, both correctness and Torch Event performance tests run simultaneously. To run correctness test only, add `--correctness-only`.

The above command defaults to using `0,1,2,3`. To specify other cards, set before the command, for example `ASCEND_RT_VISIBLE_DEVICES=3,4,5,6`; `--num-cards` cannot exceed the number of explicitly listed cards. The launcher does not check if cards are occupied by other processes.

`path/to/operator.py` is the triton operator file.
`--output` is the correctness and performance results CSV.

## 3. Execution semantics

The framework strictly separates into two phases:

1. Complete correctness testing for all shapes sequentially; failure won't terminate remaining shapes early. The framework recursively calls `torch.testing.assert_close` to judge correctness against Golden, fixed using `rtol=1e-3`, `atol=1e-3`.
2. Only if all shapes are correct, re-prepare inputs and start full performance testing.

When any shape fails, the performance CSV only writes the header without data rows, and the process eventually returns a non-zero exit code; specific failed shapes and error messages are output to terminal.

When all shapes are correct, the common timer tests each shape in two independent phases: first execute 5 Golden warmup and 50 Golden timing runs, then execute 5 custom/Triton warmup and 50 custom/Triton timing runs. The timer is also responsible for barrier outside Events, NPU synchronization, cross-rank MAX per iteration, and min/max/median/average and speedup.

## 4. CSV output fields

| Field | Meaning |
| --- | --- |
| Fields in `OperatorCase.labels` | Expanded as independent columns per user-defined names, e.g. `S`, `H`, `D` |
| `golden_*_us`, `custom_*_us` | min/max/median/average for both paths, in microseconds |
| `speedup_*` | `golden / custom` for corresponding statistics; greater than 1 indicates custom operator is faster |

Dimensions that need to appear in the performance table should be explicitly declared by the operator through `labels`.

## 5. Example

The QKV operator example is located at `perf_unified/examples/qkv_all2all_udma.py`.