Functions and Program Structure

Legacy Python API examples: platform= calls on this page are retained for compatibility. They emit DeprecationWarning; new code should use backend_type and execution_mode. See the target mapping.

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

JIT implementation structure

The JIT implementation lives in python/pypto/jit/. The public jit object and JITFunction remain in decorator.py; private analysis modules supply metadata without owning compiled-object caches or dispatching kernels.

Module	Responsibility
decorator.py	Decorators, JITFunction, dependency discovery, cache ownership, and compile/execute entry points
_arguments.py	Optional torch loading, tensor argument metadata, annotations, and signature/constexpr validation
_ast_analysis.py	Shared read-only AST cache, source names, and declared dynamic dimensions
_call_binding.py	Call-site argument mapping, constexpr variants, and dynamic-dimension propagation across dependencies
_tensor_meta.py	Source-ordered local tensor metadata, branch/loop bindings, and recursive dependency-return inference
specializer.py	Specialization context, annotation handling, program source assembly, and diagnostic source maps
_specialize_body.py	Stateful AST body rewriting, variable renaming, and lexical scope handling
_specialize_values.py	Shared constant and dtype/layout source rendering for specialization and cache keys
_specialize_types.py	Shared DynDim, ShapeDim, and TensorMeta definitions; existing imports from specializer remain valid
_source.py, cache.py, _persistent.py	Request namespace snapshots, cache keys, and persistent artifact resolution

JITFunction owns the dependency-graph state, compiled-object caches, and their lock. Key construction and specialization continue to share one request's namespace snapshot. Metadata inference receives dependency discovery as an explicit callback and carries it through recursive return analysis, so it does not import the decorator at runtime. Shared helper lookups use their owning modules, keeping AST caches and test patch targets in one place.

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), distributed enum constants (pld.ReduceOp.Max, pld.NotifyOp.AtomicAdd, pld.WaitCmp.Ge), and lists/tuples nested from those — a shape, rounding mode, or collective reduction selected by a factory, 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.

Dynamic dimension metadata

Before parsing the specialized program, the JIT infers local TensorMeta values to type arguments passed to bare-annotated helpers. Scalar bindings use int constants, existing-axis DynDim references, immutable expression tuples, or None (unknown). For width = n * 8, the snapshot is (ast.Mult, (n_dim, 8)). Assignment aliases retain that tuple; subsequent arithmetic creates a different snapshot. Rebinding an operand's source name does not change an earlier snapshot.

Reshape inference gives scalar expressions result-specific DynDim symbols when recording TensorMeta.shape; scalar snapshots stay outside tensor types. Repeated uses of the same snapshot within one result shape share a symbol; separate results get independent bindings. The callee reads the dimension from its actual tensor argument's descriptor. The original arithmetic and reshape validation remain in the program; placeholder bounds do not replace runtime extents.

synthesized records an automatically generated placeholder symbol and retains its existing annotation-overlay and helper-return behavior.

Tile parameters in JIT inline helpers

An inline helper can take an explicitly typed on-chip Tile:

@pl.jit.inline(auto_scope=False)
def keep_tile(tile: pl.Tile[[8, 64], pl.FP32, pl.MemorySpace.Vec]) -> pl.Tile:
    return tile

Specialization preserves the declared shape, dtype, and memory space, folding module and closure constants in the annotation. When a postponed annotation uses the helper's own pl.constexpr parameter, that call's constexpr binding is substituted before the annotation is evaluated; it takes precedence over a same-named module or closure constant. Different bindings produce separate helper specializations. DSL aliases such as p.Tile or a directly imported Tile are rendered as canonical pl paths, preserving TileView and MemRef metadata. An unrelated user variable named pl does not affect this rendering. For an eager Tile annotation whose factory-local constants are absent from the function closure, specialization can retain the already evaluated annotation. Independent shape axes and dtype survive even when other axes depend on helper parameters; those dependent components are substituted per call. The merged contract also contributes to the cache identity. Constexpr memory space parameters are substituted while independent shapes evaluated in the factory remain intact. Metadata replacements must retain their original kind; ambiguous replacements produce a diagnostic. If one dependent expression mixes a helper parameter with an unavailable factory local, pass that local as another constexpr parameter or capture it in the helper body. Dynamic symbols referenced only by Tile parameter types receive their own module-level pl.dynamic() declarations, even when no Tensor parameter uses them. This also applies when the emitted symbol name matches a Tile or Tensor parameter name; actual Scalar dimension parameters remain local. If a free dynamic symbol and a distinct Scalar parameter share a label, specialization reports the ambiguity; give the free symbol a unique pl.dynamic() label.

Before substituting arguments, InlineFunctions checks Tile dtype, physical rank and shape, and conflicting explicitly resolved memory spaces. Generic physical dimensions bind to the actual Tile's physical extents, including integer affine bindings such as 2*M to 16 (M=8); valid extents do not constrain that check. Known Scalar values are substituted first: n=8 and n+M=16 infer M=8. Separate Tile arguments can bind their generic physical dimensions independently, including when the helper returns its Tile parameters. Returned-parameter types also follow simple immutable aliases such as tmp = second; return tmp. Dimensions shared with Tensor types, explicit Tile views, or Scalar parameters retain their shared constraints, including in helpers with no return value. Actual Scalar values also constrain view-only dimensions: n=8 cannot match an actual valid extent of 3 when the declared valid extent is n. Scalar-dependent view expressions such as n - 1 are compared after argument substitution, including valid extents, strides, and offsets. This includes indirect dependencies through shared dimensions such as rows = 2*n. An omitted actual valid shape uses its physical shape, and an omitted offset is zero. A Scalar-dependent stride contract requires explicit actual stride metadata. JIT specialization also retains nested pl.Tuple parameter annotations, including the complete Tile types, their dynamic symbols, and per-call constexpr bindings. The same physical and view checks apply recursively to Tiles in tuple parameters. Dynamic dimension names emitted from Tile annotations must use assignable Python identifiers in normalized spelling, must not be Python keywords, and must not shadow the DSL aliases pl or pld. Inlining substitutes the caller's Tile without copying it, preserving its valid shape. Specialization also retains evaluable Tile annotations on non-inline functions. This preserves their declarations without adding Tile arguments to device dispatch interfaces; ordinary InCore entry points still reject Tile parameters.

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 entirely: scalar parameters are runtime values, so they need no value here, and dynamic extents retain the existing compile rules. A pl.constexpr parameter is the exception: its value selects the artifact, so it must be supplied as a keyword or a signature default. When calling the returned compiled object, supply its full parameter list, including scalar arguments — that is where each scalar's value is read. A constexpr parameter is absent from that list; the artifact already resolved it. 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)

Generic Function attrs print as pl.func_attr({...}) directives in the body prologue. They remain Function metadata when parsed, including on HOST SubWorkers: the captured Python InlineStmt.body excludes these directives while preserving the implementation and its docstring. Parameter-valued attrs resolve to the function's own parameters. The prologue must precede executable statements (a leading docstring is allowed); misplaced directives and duplicate keys are errors, also for SubWorkers.

Printer implementation

The private IRPythonPrinter declaration lives in src/ir/transforms/python_printer.h. Its implementations share one visitor instance, so variable renaming, signature context, indentation, and scope state remain consistent across recursive expression, type, and statement printing.

Source file in src/ir/transforms/	Responsibility
python_printer.cpp	Public entry points and format callback, name disambiguation, function/program printing and dependency order
python_printer_expr.cpp	Expression precedence, constants, calls/submits, and attribute values
python_printer_stmt.cpp	Control flow, scopes, comments, SSA yields, and window-buffer declarations
python_printer_type.cpp	Type annotations, shapes, memory references, and tensor/tile views

The public declarations remain in include/pypto/ir/transforms/printer.h. Sub-expression printers inherit the enclosing variable-name maps, preserving bindings in shape expressions, memory offsets, and slot indices.

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