Functions and Program Structure

Function declaration forms, parameter directions, cross-module reuse, and how to print IR back to Python syntax.

JIT constants and compilation reuse

A @pl.jit body is parsed into @pl.program source and re-parsed in a namespace holding only pl and pld, so any name it inherits from its own module or an enclosing function must be replaced, at its use site, by source text that evaluates back to the same value. That covers literals (int, float, bool, str, None), the pl dtype and enum constants (pl.INT8, pl.Mem.Vec, pl.PadValue.zero, pl.NZ), and lists/tuples nested from those — a shape or a rounding mode held in a constant, for instance. A value with no source form is left alone, so the name survives and the parser reports it.

Every name that folds is also in the compilation key, including constants used by transitive JIT helpers and source annotations. The key hashes the emitted text for each name, straight from the function the specializer folds with, so the two cannot drift: a constant that changes the generated source changes the key by construction, and one that does not fold contributes nothing. Rebinding a referenced constant causes a new specialization; changing an unrelated global or a name shadowed by a body-local variable does not invalidate the body dependency key.

BLOCK = 32

@pl.jit
def slice_kernel(x: pl.Tensor[[128, 128], pl.FP32]) -> pl.Tensor[[BLOCK, 128], pl.FP32]:
    with pl.at(level=pl.Level.CORE_GROUP):
        result = pl.slice(x, [BLOCK, 128], [0, 0])
    return result

first = slice_kernel.compile()
BLOCK = 64
second = slice_kernel.compile()  # A distinct specialization with 64 rows.

Name resolution, key construction, and specialization share a per-call namespace snapshot. Constants rebound after capture affect the next call, not the artifact being compiled. Each helper retains its own namespace, even when constants have the same name. Rebinding a referenced JIT helper also refreshes the dependency graph. Each concurrent call retains its captured graph throughout key construction and specialization. The hash also includes each helper's function type, level, and auto_scope setting, so rebinding an identically sourced function with different compilation attributes cannot reuse the old artifact.

Each module's globals are copied once per request; helpers use closure overlays on that shared snapshot. An empty closure cell still shadows a same-named global. Validated graphs retain their Python source hashes, and declared layouts are reused while their annotation bindings stay unchanged. External source files are still checked on each request. Referenced closure constants are covered by the source dependency hash without a separate closure-key component.

The snapshot copies bindings only: mutating arbitrary configuration objects or editing compiler/source files during compilation is not supported by this constant-tracking mechanism. Persistent reuse additionally requires the full artifact identity and cache policy.

Compile options and diagnostic requests

JIT calls and kernel.compile() resolve compile options before looking up an in-process artifact. Omitting config uses the same defaults as RunConfig(): a2a3sim, the default optimization strategy, and pass dumps disabled. Semantic options (including the effective planner, runtime ABI, and source-location emission controlled by PYPTO_EMIT_PTO_LOC) participate in the key. Runtime controls such as device_id, codegen_only, and save_kernels alone do not.

The following requests compile afresh on every call, without reading or adding an entry in the artifact cache:

• dump_passes=True or a dump level other than PassDumpLevel.NONE, and dump_ptoas_passes=True.
• compile_profiling=True, an active CompileProfiler, or profiling enabled through PYPTO_COMPILE_PROFILING.
• An explicit save_kernels_dir, or a nonempty PYPTO_PROG_BUILD_DIR.
• Explicit diagnostic settings, or an active PassContext with instruments or verification/diagnostic settings different from the pipeline defaults.

This ensures a warm kernel still emits requested dumps and reports. Repeating a request regenerates the output; a failed diagnostic compile leaves ordinary cached entries intact. Automatically allocated JIT output directories are unique per compilation, so fresh diagnostics cannot overwrite a cached artifact. Explicit output directories are used as requested. Diagnostic controls do not split compilation keys. Conflicting explicit settings and an active PassContext raise the same error on cache hits as on fresh compilation. A plain planner/runtime context can still reuse a matching cached artifact.

from pypto.runtime import RunConfig

cached = slice_kernel.compile()
slice_kernel.compile(config=RunConfig(dump_passes=True, save_kernels_dir="debug_kernel"))
assert slice_kernel.compile() is cached

Prepare binaries without executing

kernel.warmup() uses the same specialization, configuration, and in-process object cache as kernel.compile(), and also prepares every kernel and orchestration binary before returning. It initializes no NPU and creates no runtime worker. The build host still needs the target compiler, SDK, and runtime Python/native dependencies.

import pypto.language as pl
from pypto.runtime import RunConfig

@pl.jit
def add_three(
    x: pl.Tensor[[16, 16], pl.FP32],
    out: pl.Out[pl.Tensor[[16, 16], pl.FP32]],
):
    with pl.at(level=pl.Level.CORE_GROUP):
        tile = pl.load(x, [0, 0], [16, 16])
        pl.store(pl.add(tile, 3.0), [0, 0], out)
    return out

config = RunConfig(platform="a2a3")
prepared = add_three.warmup(config=config)  # No sample tensor allocation.
# Later, on a host with an available NPU:
# prepared(x, out, config=config)

You can supply sample tensors as with compile(); warmup reads their metadata, not their contents. With complete tensor annotations, omit the tensors and use scalar defaults or keyword values. pl.RUNTIME keeps a scalar unspecialized in annotation-driven mode, and dynamic extents retain the existing compile rules. When calling the returned compiled object, supply its full parameter list, including scalar arguments; JIT defaults are resolved during compilation. Execution-only settings such as codegen_only do not suppress binary preparation.

The result is the same compiled object selected by compile(), with live IR retained for fresh compilation. Warmup prepares every chip-level child of a DistributedCompiledProgram and every orchestration sub-build of a multi-orchestration CompiledProgram. It does not call the distributed object's prepare(): that method creates a live worker for execution. Compilation errors propagate to the caller; a failed binary build can be retried by calling warmup again. Diagnostic/output requests still compile afresh as described above.

With persistent caching enabled, warmup automatically publishes or reuses READY artifacts through the runtime protocol. Persistence is disabled by default. Read-only misses, unsupported inputs and storage failures can prepare private results. A restored result has .program is None; disable persistence or use specialize()/lower() to require IR. The public cache policy, statistics and metadata-only warmup CLI are documented in Persistent JIT Cache.

Functions

# Single return type
def function_name(param1: pl.INT64, param2: pl.FP32) -> pl.INT64:
    x: pl.INT64 = param1 + 1
    return x

# Multiple return types
def function_name(x: pl.INT64) -> tuple[pl.INT64, pl.INT64]:
    y: pl.INT64 = x + 1
    z: pl.INT64 = x * 2
    return y, z

# No return types
def function_name(x: pl.INT64):
    y: pl.INT64 = x + 1

# With function type
@pl.function(type=pl.FunctionType.Orchestration)
def orchestrator(n: pl.INT64) -> pl.INT64:
    return n + 1

@pl.function(type=pl.FunctionType.InCore)
def aicore_kernel(x: pl.INT64) -> pl.INT64:
    return x * 2

Function Types

Type	Usage	Description
pl.FunctionType.Opaque	Default	Unspecified function type
pl.FunctionType.Orchestration	Host/AICPU	Control flow and dependency analysis
pl.FunctionType.InCore	AICore	Sub-graph on specific AICore (unspecialized)
pl.FunctionType.AIC	Cube core	Cube core kernel (specialized InCore)
pl.FunctionType.AIV	Vector core	Vector core kernel (specialized InCore)
pl.FunctionType.Group	Multi-core	Co-scheduled group of AIC + AIV kernels
pl.FunctionType.Graph	Host/AICPU	Recordable orchestration fragment, replayed by the host_build_graph runtime (see below)

When no type is specified, functions default to Opaque.

Graph Fragments

pl.FunctionType.Graph marks a function as a recordable orchestration fragment. Under the host_build_graph runtime each call site becomes one task launch that the runtime records on the first call and replays afterwards, so N calls cost one graph build instead of N:

@pl.program
class Decoder:
    @pl.function(type=pl.FunctionType.Graph)
    def layer(self, cur, normed, next_hidden, wq, layer_base: pl.Scalar[pl.INDEX]):
        ...

    @pl.function
    def decode(self, cur, normed, next_hidden, wq):
        for i in pl.range(40):
            self.layer(cur, normed, next_hidden, wq, i * 5120)

One Graph function is one recorded topology: the runtime identifies the recording by the address of the emitted C++ function, so there is no cache key to name or keep unique.

Parameter Directions

Parameters can have In (default), Out, or InOut directions using wrapper types:

@pl.function(type=pl.FunctionType.InCore)
def kernel(
    qi: pl.Tensor[[16, 128], pl.BF16],                   # In (default)
    output: pl.InOut[pl.Tensor[[16, 128], pl.FP32]],      # InOut
    result: pl.Out[pl.Tensor[[16, 128], pl.FP32]],        # Out
    scale: pl.Scalar[pl.FP32],                             # In (default)
) -> pl.Tensor[[16, 128], pl.FP32]:
    ...

Direction	Wrapper	Description
In	None (default)	Read-only input parameter
Out	pl.Out[type]	Write-only output parameter
InOut	pl.InOut[type]	Read-write input/output parameter

Constraint: Scalar parameters cannot have InOut direction (raises ParserTypeError).

Writing an Out / InOut parameter

A bare assignment does not write the parameter. It rebinds the Python name: the parameter Var is re-pointed at a freshly computed tensor and the caller's buffer is left untouched. The program still compiles and runs. What the caller gets back depends on the direction: an Out buffer is freshly allocated and never initialised, so it reads as garbage; an InOut buffer still holds the input the caller passed in, so the result is silently stale.

out = pl.add(a, b)          # ❌ writes nothing
out[:] = pl.add(a, b)       # ✅ writes the whole tensor

Use the subscript form to write through the parameter:

Spelling	Writes	Use when
out[:] = <expr>	the whole tensor	the result is the entire output
out[<slices>] = <expr>	that sub-window	writing part of the output
out = pl.assemble(out, <expr>, <offset>)	a window at <offset>	the explicit form the subscript sugar builds
out = <expr>	nothing	never — see the warning below

Only the first and third are equivalent, and only when the slice covers the full extent and <offset> is all zeros.

The OutParamWriteDropped warning

The compiler reports a bare assignment that drops the write:

[warning] [OutParamWriteDropped] (pipeline_input) Assigning to Out parameter 'out'
in function 'main' rebinds the name only — the caller's buffer is never written.
Use 'out[:] = <expr>' to write the whole tensor, or 'out[<slices>] = <expr>' for
a sub-window. at repro.py:12:9

The check is data-flow based, not syntactic. A value can reach the parameter without naming it — through a loop carry, for instance — and that is a genuine write-through, so it stays silent:

for col, (d,) in pl.range(0, n, chunk, init_values=(data,)):
    d = pl.store(local, [0, col], d)
    staged = pl.yield_(d)
data = pld.tensor.allreduce(staged, signal, ...)   # `staged` *is* `data`; no warning

The check is deliberately conservative: a value that merely reads the parameter, such as out = pl.add(out, b), is not reported even though it also drops the write. Telling "reads the parameter" from "writes through the parameter" needs per-op write semantics the operator registry does not record, and a false warning on correct code costs more than a missed one. Write out[:] = pl.add(out, b) when the whole tensor is meant.

Disable the check with disabled_diagnostics if it is not useful for your program:

disabled = passes.DiagnosticCheckSet()
disabled.insert(passes.DiagnosticCheck.OutParamWriteDropped)
ir.compile(program, disabled_diagnostics=disabled)

How a @pl.program Class Is Located

@pl.program parses the class body from source, so it first has to find the class statement that produced the decorated object. The class name alone does not identify it: one function may define the same name in several branches, and every one of them carries the same __qualname__.

The decorator resolves this from the line numbers of the methods in the class body, so each branch is parsed from its own source:

def make(case):
    if case == "add":
        @pl.program
        class Prog:                                # parsed from *this* body
            @pl.function
            def main(self, x: pl.Tensor[[8], pl.FP32]) -> pl.Tensor[[8], pl.FP32]:
                return pl.add(x, 1.0)
        return Prog

    @pl.program
    class Prog:                                    # and this one from *this* body
        @pl.function
        def main(self, x: pl.Tensor[[8], pl.FP32]) -> pl.Tensor[[8], pl.FP32]:
            return pl.mul(x, 3.0)
    return Prog

When the definitions genuinely cannot be told apart, the decorator raises a ParserSyntaxError naming every candidate line rather than picking one — a wrong pick would compile a body you never wrote. Give the classes distinct names, or define the class once and vary it through a closure variable:

def make(scale):
    @pl.program
    class Prog:                                    # one definition, parameterised
        @pl.function
        def main(self, x: pl.Tensor[[8], pl.FP32]) -> pl.Tensor[[8], pl.FP32]:
            return pl.mul(x, scale)
    return Prog

Cross-Module Function Reuse

Functions defined outside a @pl.program class can be reused via two mechanisms.

External @pl.function Calls

An externally-defined @pl.function can be called by name inside @pl.program. The function is automatically added to the Program and an ir.Call(GlobalVar, args) is emitted.

@pl.function
def softmax(x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
    ...

@pl.program
class MyModel:
    @pl.function
    def main(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
        y: pl.Tensor[[64], pl.FP32] = softmax(x)   # ir.Call(GlobalVar("softmax"), [x])
        return y

Rules:

• Uses the function's .name as GlobalVar (aliases are transparent)
• External and internal function names must not conflict
• Two different externals with the same .name is an error
• Same external called from multiple methods is added once

@pl.inline Decorator

@pl.inline captures a function for statement-level inlining. No function is added to the Program — the body is expanded at each call site.

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

@pl.program
class MyModel:
    @pl.function
    def main(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
        y: pl.Tensor[[64], pl.FP32] = normalize(x)  # statements inlined in-place
        return y

Rules:

• Argument count must match parameter list exactly
• Closure variables from the inline definition site are available
• Inline functions can be called multiple times (each expansion is independent)
• Nested inline calls are supported

Complete Example

Tensor Operations (Loop with iter_args)

# pypto.program: my_program
import pypto.language as pl

def loop_sum(n: pl.INT64) -> pl.INT64:
    sum_init: pl.INT64 = 0
    for i, (sum,) in pl.range(n, init_values=(sum_init,)):
        sum = pl.yield_(sum + i)
    return sum

Tile Operations (Tile-based computation)

import pypto.language as pl

@pl.program
class BlockExample:
    @pl.function
    def tile_add(
        self,
        input_a: pl.Tensor[[64, 64], pl.FP32],
        input_b: pl.Tensor[[64, 64], pl.FP32],
        output: pl.Tensor[[64, 64], pl.FP32],
    ) -> pl.Tensor[[64, 64], pl.FP32]:
        tile_a: pl.Tile[[64, 64], pl.FP32] = pl.load(input_a, [0, 0], [64, 64])
        tile_b: pl.Tile[[64, 64], pl.FP32] = pl.load(input_b, [0, 0], [64, 64])
        tile_c: pl.Tile[[64, 64], pl.FP32] = pl.add(tile_a, tile_b)
        result: pl.Tensor[[64, 64], pl.FP32] = pl.store(tile_c, [0, 0], output)
        return result

Printing IR Nodes

Use as_python() on any IR node to get its Python representation:

print(stmt.as_python())          # "x: pl.Scalar[pl.INT64] = a + b" (default "pl" prefix)
print(stmt.as_python("ir"))      # "x: ir.Scalar[ir.INT64] = a + b" (custom prefix)

Concise Mode

Pass concise=True to omit intermediate type annotations. Function signature types (parameters and return) are always preserved:

print(func.as_python())                  # verbose (default): type on every assignment
print(func.as_python(concise=True))      # concise: omits intermediate type annotations

Verbose output:

def main(self, x: pl.Tensor[[64, 128], pl.FP32]) -> pl.Tensor[[64, 128], pl.FP16]:
    y: pl.Tensor[[64, 128], pl.FP32] = pl.some_op(x)
    result: pl.Tensor[[64, 128], pl.FP16] = pl.cast(y, pl.FP16)
    return result

Concise output:

def main(self, x: pl.Tensor[[64, 128], pl.FP32]) -> pl.Tensor[[64, 128], pl.FP16]:
    y = pl.some_op(x)
    result = pl.cast(y, pl.FP16)
    return result

The free function ir.python_print(node) is also available and supports the same parameters.

See Also

• Python IR Syntax Specification — types and expressions
• Statements and Control Flow — statement forms inside a function body
• Integrating Hand-Written C++ Kernels — calling external kernels