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.

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

What does not resolve is a returned tensor whose shape the specializer cannot reach at all: a pl.reshape whose target shape is not static (a reshape is constrained by its source's element count, so no dynamic dimension can stand in for it), 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.Ascend910B)

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.