Functions and Programs

How a Python function becomes an IR function, which decorator to reach for, and how functions call each other.

Prerequisites: Types.

Concept

A decorator does not wrap your function — it parses its source. The body never executes as Python. That single fact explains most of what follows: why closure variables behave the way they do, why pl.yield_ is meaningful only inside one, and why an error in a kernel body is reported at parse time with a line number rather than at call time with a traceback.

@pl.jit is how you write PyPTO kernels. Types come from the arguments at the first call, the function specializes, and sub-functions are discovered automatically — you call them by name and the decorator finds them. It is what examples/ uses and what the rest of this manual uses.

You will also meet @pl.function inside @pl.program: a class where each method is one IR function and calls between them are written self.other(...). That form is a one-to-one transcription of the IR, and it exists mainly for writing compiler test cases, where a test needs to state a program's exact shape without ever running it. As a user you do not need it — the section below is there for when you read a compiler test or a piece of printed IR.

Quickstart: an entry point and a device kernel

import pypto.language as pl

@pl.jit.incore
def add_kernel(
    a: pl.Tensor[[128, 128], pl.FP32],
    b: pl.Tensor[[128, 128], pl.FP32],
    out: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
):
    out[:] = pl.add(a, b)
    return out

@pl.jit
def entry(
    a: pl.Tensor[[128, 128], pl.FP32],
    b: pl.Tensor[[128, 128], pl.FP32],
    out: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
):
    out = add_kernel(a, b, out)      # sub-function discovered automatically
    return out

Line	What it does
@pl.jit.incore	Marks a device kernel — the execution plane, where operators live
@pl.jit	Marks the chip-level entry point — the control plane, which dispatches
add_kernel(a, b, out)	A plain call; the decorator discovers the callee and wires it in
out: pl.Out[...]	Declares the direction, which is how the compiler orders this task

To read the IR this becomes, call entry.lower(*args) for the post-pass ir.Program, or entry.compile(*args) and print compiled.program.as_python().

Mechanics

The @pl.jit family

Five variants, one per IR function kind, so a single program can span host, chip, and core levels:

Decorator	IR target	Use for
@pl.jit	Orchestration	Chip-level entry point that dispatches InCore work
@pl.jit.host	level=HOST, role=Orchestrator	HOST entry — allocates window buffers, dispatches chip orchestrators per rank
@pl.jit.incore	InCore	A device kernel (accepts level= to target a specific hierarchy level)
@pl.jit.inline	Inline	Helper spliced into every call site by InlineFunctions
@pl.jit.opaque	Opaque	A separate IR function that may hold orchestration loops and pl.at scopes
@pl.jit.graph	Graph	A recordable orchestration fragment — the host_build_graph runtime records its task topology on the first call and replays it after, so N calls cost one graph build rather than N. Requires compiling under RuntimeKind.HOST_BUILD_GRAPH

@pl.jit.graph has a scope form: with pl.graph("name"): marks a region in place instead of splitting it into a function. The two compile to the same thing — the region is outlined into a Graph function named after name — so the choice is ergonomic. Use the decorator when the layer is already its own function; use the scope when the region is a slice of a larger orchestration body you would rather not split up:

@pl.jit
def decode(w: pl.Tensor, hidden: pl.InOut[pl.Tensor]):
    for layer in pl.range(40):
        with pl.graph("decoder_layer"):        # recorded once, replayed 39 times
            ...
    return hidden

The name is required and becomes the recorded graph's identity, so keep it stable. A Graph region cannot nest inside another Graph region, nor inside pl.at / pl.cluster / pl.spmd — those become a single device task, whereas a Graph region records a topology of them. All three are compile errors.

Sub-function dependencies (.incore / .inline / .opaque / .graph) are auto-discovered from the entry's body — call them by name. The name you call is resolved in the entry's own namespace, so an aliased import (from kernels import matmul as mm, or a plain mm = matmul rebinding) is discovered like any other binding; the generated program still names the function after its def. When two distinct sub-functions share that name — two modules each defining helper, or two kernels from one factory — the second one generated is suffixed (helper, helper__2), so both specializations survive; the entry keeps its own name. A @pl.jit.host entry additionally discovers @pl.jit chip-orchestration dependencies, so a full distributed program needs no @pl.program class.

The fragment below shows only the discovery structure — the kernel bodies are elided, and the distributed types it names are covered in the distributed chapter, which is not written yet:

import pypto.language.distributed as pld

@pl.jit.inline
def reduce_step(local, peer, out): ...

@pl.jit
def chip_orch(inp: pl.Tensor, out: pl.Out[pl.Tensor],
              data: pl.InOut[pld.DistributedTensor], peer: pl.Scalar[pl.INT32]):
    return reduce_step(inp, peer, out)      # auto-discovered sub-function

@pl.jit.host
def host_orch(
    inputs: pl.Tensor[[2, 1, 256], pl.FP32],
    outputs: pl.Out[pl.Tensor[[2, 1, 256], pl.FP32]],
):
    data_buf = pld.alloc_window_buffer(256 * pl.FP32.get_byte())
    for r in pl.range(pld.world_size()):
        data = pld.window(data_buf, [1, 256], dtype=pl.FP32)
        chip_orch(inputs[r], outputs[r], data, (r + 1) % pld.world_size(), device=r)
    return outputs

Plain @pl.jit entries do not discover other @pl.jit entries — only .host reaches across the chip boundary. That keeps two unrelated top-level kernels from silently folding into one program.

@pl.jit.host rejects level= (HOST is implicit).

What a sub-function hands back keeps its shape and dtype

Specialization stamps every generated parameter with a concrete shape and dtype, so a tensor a sub-function returns must be traceable to one. Both conventions work, and you can mix them in the same entry:

@pl.jit.inline
def make_pair(x: pl.Tensor[[1, 8], pl.FP32]):
    a = pl.create_tensor([1, 8], dtype=pl.FP32)   # helper allocates its own results
    b = pl.create_tensor([1, 8], dtype=pl.FP32)
    with pl.at(level=pl.Level.CORE_GROUP):
        a[:, :] = x[:, :]
        b[:, :] = pl.mul(x[:, :], 2.0)
    return a, b

@pl.jit.incore
def relu_kernel(x: pl.Tensor, out: pl.Out[pl.Tensor]):   # caller allocates, kernel fills
    ...

@pl.jit
def entry(x: pl.Tensor[[1, 8], pl.FP32], out: pl.Out[pl.Tensor[[1, 8], pl.FP32]]):
    a, b = make_pair(x)          # metas read out of make_pair's own body
    buf = pl.create_tensor([1, 8], dtype=pl.FP32)
    mid = relu_kernel(a, buf)    # mid aliases buf, so it inherits buf's meta
    ...

A device kernel (@pl.jit.incore) cannot allocate, so it always takes the second form — pl.create_tensor belongs on the control plane. An @pl.jit.inline helper is spliced into the caller, so either form is available to it.

A complete one-dimensional view of an NZ tensor also keeps inferable metadata. For w: pl.Tensor[[256, 512], pl.INT8, pl.NZ], w_flat = pl.reshape(w, [131072]) can be passed to a helper declared with src: pl.Tensor; specialization supplies src: pl.Tensor[[131072], pl.INT8] (ND). The reshape creates a view of the existing physical bytes without layout conversion or copying. Every source dimension and the target length must be static, and the target length must equal the source element count. Other NZ reshape targets and symbolic NZ flattening remain unsupported.

An extent the specializer cannot compute statically is not by itself a problem. A pl.create_tensor sized from a value only the device knows — pl.tensor.read(cfg, [0]), pld.world_size() — becomes a dynamic dimension and keeps flowing, and the shared pass pipeline judges whatever the program then does with it (loading the whole tensor as a tile, say, gets you InitMemRef requires static shape — the same error the equivalent @pl.program earns).

pl.reshape also preserves metadata when its target dimensions use known tensor dimensions and integer arithmetic (+, -, *, //, and unary -). For example, pl.reshape(src, [src.shape[0] * src.shape[1]]) keeps a dynamic extent when either source dimension is dynamic. The consuming JIT function reads the actual extent from the reshaped tensor's descriptor. The original reshape expression remains in IR, and existing reshape validation is unchanged. Calls to pl.tensor.dim support positional or keyword tensor/axis arguments and static constant or negative axes. Unknown runtime axes remain unresolved. Shape subscripts use nonnegative integer literal indices; use pl.tensor.dim(src, -1) for a negative axis. Dynamic unary + remains unsupported.

Intermediate scalar aliases work as well: total = src.shape[0] * src.shape[1] followed by flat = pl.reshape(src, [total]) preserves the dynamic extent. Scalar calculations stay scalar: inference records immutable expression snapshots and creates dynamic-axis symbols when recording a tensor shape. Rebinding a source name does not change an already computed alias. Chains of simple names such as total = half = total // 2 evaluate the right-hand side once and are lowered to individual assignments before parsing. Static integer aliases fold normally; unknown runtime values and invalid arithmetic remain unresolved. Aliases rebound under conditionals or loops are handled conservatively across control-flow boundaries, rather than assuming the value from one syntactic execution. After a tensor variable is rebound, shape reads use its current value, including whole-shape unpacking and results of tensor operations. Shape and dtype arguments such as shape=src.shape and dtype=src.dtype also use the current tensor type. Unrelated branches and loops keep existing aliases; a loop-carried alias becomes unknown until a definite assignment establishes its value in the current iteration. Mutually exclusive if/elif/else arms start from the same incoming scalar bindings and tensor metadata. Conditional scalar writes remain unknown after the branches join. Tensor metadata established identically in both arms remains usable; conflicting or missing tensor metadata remains unknown. Metadata keeps independent derived dimensions separate when a Python variable is reused or a helper is called repeatedly. Use explicit assignments such as n = n + step; augmented assignment (n += step) remains unsupported.

What does not resolve is a returned tensor whose shape the specializer cannot reach at all: a pl.reshape sized from an unknown runtime scalar, such as pl.tensor.read(cfg, [0]), or a result rebound through an operation the specializer does not model. That surfaces as missing inferred tensor metadata for parameter '<name>' at the next call that consumes it — the error names the consumer, but the fix belongs at the producer: give the producing statement a statically inferable shape, or pass the buffer in as a pl.Out[...] parameter.

Three constraints that decide whether a jit kernel compiles

These are the failures new @pl.jit code hits, in the order it hits them.

1. A @pl.jit entry body cannot hold operators. It is an Orchestration function — the control plane. Put the operators inside with pl.at(level=pl.Level.CORE_GROUP):, or move them into a @pl.jit.incore sub-function.

@pl.jit
def bad(x: pl.Tensor[[64, 64], pl.FP32], out: pl.Out[pl.Tensor[[64, 64], pl.FP32]]):
    out[:] = pl.add(x, x)        # ✗ Misplaced tensor op ... should be inside InCore block
    return out

@pl.jit
def good(x: pl.Tensor[[64, 64], pl.FP32], out: pl.Out[pl.Tensor[[64, 64], pl.FP32]]):
    with pl.at(level=pl.Level.CORE_GROUP):
        out[:] = pl.add(x, x)    # ✓
    return out

2. JITFunction has no as_python(). The IR does not exist until a specialization does. Call lower(*args) for the post-pass ir.Program, or compile(*args) and read compiled.program.as_python() for the specialized, pre-pass IR.

3. compile() takes the kernel's own arguments, not compile options. Compile-time knobs go through config=RunConfig(...). A stray compile(skip_ptoas=True) is bound against the kernel's signature and raises TypeError: got an unexpected keyword argument. @pl.jit detects whether ptoas is available on its own, so skip_ptoas is not something you need to pass.

A device kernel cannot return a scalar. The runtime's two task-argument channels are disjoint: scalars go in by value, and only tensors come back. So a value computed on device and needed by the caller — the orchestration function that launched the kernel, and the host beyond it — has to travel in a tensor. That covers a scalar tucked inside a pl.Tuple[...] return, too.

@pl.jit.incore
def bad(x: pl.Tensor[[64], pl.FP32]) -> pl.Scalar[pl.INDEX]:   # ✗ no carrier for the return
    ...

@pl.jit.incore                                                  # ✓ hand it back in a [1] tensor
def good(x: pl.Tensor[[64], pl.FP32], n_out: pl.Out[pl.Tensor[[1], pl.INT32]]):
    ...
# then, in the entry body, after the launch:
n = pl.tensor.read(n_out, [0])

The same rule applies to a scalar you assign inside with pl.at(...) and read after it — that is a kernel return in disguise. When the value only depends on things the entry body already has (a loop variable, a scalar parameter), the compiler moves the computation out of the scope for you; when it depends on device data, it asks you to route it through a tensor as above.

A helper that computes a scalar for the kernel is fine — write it @pl.jit.inline (FunctionType.Inline). It is spliced at the call site, so it is not a task and the rule does not apply to it.

@pl.function and @pl.program

You reach for this form when writing a compiler test case, not when writing a kernel. It describes the IR one-to-one: the class is the program, each method is a function, and the call graph is written out rather than discovered. @pl.jit specializes into exactly this shape — printing a compiled program shows you @pl.program source.

@pl.program
class Adder:
    @pl.function(type=pl.FunctionType.InCore)
    def add_kernel(self, a, b, out): ...

    @pl.function(type=pl.FunctionType.Orchestration)
    def entry(self, a, b, out):
        out = self.add_kernel(a, b, out)     # explicit cross-function call
        return out

Every method takes self (it is stripped from the IR), and Adder becomes an ir.Program — not a Python class you can instantiate. Adder.as_python() prints it.

type= names the plane each function belongs to:

Function type	Plane	Typical use
Opaque (default)	none yet	Standalone building block; takes its plane from where it is used
InCore	Execution	Load / compute / store kernel
Orchestration	Control	Creates tensors, dispatches InCore tasks
Inline	none	Spliced at every call site; leaves no function behind

A standalone @pl.function called from inside a @pl.program is added to the program as a separate function. @pl.inline (and @pl.jit.inline) instead expand at the call site and leave no function behind.

@pl.inline
def normalize(x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
    return pl.mul(x, 2.0)

The decorated object is a pl.InlineFunction — a template the parser splices, not a function you can call from Python.

Function attributes: pl.func_attr

Metadata about the function as a whole is declared with pl.func_attr({...}) as the first statement of the body:

@pl.program
class Kernels:
    @pl.function(type=pl.FunctionType.InCore)
    def kernel(self, x: pl.Tensor[[64, 64], pl.FP32], w: pl.Tensor[[64, 64], pl.FP32],
               out: pl.Out[pl.Tensor[[64, 64], pl.FP32]]):
        pl.func_attr({"stationary": w, "split": pl.SplitMode.UP_DOWN})
        ...

It reads oddly for a function-level declaration to sit inside the body, so it is worth saying why. A decorator is evaluated before the signature binds any name, so @pl.function(attrs={"stationary": w}) cannot be written at all — w does not exist yet. Body position places the declaration after the parameters are bound, which is what makes an attribute that references a parameter expressible. The alternatives are a positional index ({"stationary_param": 1}, which breaks the moment a pass reorders parameters) or a naming convention nothing enforces.

Rules worth knowing:

Rule	Why
Must precede every other statement	An attribute describes the whole function; it must not appear to start applying partway down a body. This also bounds what it can reference to the parameters.
A bare name is always a parameter	pl.func_attr({"n": k}) records the parameter k, never a same-named Python variable from the enclosing scope. Write Python constants as literals.
Multiple calls merge	A key declared twice is an error naming the key, so which value wins is never a matter of parse order.
auto_scope= and external_source= stay on the decorator	The parser reads them before it walks the body, so a body-position declaration would arrive too late to take effect.

@pl.function(attrs={...}) is deprecated and emits a DeprecationWarning. It still parses and behaves identically, but it can only ever carry values that reference nothing. Printed IR always uses a non-deprecated spelling — the pl.func_attr prologue, or the dedicated auto_scope= / external_source= keywords — so reparsing compiler output never warns.

Splitting compile from dispatch

@pl.jit kernels normally fuse specialize + compile + dispatch into one kernel(*args) call. JITFunction.compile(*sample_args) stops after compilation and hands back the CompiledProgram — for driving ChipWorker yourself, inspecting artifacts under compiled.output_dir, or validating codegen ahead of time.

compiled = my_kernel.compile(sample_x, sample_w, sample_out)
print("artifacts in:", compiled.output_dir)

The returned object is the same one the JIT cache holds, so a later call with the same specialization key returns the identical instance.

lower(*sample_args) stops one stage earlier: it runs the passes and returns the post-pass ir.Program, with no code generation, no ptoas, no artifacts, and no cache write. Use it to read lowered IR; use compile() when codegen itself is what you want to check. Both accept config=RunConfig(...), but lower() ignores the runtime and artifact fields. Compile options are in Compiling and the runtime surface in Running.

specialize(*sample_args) stops one stage earlier still: entry and deps are specialized into @pl.program source and parsed, and the pre-pass ir.Program comes back untransformed. Reach for it only when something downstream runs the pass pipeline itself — most of all ir.compile(program, output_dir=...), which runs passes and code generation, so handing it lower()'s output would run the pipeline twice.

program = my_kernel.specialize(sample_x, sample_w, sample_out)
ir.compile(program, output_dir="build/out", backend_type=BackendType.Ascend910)

It takes no config=: no pass runs here, so a RunConfig would have nothing to configure. Sample arguments may be omitted only when every tensor parameter is annotated with full shapes — a bare pl.Tensor has none to read.

Two kernels that specialize to the same program compare equal after passes, not before: the specializer renames SSA-rebound locals (out becomes out_v1), and canonicalization removes the difference. Compare lower() output when asserting equivalence against a hand-written @pl.program.

Three accessors describe the signature without specializing anything: param_names (in declaration order), output_param_names (the pl.Out[...] and pl.InOut[...] params, also in declaration order), and __name__.

External C++ kernels

A hand-written C++ kernel can be called like any other function. See Integrating Hand-Written C++ Kernels.

Edge Cases

Fatal pitfall: verify a new @pl.jit example with a full compile(), never with lower() alone. lower() stops after the passes, so the "operators in an Orchestration body" error above never fires — the kernel appears to pass and fails only when someone runs it for real.

Symptom	Likely cause	Fix
Misplaced tensor op ... should be inside InCore block	Operators directly in a @pl.jit body	Wrap in with pl.at(level=pl.Level.CORE_GROUP): or move to @pl.jit.incore
AttributeError: 'JITFunction' object has no attribute 'as_python'	Printing IR that does not exist yet	f.lower(*args), or f.compile(*args) then compiled.program.as_python()
lower() passes but compile() fails	lower() runs no code generation	Expected — use compile() to check codegen
TypeError: got an unexpected keyword argument	A compile option was passed to compile(), which binds against the kernel's signature	Pass config=RunConfig(...)
A second top-level kernel is missing from the program	Plain @pl.jit does not discover other @pl.jit entries	Use @pl.jit.host, or make the callee .incore / .opaque
auto_scope=False rejected	Used on .incore / .opaque	Put it on the entry or on an .inline helper
self missing from a @pl.program method	Every method needs it	Add self; it is stripped from the IR
A task cannot return a scalar	A device kernel declares a pl.Scalar return	Write it into a [1] tensor output and pl.tensor.read(t, [0]) after the launch
cannot return a scalar naming a variable in a scope	A scalar assigned inside pl.at(...) is read after it, and depends on device data	Route it through a [1] tensor, or move the computation out of the scope

Worked example

examples/utils/cross_function_calls.py — @pl.jit.inline helpers auto-discovered as deps of a @pl.jit entry and spliced at the call site.

See Also

• Control Flow — loops and conditionals inside these bodies.
• Scopes and Placement — pl.at and the other placement scopes.
• Quickstart — the same decorators in a worked example.
• InlineFunctions — how Inline bodies are spliced.
• Integrating Hand-Written C++ Kernels — calling external kernels.