Background（背景信息）

We write compute kernels in PyPTO, but today we can only run them standalone from a Python script. There is no supported way to put one inside a real model and ship it.

To get a PyPTO kernel into a production inference model, someone on our side has to hand-write the framework plumbing for every single operator - registering it, making it survive model export, and describing it to the graph compiler. That is several hundred lines per operator, none of it about the computation we care about. It is slow, only two or three people can do it, and whenever a framework interface changes, every operator we already wrote has to be fixed one by one.

We need a PyPTO kernel to become a deployable operator without any of that. The author should describe the kernel and how its output shape and dtype are derived; everything the frameworks need should be produced from that description, so the operator can be exported into a model, compiled with the rest of the graph, and executed on device.

Origin（信息来源）

Requested by the team deploying PyPTO-authored operators into customer inference models on Kirin, where the model is compiled ahead of time and executed on device.

Benefit / Necessity（价值/作用）

Why we need it

• A custom operator is a specialist task measured in days. It should be something a kernel author finishes alone, in one file, without learning the framework interfaces.
• The framework code is copied per operator today, so an interface change means a migration across every operator we own. That cost should be paid once, centrally.
• "Run the kernel" and "ship the kernel" are currently two separate pieces of work. One description should serve both.

Value

• A much lower barrier to contributing an operator, so more of our model-level performance ideas actually get tried.
• One operator serving many configurations through typed parameters, instead of a new kernel per configuration.
• A shipped operator that stays editable and redeployable in place, so on-device debugging does not mean rebuilding from scratch.

Where we would use it

• A custom fused operator (attention, a normalization, a fused elementwise chain) inside an otherwise standard inference model.
• Pipelines where one custom operator feeds the next, and models with several of them.
• Mixed models, where our operators sit next to standard operators or operators authored by other teams.

Delivered means

• Worked examples covering these shapes, each runnable end to end.
• Those examples verified on real hardware: exported, compiled, executed on device, and the device output compared against a CPU reference.

Design（设计方案）

Overall idea: the operator is described once, in Python, by the person who wrote the kernel. Everything the frameworks need is derived from that description - nothing framework-specific is written by hand.

1. The author describes the operator in one place

my_op = CustomOp(
    kernel      = my_kernel,        # the PyPTO kernel, kept as an ordinary editable file
    infer_shape = my_infer_shape,   # output shape, from the input shapes
    infer_dtype = my_infer_dtype,   # output dtype, from the input dtypes
    reference   = my_cpu_version,   # used on CPU, and as the correctness baseline
    parameters  = [ scale: float, channels: int ],
)

The shape and dtype rules are given once and used both while the model is traced and by the deployed operator, so they cannot drift apart.

2. Export and build then just work

The author exports the model as usual, and the operator is built and compiled together with the rest of the graph:

describe operator -> export model -> build operator -> compile model -> run on device

3. The same description runs everywhere

Which one runs is chosen by how the model is run, rather than by editing the description: on CPU it uses the reference version, on the simulator and on device it uses the real kernel. That is what lets us compare device output against CPU output.

4. The shipped kernel file stays the real kernel

Editing it and redeploying is a single step.