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):

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:

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):

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:

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:

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

2. Running

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:

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.
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.