Skip to content

pl

The names you reach directly on pl — decorators, types, control-flow constructs, and the type-dispatched operator wrappers. An operator here dispatches on what you hand it: pl.add of two tiles is pl.tile.add, of two tensors pl.tensor.add. See Choosing a namespace.

PyPTO Language module - Type-safe DSL API for writing IR functions.

This module provides: - function decorator for parsing DSL functions to IR - Tensor type for tensor annotations and runtime wrapping - Tile type for tile annotations and runtime wrapping - Type-safe operation wrappers (tensor., tile., system.*, and unified ops) - DSL helpers (range, yield_) - DataType constants

Typical usage

import pypto.language as pl

@pl.function def my_func(x: pl.Tensor[[64, 128], pl.FP16]) -> pl.Tensor[[64, 128], pl.FP32]: result: pl.Tensor[[64, 128], pl.FP32] = pl.create_tensor([64, 128], dtype=pl.FP32) return result

@pl.function def block_func(x: pl.Tensor[[64, 64], pl.FP32]) -> pl.Tensor[[64, 64], pl.FP32]: tile: pl.Tile[[64, 64], pl.FP32] = pl.load(x, [0, 0], [64, 64]) result: pl.Tile[[64, 64], pl.FP32] = pl.add(tile, tile) return pl.store(result, [0, 0], x)

@pl.function def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: return x

jit = _JITDecorator() module-attribute

JITFunction

A JIT-compiled function with shape specialization and caching.

Created by the @jit or @jit.incore decorators.

Attributes:

Name Type Description
_func

Original Python function.

_func_type

'orchestration' | 'host' | 'incore' | 'inline' | 'opaque'. 'host' is the HOST-level orchestrator produced by @pl.jit.host — it owns pld.alloc_window_buffer / pld.window / pld.world_size() and the per-rank device= dispatch loop. End-to-end runtime dispatch works when the caller supplies config=RunConfig(distributed_config=...): the config is forwarded through _compileir.compile() (see RunConfig.compile_kwargs), which yields a DistributedCompiledProgram that __call__ dispatches per-rank.

_level

pl.Level or None.

_auto_scope

Whether the compiler auto-inserts AUTO runtime scopes (SIMPLER_SCOPE) around the body and each for/if body. True by default; set False via @pl.jit(auto_scope=False) / @pl.jit.host(auto_scope=False) to place scopes by hand with with pl.scope(). Also accepted on @pl.jit.inline(auto_scope=False) — after the InlineFunctions pass splices the body, hand-placed scopes land in the caller. incore / opaque kinds reject it (they outline into separate kernels, so scopes never land in the caller).

_dep_graph_state _CachedDepGraph | None

Last resolved graph and its validating bindings, published together. Each call pins its own validated graph.

_cache dict[CacheKey, Any]

L1 in-memory cache: CacheKey → CompiledProgram (post-pass ir.Program wrapped).

compile(*args, **kwargs)

Specialize + compile for the shape/dtype combination implied by args, and return the underlying CompiledProgram.

Same specialization / cache pipeline as __call__, minus the on-device dispatch. Use this when you want to drive execution through the runtime worker API directly:

  • pypto.runtime.ChipWorker.run / register for explicit L2 dispatch.
  • CompiledProgram.chip_callable / runtime_name / runtime_config to drive a hand-constructed simpler.worker.Worker.

config=RunConfig(...) is still consumed (and its compile-side knobs forwarded to ir.compile()) so the returned CompiledProgram honours the same options as a direct kernel(*args, config=...) call. Runtime-side fields on the RunConfig (device_id, DFX flags, ...) do not apply here — they affect dispatch, not the compiled artefact.

Subsequent calls (either __call__ or compile) with the same specialization and compatible cache policy return the same CompiledProgram instance. With persistence enabled, a disk-restored object has program is None; disable persistence when IR is required. Dump, compile-profiling, explicit output, and custom pass-diagnostic requests always compile afresh, preserving ordinary cached entries. Omitting config uses RunConfig defaults, including disabled dumps.

Compiling without tensors. When called with no tensor arguments — neither positional nor keyword — the shape/dtype contract is read directly from the kernel's own parameter annotations, so no throwaway torch.empty(...) dummies are needed. (Passing tensors by keyword, e.g. compile(a=sample_a), still binds them normally.) This requires every tensor parameter to carry a full pl.Tensor[[...], dtype] annotation (a bare pl.Tensor has no shape to read and raises). Dynamic dims (pl.dynamic / bind_dynamic) need no value — the artifact is extent-independent. Scalar parameters have no value in the signature, so pass them as keyword args (or via a signature default); a literal specializes that value into the artifact, while pl.RUNTIME leaves the parameter unspecialized — it stays a real pl.Scalar parameter supplied at dispatch and, like a dynamic dim, drops out of the cache key. A signature-mode call shares a cache entry with an equivalent compile(*sample_tensors) call whenever the two agree on every specialized scalar; pl.RUNTIME is its own specialization and is rejected on the sample-argument path, which always specializes the scalar value it is handed.

Example::

M = pl.dynamic("M")

@pl.jit
def my_kernel(
    x: pl.Tensor[[M, 4096], pl.BF16],
    w: pl.Tensor[[4096, 4096], pl.BF16],
    out: pl.Out[pl.Tensor[[M, 4096], pl.BF16]],
    num_tokens: pl.Scalar[pl.INT32],
):
    ...

worker = ChipWorker(config=RunConfig(platform="a2a3"))

# From sample tensors (shape/dtype read; contents ignored):
compiled = my_kernel.compile(sample_x, sample_w, sample_out, 128)

# Or straight from the (fully-annotated) signature — no tensors.
# num_tokens varies per launch, so keep it out of the artifact: its
# value is supplied on each dispatch through the compiled artifact
# (below), not by calling my_kernel(...) directly — an eager call
# re-specializes and compiles a separate artifact.
compiled = my_kernel.compile(num_tokens=pl.RUNTIME)

w_dev = worker.alloc_tensor(real_w.shape, real_w.dtype, init=real_w)
h = worker.register(compiled)
for batch in stream:
    h(batch.x, w_dev, batch.out, batch.num_tokens)

Parameters:

Name Type Description Default
*args Any

Positional arguments matching the decorated function's params. Tensor values are inspected for shape/dtype only; their contents are not read. Omit all positional args to compile straight from the signature annotations instead.

()
**kwargs Any

Keyword arguments. A config keyword, if present, is a RunConfig. In signature mode, scalar parameter values are also passed here (by name) — a literal to specialize it, or pl.RUNTIME to leave it unspecialized.

{}

Returns:

Type Description
Any

The cached CompiledProgram for this specialization.

warmup(*args, **kwargs)

Compile and prepare all device binaries without executing the kernel.

Uses the same arguments, configuration, specialization, and in-process cache as :meth:compile. Fully annotated tensors need no sample allocation; scalar defaults, keyword values, and pl.RUNTIME follow the same rules as annotation-driven compilation.

Unlike :meth:compile, this also assembles the kernel and orchestration binaries for every chip-level build before returning. It creates no runtime worker, initializes no NPU, and executes no kernel. The build host still needs the target compiler, SDK, and runtime dependencies. This method does not enable automatic persistent caching.

Parameters:

Name Type Description Default
*args Any

Optional sample arguments accepted by :meth:compile. Tensor contents are not read.

()
**kwargs Any

Kernel arguments and an optional config=RunConfig(...). Omit tensor arguments to use the complete tensor annotations.

{}

Returns:

Type Description
Any

The same CompiledProgram or DistributedCompiledProgram

Any

selected by :meth:compile, with all device binaries prepared.

Any

Execution remains a separate operation on the returned object.

Raises:

Type Description
TypeError

The compiled result does not support device-free warmup.

RuntimeError

A required binary cannot be prepared. Compiler and configuration errors propagate; a later warmup can retry.

specialize(*args, **kwargs)

Specialize this JIT function and return its pre-pass IR.

The step lower takes before it runs the pass pipeline: entry and every transitive dep are specialized into @pl.program source and parsed, and the parsed program is returned untransformed. No passes, no code generation, no ptoas, no device, and the compiled-program cache is neither read nor written.

Use this to hand a JIT kernel to a consumer that wants to drive the pass pipeline itself — most notably ir.compile(program, output_dir=...), which runs passes and code generation and so must be given the program before any pass has touched it. Passing lower()'s result there would run the pipeline a second time.

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

Parameters:

Name Type Description Default
*args Any

Positional sample arguments matching the decorated function. Omit tensor samples to specialize from the annotations, which then must carry full shapes.

()
**kwargs Any

Keyword sample arguments. Unlike lower() there is no config: the pass pipeline never runs here, so a RunConfig would have nothing to configure.

{}

Returns:

Type Description
Program

The parsed ir.Program, before any pass has run.

Raises:

Type Description
TypeError

config= was passed. Accepting it silently would let a caller believe a strategy or diagnostics setting shaped the returned IR, when no pass ran to read it.

param_names property

Declared parameter names, in signature order (self excluded).

output_param_names property

Parameters the kernel writes — pl.Out[...] and pl.InOut[...].

In declaration order, so the tuple stays aligned with the callee's return order. A caller that materialises tensors for this kernel uses it to decide which of them are results to validate.

lower(*args, **kwargs)

Specialize this JIT function and return its post-pass IR.

A config=RunConfig(...) keyword controls pass execution through its strategy, diagnostics, dependency-analysis, memory-planner, and platform fields. Runtime and artifact fields are ignored. This method does not run code generation, invoke ptoas, execute on a device, write artifacts, or access the compiled-program cache.

Parameters:

Name Type Description Default
*args Any

Positional sample arguments matching the decorated function. Omit tensor samples to specialize from fully shaped annotations.

()
**kwargs Any

Keyword sample arguments and an optional config. In signature mode a scalar parameter takes a literal (specialized into the IR) or pl.RUNTIME (left unspecialized).

{}

Returns:

Type Description
Program

The specialized ir.Program after configured passes.

function(func=None, *, type=ir.FunctionType.Opaque, level=None, role=None, attrs=None, auto_scope=True, strict_ssa=False, external_source=None)

function(func: Callable[..., Any], *, type: ir.FunctionType = ir.FunctionType.Opaque, level: ir.Level | None = None, role: ir.Role | None = None, attrs: dict[str, Any] | None = None, auto_scope: bool = True, strict_ssa: bool = False, external_source: str | Path | None = None) -> ir.Function
function(func: None = None, *, type: ir.FunctionType = ir.FunctionType.Opaque, level: ir.Level | None = None, role: ir.Role | None = None, attrs: dict[str, Any] | None = None, auto_scope: bool = True, strict_ssa: bool = False, external_source: str | Path | None = None) -> FunctionDecorator

Decorator that parses a DSL function and returns IR Function.

This decorator analyzes the decorated function's AST, parses the DSL constructs (type annotations, pl.range, pl.yield_, etc.), and builds an IR Function object.

Parameters:

Name Type Description Default
func Callable[..., _R_co] | None

Python function decorated with @pl.function

None
type FunctionType

Function type (Opaque, Orchestration, or InCore)

Opaque
level Level | None

Hierarchy level (e.g. pl.Level.HOST)

None
role Role | None

Function role (e.g. pl.Role.SubWorker)

None
attrs dict[str, Any] | None

Function-level attributes dict (e.g. {"split": pl.SplitMode.UP_DOWN})

None
auto_scope bool

If True (default), the compiler inserts AUTO runtime scopes (SIMPLER_SCOPE) around the function body and each for/if body. Set False to place scopes by hand with with pl.scope() (only meaningful for Orchestration functions).

True
strict_ssa bool

If True, enforce SSA (single assignment per variable). If False (default), allow variable reassignment (non-SSA mode).

False
external_source str | Path | None

Path to a hand-written C++ kernel .cpp that backs this function. Only valid on FunctionType.AIC / FunctionType.AIV functions, whose body must be a bare ... (signature only). The orchestration calls the kernel as usual, but the compiler skips PyPTO codegen for it and compiles the referenced source as the InCore kernel instead. Relative paths resolve against the defining file's directory.

None

Returns:

Type Description
DSLFunction[_R_co] | Callable[[Callable[..., _R_co]], DSLFunction[_R_co]]

IR Function object (or decorator if used with parameters)

Example

@pl.function ... def my_func(x: pl.Tensor[[64, 128], pl.FP16]) -> pl.Tensor[[64, 128], pl.FP32]: ... result = pl.create_tensor([64, 128], dtype=pl.FP32) ... return result @pl.function(level=pl.Level.HOST, role=pl.Role.SubWorker) ... def sub_worker(x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]: ... return x

inline(func)

Decorator that captures a function for inlining at call sites.

Unlike @pl.function which parses to an ir.Function immediately, @pl.inline defers parsing until the function is called within a @pl.program. The body is expanded in-place at each call site.

Parameters:

Name Type Description Default
func Callable[..., Any]

Python function to capture for inlining

required

Returns:

Type Description
InlineFunction

InlineFunction object with captured AST and metadata

Example

@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

program(cls=None, *, strict_ssa=False)

program(cls: type) -> ir.Program
program(cls: None = None, *, strict_ssa: bool = False) -> ProgramDecorator

Decorator that parses a class with @pl.function methods into a Program.

The class should contain one or more methods decorated with @pl.function. Each method is parsed as a separate function and added to the program. Methods must have 'self' as the first parameter (standard Python syntax), which is automatically stripped from the IR.

Parameters:

Name Type Description Default
cls type[Any] | None

Class with @pl.function decorated methods

None
strict_ssa bool

If True, enforce SSA (single assignment per variable). If False (default), allow variable reassignment (non-SSA mode).

False

Returns:

Type Description
Program | Callable[[type[Any]], Program]

IR Program object (or decorator if used with parameters)

Example

@pl.program ... class MyProgram: ... @pl.function ... def add(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]: ... result: pl.Tensor[[64], pl.FP32] = pl.add(x, 1.0) ... return result ... ... @pl.function ... def mul(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]: ... result: pl.Tensor[[64], pl.FP32] = pl.mul(x, 2.0) ... return result

MyProgram is now an ir.Program object

InlineFunction dataclass

Stores AST and metadata for a function to be inlined at call sites.

name instance-attribute

func_def instance-attribute

param_names instance-attribute

source_file instance-attribute

source_lines instance-attribute

line_offset instance-attribute

col_offset instance-attribute

closure_vars instance-attribute

parse(code, filename='<string>', source_map=None)

Parse a DSL function or program from a string.

This function takes Python source code containing a @pl.function decorated function or @pl.program decorated class and parses it into an IR Function or Program object. The code is executed dynamically, automatically importing pypto.language as pl if not already present.

Parameters:

Name Type Description Default
code str

Python source code containing @pl.function or @pl.program

required
filename str

Optional filename for error reporting (default: "")

'<string>'
source_map dict[int, tuple[str, int, int]] | None

Optional generated_line → (orig_file, orig_line, orig_col) map. When provided, spans whose emitted line is present are remapped to the original source location, so diagnostics point at the user's real file rather than this (possibly generated) code. Used by @pl.jit to recover provenance through its specialize→reparse round-trip (issue #1612).

None

Returns:

Type Description
Function | Program

Parsed ir.Function or ir.Program object (auto-detected)

Raises:

Type Description
ValueError

If the code contains nothing to parse or multiple items

ParserError

If parsing fails (syntax errors, type errors, etc.)

Warning

This function uses exec() to execute the provided code string. It should only be used with trusted input, as executing untrusted code can lead to arbitrary code execution vulnerabilities.

Examples:

>>> import pypto.language as pl
>>> # Parse a function
>>> func_code = '''
... @pl.function
... def add(x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
...     result: pl.Tensor[[64], pl.FP32] = pl.add(x, 1.0)
...     return result
... '''
>>> func = pl.parse(func_code)
>>> print(func.name)
add
>>> # Parse a program
>>> prog_code = '''
... @pl.program
... class MyProgram:
...     @pl.function
...     def add(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
...         return pl.add(x, 1.0)
... '''
>>> prog = pl.parse(prog_code)
>>> print(prog.name)
MyProgram

loads(filepath)

Load a DSL function or program from a file.

This function reads a Python file containing a @pl.function decorated function or @pl.program decorated class and parses it into an IR Function or Program object (auto-detected).

Parameters:

Name Type Description Default
filepath str

Path to Python file containing @pl.function or @pl.program

required

Returns:

Type Description
Function | Program

Parsed ir.Function or ir.Program object (auto-detected)

Raises:

Type Description
FileNotFoundError

If the file does not exist

ValueError

If the file contains nothing to parse or multiple items

ParserError

If parsing fails (syntax errors, type errors, etc.)

Warning

This function reads a file and executes its contents. It should only be used with trusted files, as executing code from untrusted sources can lead to arbitrary code execution vulnerabilities.

Examples:

>>> import pypto.language as pl
>>> # Load a function
>>> func = pl.loads('my_kernel.py')
>>> print(func.name)
>>> # Load a program
>>> prog = pl.loads('my_program.py')
>>> print(prog.name)

parse_program(code, filename='<string>')

Parse a DSL program from a string.

Deprecated

Use parse instead, which auto-detects functions and programs.

This is now an alias for parse that validates the result is a Program.

Parameters:

Name Type Description Default
code str

Python source code containing a @pl.program decorated class

required
filename str

Optional filename for error reporting (default: "")

'<string>'

Returns:

Type Description
Program

Parsed ir.Program object

Raises:

Type Description
ValueError

If the code contains a function instead of a program

ParserError

If parsing fails (syntax errors, type errors, etc.)

loads_program(filepath)

Load a DSL program from a file.

Deprecated

Use loads instead, which auto-detects functions and programs.

This is now an alias for loads that validates the result is a Program.

Parameters:

Name Type Description Default
filepath str

Path to Python file containing @pl.program decorated class

required

Returns:

Type Description
Program

Parsed ir.Program object

Raises:

Type Description
FileNotFoundError

If the file does not exist

ValueError

If the file contains a function instead of a program

ParserError

If parsing fails (syntax errors, type errors, etc.)

Tensor

Tensor type for PyPTO Language DSL.

This class serves dual purposes: 1. Type annotation helper for function signatures 2. Runtime wrapper around IR Expr/Call objects

Annotation mode (used in type hints): x: pl.Tensor[[64, 128], pl.FP16] y: pl.Tensor[[64, 128], pl.FP16, pl.NZ]

Runtime mode (wraps IR expressions): tensor = pl.create_tensor([64, 128], dtype=pl.FP32) # Returns Tensor wrapping the Call expression

Examples:

>>> import pypto.language as pl
>>>
>>> @pl.function
... def my_func(x: pl.Tensor[[64, 128], pl.FP16, pl.NZ]) -> pl.Tensor[[64, 128], pl.FP32]:
...     result: pl.Tensor[[64, 128], pl.FP32] = pl.create_tensor([64, 128], dtype=pl.FP32)
...     return result

shape = shape instance-attribute

dtype = dtype instance-attribute

layout = layout instance-attribute

memref = memref instance-attribute

unwrap()

Get underlying IR expression.

Returns:

Type Description
Expr

The wrapped Expr/Call object

Raises:

Type Description
ValueError

If called on an annotation-only Tensor

bind_dynamic(dim, var)

Mark a tensor dimension as runtime-dynamic for @pl.jit specialization.

This is a no-op at runtime. The @pl.jit specializer reads this call statically from the AST to determine which dimensions should be represented as DynVar nodes (ir.Var) in the generated type annotation rather than as compile-time constants.

Use bind_dynamic with bare pl.Tensor parameters, where the annotation carries no shape to hold the DynVar. When a parameter is already annotated with an explicit shape, prefer placing the pl.dynamic() variable directly in the annotation instead — that form matches @pl.program and needs no bind_dynamic call::

M = pl.dynamic("M")

@pl.jit
def kernel(
    a: pl.Tensor[[M, 128], pl.FP32],          # dim 0 dynamic via annotation
    c: pl.Out[pl.Tensor[[M, 128], pl.FP32]],  # shares the same DynVar
):
    K = pl.tensor.dim(a, 1)                    # dim 1 stays constant (128)
    ...

Both forms are honoured and may be combined; the dynamic dimensions are the union of the two sources.

Parameters:

Name Type Description Default
dim int

Zero-based dimension index to mark as dynamic.

required
var Any

The DynVar object (created with pl.dynamic()) to bind.

required

Example::

@pl.jit
def kernel(a: pl.Tensor, c: pl.Out[pl.Tensor]):
    M = pl.dynamic("M")
    a.bind_dynamic(0, M)   # dim 0 of a is runtime-dynamic
    c.bind_dynamic(0, M)   # dim 0 of c shares the same DynVar
    K = a.shape[1]         # dim 1 is compile-time constant
    ...

Tile

Tile type for PyPTO Language DSL.

Tile represents a tile in unified buffer (UB) memory. It is used for tile-level programming with operations like load, store, add, mul, etc.

Annotation mode (used in type hints): x: pl.Tile[[64, 64], pl.FP32]

Runtime mode (wraps IR expressions): tile = pl.load(tensor, [0, 0], [64, 64]) # Returns Tile wrapping the Call expression

Examples:

>>> import pypto.language as pl
>>>
>>> @pl.function
... def my_func(input: pl.Tensor[[64, 64], pl.FP32]) -> pl.Tensor[[64, 64], pl.FP32]:
...     tile: pl.Tile[[64, 64], pl.FP32] = pl.load(input, [0, 0], [64, 64])
...     result: pl.Tile[[64, 64], pl.FP32] = pl.add(tile, tile)
...     return pl.store(result, [0, 0], input)

shape = shape instance-attribute

dtype = dtype instance-attribute

memref = memref instance-attribute

memory_space = memory_space instance-attribute

tile_view = tile_view instance-attribute

unwrap()

Get underlying IR expression.

Returns:

Type Description
Expr

The wrapped Expr/Call object

Raises:

Type Description
ValueError

If called on an annotation-only Tile

Scalar

Scalar type for PyPTO Language DSL.

This class serves dual purposes: 1. Type annotation helper for function signatures 2. Runtime wrapper around IR Expr/Call objects

Annotation mode (used in type hints): x: pl.Scalar[pl.FP32] count: pl.Scalar[pl.INT32]

Runtime mode (wraps IR expressions): scalar_value = pl.scalar.create(3.14, dtype=pl.FP32) # Returns Scalar wrapping the Call expression

Examples:

>>> import pypto.language as pl
>>>
>>> @pl.function
... def add_scalar(
...     x: pl.Tensor[[64], pl.FP32],
...     scalar: pl.Scalar[pl.FP32]
... ) -> pl.Tensor[[64], pl.FP32]:
...     result: pl.Tensor[[64], pl.FP32] = pl.add(x, scalar)
...     return result

dtype = dtype instance-attribute

expr = None instance-attribute

unwrap()

Unwrap to get the underlying IR expression.

Returns:

Type Description
Expr

The wrapped IR expression

Raises:

Type Description
RuntimeError

If this is an annotation-only instance

Array

On-core array wrapper.

Annotation mode (used in type hints — rare; Arrays don't cross function boundaries in v1, but the annotation is available for clarity and forward-compat)::

arr: pl.Array[16, pl.INT32]

Runtime mode (the common path)::

arr = pl.array.create(16, pl.INT32)
arr[i] = value       # -> array.update_element (functional, rebinds arr)
x = arr[i]           # -> array.get_element

extent = extent instance-attribute

dtype = dtype instance-attribute

unwrap()

Tuple

Tuple type for PyPTO Language DSL.

Used exclusively as a type annotation helper in function signatures. The parser reads the AST (not the runtime value), so this class only needs to make pl.Tuple[T1, T2, ...] evaluate without error.

Annotation syntax

result: pl.Tuple[pl.Tensor[[64], pl.FP32], pl.Scalar[pl.INT32]]

Examples:

>>> import pypto.language as pl
>>>
>>> @pl.function
... def multi_return(
...     x: pl.Tensor[[64], pl.FP32]
... ) -> pl.Tuple[pl.Tensor[[64], pl.FP32], pl.Scalar[pl.INT32]]:
...     ...

DynVar

Bases: Scalar

Dynamic shape variable for use in type annotations.

Creates a symbolic dimension that becomes an ir.Var node in the IR shape. Inherits from Scalar so that DynVar is accepted wherever Scalar/IntLike is expected (e.g. shape parameters, TensorView valid_shape).

Example

M = pl.dynamic("M") N = pl.dynamic("N")

@pl.function def func(a: pl.Tensor[[M, N], pl.FP32]) -> ...: ...

name = name instance-attribute

dtype = DataType.INDEX instance-attribute

expr = None instance-attribute

unwrap()

Return the underlying ir.Var, creating it eagerly if needed.

This allows DynVar to participate in _normalize_intlike() and other Scalar-consuming paths without requiring prior TypeResolver resolution.

InOut

Bases: _DirectionWrapper

Wrapper for InOut parameter direction in type annotations.

Usage::

def kernel(output: pl.InOut[pl.Tensor[[16, 128], pl.FP32]]) -> ...:

IntLike = int | Scalar | Expr module-attribute

Type alias for shape/offset parameters that accept int literals, Scalar DSL values, or raw Expr.

Out

Bases: _DirectionWrapper

Wrapper for Out parameter direction in type annotations.

Usage::

def kernel(result: pl.Out[pl.Tensor[[16, 128], pl.FP32]]) -> ...:

RUNTIME = RuntimeScalarMarker() module-attribute

Singleton RuntimeScalarMarker — see the class docstring.

const(value, dtype)

Create a typed constant with an explicit dtype.

Used by the printer to preserve non-default constant dtypes in round-trip. The parser intercepts pl.const() calls and creates ConstInt/ConstFloat with the specified dtype.

Statically typed -> Scalar to mirror its IR semantics: pl.const builds a ConstInt/ConstFloat expression, so it can be returned from a -> pl.Scalar function and combined with other scalars. At runtime the stub returns the numeric value unchanged (cast is a no-op).

Parameters:

Name Type Description Default
value int | float

Numeric value (int or float)

required
dtype Any

DataType for the constant

required

Returns:

Type Description
Scalar

Parser builds ConstInt/ConstFloat IR; the runtime stub returns

Scalar

the numeric value unchanged (type checking sees a Scalar).

range(*args, init_values=None)

range(*args: RangeArg, init_values: None = None) -> RangeIterator[Scalar]
range(*args: RangeArg, init_values: tuple[T1]) -> RangeIterator[tuple[Scalar, tuple[T1]]]
range(*args: RangeArg, init_values: tuple[T1, T2]) -> RangeIterator[tuple[Scalar, tuple[T1, T2]]]
range(*args: RangeArg, init_values: tuple[T1, T2, T3]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3]]]
range(*args: RangeArg, init_values: tuple[T1, T2, T3, T4]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4]]]
range(*args: RangeArg, init_values: tuple[T1, T2, T3, T4, T5]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4, T5]]]

Create a range iterator for for loops.

Supports several patterns

Simple: for i in pl.range(10): Iter args: for i, (var1, var2) in pl.range(16, init_values=(init1, init2)):

For software pipelining (body replication for ping-pong buffering), use pl.pipeline(N, stage=F) instead — it is a sibling loop iterator.

Args can be int literals or Scalar variables

for i in pl.range(n): # n is pl.Scalar[pl.INT64] for i in pl.range(0, n, 1): for i in pl.range(n * 2 + 1):

Parameters:

Name Type Description Default
*args RangeArg

Positional arguments (stop) or (start, stop) or (start, stop, step). Each argument can be an int literal or a pl.Scalar value.

()
init_values tuple[Any, ...] | None

Initial values for iteration arguments

None

Returns:

Type Description
RangeIterator[Scalar] | RangeIterator[tuple[Scalar, tuple[Any, ...]]]

If no init_values: RangeIterator yielding loop variable (Scalar)

RangeIterator[Scalar] | RangeIterator[tuple[Scalar, tuple[Any, ...]]]

If init_values: RangeIterator yielding (loop_var, (iter_args...))

parallel(*args, init_values=None)

parallel(*args: RangeArg, init_values: None = None) -> RangeIterator[Scalar]
parallel(*args: RangeArg, init_values: tuple[T1]) -> RangeIterator[tuple[Scalar, tuple[T1]]]
parallel(*args: RangeArg, init_values: tuple[T1, T2]) -> RangeIterator[tuple[Scalar, tuple[T1, T2]]]
parallel(*args: RangeArg, init_values: tuple[T1, T2, T3]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3]]]
parallel(*args: RangeArg, init_values: tuple[T1, T2, T3, T4]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4]]]
parallel(*args: RangeArg, init_values: tuple[T1, T2, T3, T4, T5]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4, T5]]]

Create a parallel range iterator for parallel for loops.

Behaves identically to range() at runtime. The distinction is used by the parser to emit ForKind.Parallel instead of ForKind.Sequential.

Parameters:

Name Type Description Default
*args RangeArg

Positional arguments (stop) or (start, stop) or (start, stop, step). Each argument can be an int literal or a pl.Scalar value.

()
init_values tuple[Any, ...] | None

Initial values for iteration arguments

None

Returns:

Type Description
RangeIterator[Scalar] | RangeIterator[tuple[Scalar, tuple[Any, ...]]]

If no init_values: RangeIterator yielding loop variable (Scalar)

RangeIterator[Scalar] | RangeIterator[tuple[Scalar, tuple[Any, ...]]]

If init_values: RangeIterator yielding (loop_var, (iter_args...))

unroll(*args)

Create an unroll range iterator for compile-time loop unrolling.

Behaves identically to range() at runtime. The distinction is used by the parser to emit ForKind.Unroll instead of ForKind.Sequential.

Unrolled loops do not support init_values (loop-carried state).

Parameters:

Name Type Description Default
*args RangeArg

Positional arguments (stop) or (start, stop) or (start, stop, step). Each argument must be an int literal (compile-time constant).

()

Returns:

Type Description
RangeIterator[Scalar]

RangeIterator yielding loop variable (Scalar)

Examples:

>>> for i in pl.unroll(4):
...     x = pl.add(x, 1.0)
>>> for i in pl.unroll(0, 6, 2):
...     x = pl.add(x, i)

pipeline(*args, stage, init_values=None)

pipeline(*args: RangeArg, stage: int, init_values: None = None) -> RangeIterator[Scalar]
pipeline(*args: RangeArg, stage: int, init_values: tuple[T1]) -> RangeIterator[tuple[Scalar, tuple[T1]]]
pipeline(*args: RangeArg, stage: int, init_values: tuple[T1, T2]) -> RangeIterator[tuple[Scalar, tuple[T1, T2]]]
pipeline(*args: RangeArg, stage: int, init_values: tuple[T1, T2, T3]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3]]]
pipeline(*args: RangeArg, stage: int, init_values: tuple[T1, T2, T3, T4]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4]]]
pipeline(*args: RangeArg, stage: int, init_values: tuple[T1, T2, T3, T4, T5]) -> RangeIterator[tuple[Scalar, tuple[T1, T2, T3, T4, T5]]]

Create a software-pipelined loop iterator.

Replicates the loop body stage times per outer iteration to enable ping-pong buffering. The outer loop advances in strides of stage * step; a tail dispatch covers the remainder when the trip count is not divisible by stage. Lowered by the LowerPipelineLoops pass at the tile level.

Positional args match pl.range: (stop) / (start, stop) / (start, stop, step). The stage kwarg is required and must be a positive integer.

Parameters:

Name Type Description Default
*args RangeArg

1-3 positional args — same shape as pl.range().

()
stage int

Pipeline depth (positive integer, typically 2-4).

required
init_values tuple[Any, ...] | None

Loop-carried state, same semantics as pl.range().

None

Examples:

>>> for i in pl.pipeline(64, stage=4):
...     tile = pl.tile.load(...)
...     pl.tile.store(...)

while_(*, init_values=None)

while_(*, init_values: tuple[T1]) -> WhileIterator[tuple[T1]]
while_(*, init_values: tuple[T1, T2]) -> WhileIterator[tuple[T1, T2]]
while_(*, init_values: tuple[T1, T2, T3]) -> WhileIterator[tuple[T1, T2, T3]]
while_(*, init_values: tuple[T1, T2, T3, T4]) -> WhileIterator[tuple[T1, T2, T3, T4]]
while_(*, init_values: tuple[T1, T2, T3, T4, T5]) -> WhileIterator[tuple[T1, T2, T3, T4, T5]]

Create a while iterator for while loops.

Always requires init_values to specify loop-carried state. The loop condition must be specified as the first statement in the loop body using pl.cond().

Pattern

for (var1, var2) in pl.while_(init_values=(init1, init2)): pl.cond(condition) # loop body var1_out, var2_out = pl.yield_(var1_updated, var2_updated)

Parameters:

Name Type Description Default
init_values tuple[ExprType, ...] | None

Initial values for iteration arguments (required)

None

Returns:

Type Description
WhileIterator[tuple[ExprType, ...]]

WhileIterator yielding tuple of iter_args

Raises:

Type Description
ValueError

If init_values is not provided

Examples:

>>> for (x,) in pl.while_(init_values=(0,)):
...     pl.cond(x < 10)
...     x = x + 1
...     x_out = pl.yield_(x)
>>>
>>> for (x, y) in pl.while_(init_values=(0, 1)):
...     pl.cond(x < n)
...     x_new = x + 1
...     y_new = y * 2
...     x_out, y_out = pl.yield_(x_new, y_new)

yield_(*values)

yield_(value: T1) -> T1
yield_(v1: T1, v2: T2) -> tuple[T1, T2]
yield_(v1: T1, v2: T2, v3: T3) -> tuple[T1, T2, T3]
yield_(v1: T1, v2: T2, v3: T3, v4: T4) -> tuple[T1, T2, T3, T4]
yield_(v1: T1, v2: T2, v3: T3, v4: T4, v5: T5) -> tuple[T1, T2, T3, T4, T5]

Yield values from a scope (for, if).

This function is used to explicitly return values from nested scopes and create SSA phi nodes.

Parameters:

Name Type Description Default
*values Any

Values to yield

()

Returns:

Type Description
Any | tuple[Any, ...]

The yielded value(s). For single value, returns the value.

Any | tuple[Any, ...]

For multiple values, returns tuple.

Examples:

>>> # Single value yield
>>> result = pl.yield_(x + 1)
>>>
>>> # Multiple value yield
>>> a, b = pl.yield_(x, y)

cond(condition)

Specify the condition for a pl.while_() loop.

This function must be the first statement in a pl.while_() loop body. It is purely syntactic - the parser extracts the condition and sets it on the WhileStmt.

Parameters:

Name Type Description Default
condition CondArg

While loop condition (bool literal or Scalar variable)

required

Examples:

>>> for (x,) in pl.while_(init_values=(0,)):
...     pl.cond(x < 10)
...     x = x + 1
...     x_out = pl.yield_(x)

static_print(*args)

Print compile-time information about IR objects.

At parse time, prints type/value info to stdout. At runtime, no-op.

Parameters:

Name Type Description Default
*args Any

Values to print (variables, expressions, string labels)

()

static_assert(condition, msg='')

Assert a condition at compile time (parse time).

At parse time, evaluates condition. If false, raises ParserError. At runtime, this is a no-op (all semantics are handled by the parser).

Notes
  • This is a statement-only construct. It must be used as a standalone statement, not as part of an expression.
  • The msg argument must be a string literal at the call site. Passing a variable or expression for msg will raise ParserSyntaxError.
  • The check is evaluated at parse time only; it does not run at execution time.

Parameters:

Name Type Description Default
condition Any

Condition to check (must be compile-time evaluable)

required
msg str

Optional error message as a string literal

''

func_attr(attrs)

Attach function-level attributes from inside the function body.

Written as the first statement of a function body, pl.func_attr declares metadata about the function as a whole::

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

Why the body and not the decorator: a decorator is evaluated before the signature binds any name, so @pl.function(attrs={"stationary": w}) cannot be written — 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 at all.

Semantics:

  • Prologue only. Every pl.func_attr call must precede every other statement in the body. An attribute describes the whole function, so it must not appear to start applying partway down a body; pinning it to the prologue also keeps the printed form deterministic. Consequently only parameters — not body-locals — are referenceable.
  • A bare name is always a parameter reference. pl.func_attr({"n": k}) records the parameter k, never the value of an enclosing Python variable named k. Write Python-level constants as literals.
  • Multiple calls merge. A key repeated across two calls — or between pl.func_attr and a decorator attrs= — is a ParserSyntaxError naming the key.
  • Consumed at parse time. The dict lands in Function.attrs and emits no IR statement of its own.

Attributes the parser must read before it can parse the body stay on the decorator as dedicated keywords: auto_scope= and external_source=.

Parameters:

Name Type Description Default
attrs dict[str, Any]

Attribute dict. Keys must be string literals.

required

at(level, role=None, *, optimizations=None, deps=None, no_dep_args=None, dumps=None, allow_early_resolve=False, predicate=None, name_hint='', windowize=False)

Mark a region of code for execution at a specific hierarchy level.

With level=pl.Level.CORE_GROUP, the optimizations= list controls the resulting scope kind:

  • no entries → ScopeStmt(InCore)
  • pl.split(mode)ScopeStmt(InCore, split=mode)

For all other levels, this creates a Hierarchy scope.

Parameters:

Name Type Description Default
level Level

Target hierarchy level (e.g. pl.Level.HOST, pl.Level.CORE_GROUP).

required
role Role | None

Function role (Orchestrator or Worker). Default: None.

None
optimizations list[Optimization] | None

Optional list literal of optimization entries. Each entry must be pl.split(mode) — written inline at the call site, since the DSL parser inspects the AST and does not accept dynamically built variables here.

None
deps list[Any] | None

Optional explicit producer-edge list (TaskId Vars and/or None sentinels). Lowered to the resulting Call's manual_dep_edges attr, which codegen packs into a set_dependencies(...) invocation. Operates at the TaskId level — orthogonal to no_dep_args=, which operates at the arg-slot level on captured tensors.

None
no_dep_args list[Any] | None

Optional list literal of outer-scope tensor names captured by the scope body. Each entry must be a bare tensor name; the parser resolves it to a Var, the outliner translates the Var list into positional indices into the synthesised Call's args, and DeriveCallDirections overwrites those slots to ArgDirection.NoDep. Equivalent to wrapping the same tensor with pl.no_dep(t) at an explicit kernel call site — use this form when the kernel call is synthesised by the pl.at outliner and there is no syntactic call-arg slot to wrap. Legal for both read-only captures and captures that the scope body mutates via pl.assemble / pl.store (the outliner classifies the latter as InOut on the synthesised kernel, and NoDep overrides InOut just as it overrides Input); the user is asserting that sibling fan-outs touch disjoint regions of the tensor and therefore do not need OverlapMap dep tracking. Note: deps= takes TaskIds, while no_dep_args= takes tensors — they describe different things.

None
dumps list[Any] | None

Optional list literal of outer-scope tensor names to mark for selective tensor dump on the synthesised kernel dispatch. The scope-level selective-dump surface, symmetric with deps=: the outliner translates the captured-tensor entries into the dispatch's dump_vars by Var identity. Equivalent to declaring the same tensors with pl.dump_tag(t) before the scope; use dumps= when you want the dump targets listed explicitly at the scope.

None
allow_early_resolve bool

Opt the outlined dispatch in as a speculative early-dispatch producer (simpler#1065). Same hint as pl.submit(..., allow_early_resolve=True): the scheduler may pre-stage this task's consumers before it completes. Forces the scope to lower to an ir.Submit (even without as tid) so the flag can ride to codegen, where it emits Arg::set_allow_early_resolve(true). Pure scheduling hint.

False
predicate Any

Optional dispatch predicate — a single comparison of one tensor element against an integer literal, e.g. predicate=(row_count[e] > 0). Same surface, validation and lowering as pl.submit(..., predicate=...) / pl.spmd(..., predicate=...): the scheduler evaluates it at the dispatch point (after this scope's dependencies are satisfied, so the value is current) and retires the task inline — never dispatching it to a core — when it is false, while still settling fanin/fanout so downstream consumers unlock. The comparison is parsed but never evaluated in orchestration; reading it there would stall on wait_for_tensor_ready, exactly what the predicate avoids. Like allow_early_resolve it forces the scope to lower to an ir.Submit (even without as tid).

Only valid with level=pl.Level.CORE_GROUP and only on a scope that is not nested inside pl.cluster() / pl.spmd() / another pl.at(level=pl.Level.CORE_GROUP). Every other placement either never lowers to a Submit (Hierarchy scopes are not outlined into a task dispatch) or is folded into an enclosing wrapper dispatch, so the predicate would be silently dropped; the parser rejects those instead.

Contract: the operand tensor's producing task must be one of deps= — otherwise the predicate may read a stale value. Unlike pl.spmd, deps= is available on every pl.at form. The parser makes a best-effort check (see pl.submit); getting deps= right remains the author's responsibility.

None
name_hint str

Optional name hint for the outlined function (must be a valid identifier).

''
windowize bool

Explicitly allow local windowization for the outlined InCore kernel. The default is False.

False

Returns:

Type Description
AtContext

Context manager for the appropriate scope.

Examples:

>>> # InCore scope:
>>> with pl.at(level=pl.Level.CORE_GROUP):
...     y = pl.ops.add(x, x)
>>> # InCore scope with split hint:
>>> with pl.at(level=pl.Level.CORE_GROUP,
...            optimizations=[pl.split(pl.SplitMode.UP_DOWN)]):
...     y = pl.ops.add(x, x)
>>> # Hierarchy scope (unchanged behavior):
>>> with pl.at(level=pl.Level.HOST, role=pl.Role.SubWorker):
...     y = pl.add(x, x)
>>> # Conditionally dispatched InCore scope — the scheduler reads
>>> # rc[0, 0] at the dispatch point and skips the task when it is 0:
>>> with pl.at(level=pl.Level.CORE_GROUP) as gate_tid:
...     rc = pl.store(pl.load(rc, [0, 0], [128, 128]), [0, 0], rc)
>>> with pl.at(level=pl.Level.CORE_GROUP,
...            deps=[gate_tid], predicate=(rc[0, 0] > 0)) as tid:
...     out = pl.store(pl.load(x, [0, 0], [128, 128]), [0, 0], out)

cluster(*, name_hint='')

Mark a region of code as belonging to a Cluster execution context.

A cluster groups co-scheduled AIC (Cube) and AIV (Vector) kernels that share the same physical cluster resources. The OutlineClusterScopes pass extracts Cluster scopes into separate Group-typed functions.

Parameters:

Name Type Description Default
name_hint str

Optional name hint for the outlined function.

''

Returns:

Type Description
ClusterContext

Context manager for Cluster scope

Examples:

>>> with pl.cluster():
...     with pl.at(level=pl.Level.CORE_GROUP):
...         y = pl.add(x, x)

graph(name)

Mark a repeated region of orchestration as one recordable graph.

Under runtime="host_build_graph" the runtime records the region's task topology the first time it executes and replays the recording on every later execution. A decoder stack of N structurally identical layers therefore costs one graph build instead of N, and occupies N outer task slots instead of N x (tasks per layer).

pl.graph is the in-place form of @pl.jit.graph: OutlineGraphScopes extracts the region into a Graph function named after name, so the two surfaces compile to the same thing. Reach for the scope form when the region is a slice of a larger orchestration body that you would rather not split into a separate function; reach for the decorator when the layer is already its own function.

The recorded topology is fixed after the first execution: only tensor addresses and boundary scalars are refreshed per replay. LegalizeGraphBoundary proves that statically and rejects at compile time whatever it cannot prove — it never silently degrades to a wrong answer.

Parameters:

Name Type Description Default
name str

Region name. Must be a valid Python identifier; it becomes the outlined function's name and hence the runtime's graph key, so keep it stable across edits.

required

Returns:

Type Description
GraphContext

Context manager for the Graph scope.

Raises:

Type Description
ParserSyntaxError

if written in a device-kernel body (InCore / AIC / AIV / Group / Spmd), nested inside pl.at / pl.cluster / pl.spmd, or nested inside another pl.graph.

Examples:

>>> for layer in pl.range(40):
...     with pl.graph("decoder_layer"):
...         h = attention(x, wq, h)
...         x = mlp(h, w1, x)

spmd(core_num, *, sync_start=False, name_hint='', optimizations=None, deps=None, allow_early_resolve=False, predicate=None)

Dispatch a kernel with SPMD (Single Program Multiple Data) multi-block execution.

The first argument is the number of blocks and is positional — mirroring range(n). Loop start is fixed at 0 and step at 1; each block gets an index i in [0, core_num).

Three usage forms at a glance:

Form Description
with pl.spmd(n): Dispatch or inline block (no captured TaskId).
for i in pl.spmd(n): Loop-style; i = per-block index, body is auto-outlined to InCore.
with pl.spmd(n) as tid: Same body as form 1, plus the dispatch TaskId in tid.

deps=[...] works on all three: capturing the TaskId is only needed when a later task must wait on this one.

Usage forms:

  1. with pl.spmd(n): — body is either a dispatch body calling a pre-defined InCore kernel, or an inline block auto-outlined into a synthetic InCore kernel (like the loop form, minus the auto-bound index). Which one is decided semantically, not by statement count: a body reading the per-block index via pl.tile.get_block_idx() is inline; otherwise it is a dispatch body, however many statements it holds — though it may launch only one kernel (the lowering stops at the first call). A body that neither reads the index nor dispatches a kernel is rejected — every block would run identical work. An explicit with pl.at(<CORE_GROUP level>, ...): as the sole body statement is the InCore carrier and is not wrapped again; with such a body optimizations= must go on that pl.at(...) rather than on the pl.spmd(...) line. Captures no producer TaskId (use form 3 for that), but still accepts deps=. Can stand alone (implicit cluster) or nest inside pl.cluster() — but only a scope carrying no Submit-only metadata may nest: deps= / allow_early_resolve= / predicate= require the standalone form (a cluster-nested pl.spmd is unwrapped into the Group function and never produces the Submit that carries them).

  2. for i in pl.spmd(n): — loop-style. The iteration variable binds the per-block index (equivalent to pl.tile.get_block_idx()); the body is auto-outlined into a synthetic InCore function, so inline tile/tensor ops work without a separate @pl.function(type=InCore) declaration.

  3. with pl.spmd(n, deps=[...]) as tid: — same body shapes as form 1, and additionally captures the grid dispatch's producer Scalar[TASK_ID] in tid (mirroring with pl.at(...) as tid:), usable as a deps= edge on later tasks, stored into a pl.array.create(N, pl.TASK_ID), or crossing into pl.manual_scope. TaskId capture is the only thing this adds over form 1 — it is orthogonal to both the inline body and deps=.

Optional optimizations=[pl.split(mode)] applies to the inner InCore scope (auto-generated for the for-form and the as tid form, wrapped around the call for the plain with-form).

Parameters:

Name Type Description Default
core_num RangeArg

Number of blocks for SPMD dispatch. Positional; accepts a Python int or any ir.Expr of integer type. Closure-captured integer constants and closure arithmetic are folded to ConstInt by the parser and Simplify; non-foldable expressions flow through to codegen unchanged. Pass pl.system.available_cluster_count() (mixed / cube-only kernel) or pl.system.available_aiv_count() (vector-only kernel) to size the launch on the run's own device geometry — the only spelling that stays at full occupancy across devices, which a hard pl.system.syncall requires. Pass such a query inline as shown; binding it to a name first compiles and lowers correctly, but the printed IR of the outlined Spmd wrapper then references a variable defined in the caller and cannot be re-parsed.

required
sync_start bool

If True, all blocks start execution simultaneously (default: False).

False
name_hint str

Optional name hint for the outlined function.

''
optimizations list[Optimization] | None

Optional list literal containing only pl.split(mode) entries (the parser inspects the AST).

None
deps list[Any] | None

Optional explicit producer-edge list (TaskId Vars and/or None sentinels), accepted on all three forms. Lowered to the outlined Submit's manual_dep_edges; codegen packs it into a set_dependencies(...) invocation (union'd with auto-deps). Like allow_early_resolve it forces the dispatch to lower to an ir.Submit even without as tid — the Spmd outliner synthesises the (unused) producer TaskId Var that shape needs. Rejected on a pl.cluster()-nested pl.spmd, which is unwrapped into the Group function and never produces a Submit.

None
allow_early_resolve bool

Opt the grid dispatch in as a speculative early-dispatch producer (simpler#1065). Same hint as pl.submit(..., allow_early_resolve=True) / pl.at(..., allow_early_resolve=True): the scheduler may pre-stage this task's consumers before it completes. Forces the dispatch to lower to an ir.Submit (even without as tid) so the flag rides to codegen, where it emits Arg::set_allow_early_resolve(true). Pure scheduling hint. Rejected on a pl.cluster()-nested pl.spmd (such a scope is unwrapped into the Group function and never produces a Submit, so the hint would be lost).

False
predicate Any

Optional dispatch predicate — a single comparison of one tensor element against an integer literal, e.g. predicate=(row_count[e] > 0). Same surface, validation and lowering as pl.spmd_submit(..., predicate=...): the scheduler evaluates it at the dispatch point (after this scope's dependencies are satisfied, so the value is current) and retires the task inline — never dispatching it to a core — when it is false, while still settling fanin/fanout so downstream consumers unlock. The comparison is parsed but never evaluated in orchestration; reading it there would stall on wait_for_tensor_ready, exactly what the predicate avoids. Like allow_early_resolve it forces the dispatch to lower to an ir.Submit (even without as tid) and is rejected on a pl.cluster()-nested pl.spmd.

Contract: the operand tensor's producing task must be one of deps= — otherwise the predicate may read a stale value. The parser makes a best-effort check (see pl.spmd_submit); getting deps= right remains the author's responsibility.

None

Returns:

Type Description
SpmdContext

Context manager / loop iterator for the SPMD scope.

Examples:

>>> # Single-kernel context-manager form (direct dispatch)
>>> with pl.spmd(4):
...     out = self.kernel(a, b, out)
>>>
>>> # Inline context-manager form — no TaskId; read the block index yourself
>>> with pl.spmd(4):
...     i = pl.tile.get_block_idx()
...     offset = i * 128
...     tile_a = pl.load(a, [offset, 0], [128, 128])
...     tile_b = pl.load(b, [offset, 0], [128, 128])
...     out = pl.store(pl.add(tile_a, tile_b), [offset, 0], out)
>>>
>>> # Loop form — body runs per-block with i = tile.get_block_idx()
>>> for i in pl.spmd(4):
...     offset = i * 128
...     tile_a = pl.load(a, [offset, 0], [128, 128])
...     tile_b = pl.load(b, [offset, 0], [128, 128])
...     out = pl.store(pl.add(tile_a, tile_b), [offset, 0], out)
>>>
>>> # Capture-form: inline body + producer TaskId for explicit dep wiring
>>> with pl.spmd(4, name_hint="stage1") as tid_a:
...     i = pl.tile.get_block_idx()
...     offset = i * 128
...     out = pl.store(pl.add(a[offset], b[offset]), [offset, 0], out)
>>> with pl.spmd(4, name_hint="stage2", deps=[tid_a]) as tid_b:
...     i = pl.tile.get_block_idx()
...     out2 = pl.store(pl.relu(out[i * 128]), [i * 128, 0], out2)
>>>
>>> # With-form with split hint on the inner InCore wrapper
>>> with pl.spmd(4, optimizations=[pl.split(pl.SplitMode.UP_DOWN)]):
...     out = self.kernel(a, b, out)
>>>
>>> # SPMD inside cluster (mixed kernel)
>>> with pl.cluster():
...     with pl.spmd(4, sync_start=True):
...         out = self.kernel(a, b, out)
>>>
>>> # Device-sized launch — required for a hard pl.system.syncall
>>> with pl.spmd(pl.system.available_cluster_count(), sync_start=True):
...     out = self.mixed_kernel(a, b, out)
>>>
>>> # Dispatch predicate — skip the expert entirely when its row count is 0
>>> with pl.spmd(1) as gate_tid:
...     row_count = self.gate(row_count)
>>> with pl.spmd(4, deps=[gate_tid], predicate=(row_count[0, 0] > 0)) as tid:
...     out = self.expert(x, out)

split_aiv(n, *, mode)

Open an explicit AIV-split region as an SPMD-style loop.

Usage::

for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
    ...  # body runs per AIV lane; aiv_id = pl.tile.get_subblock_idx()

The loop builds a first-class SplitAivScopeStmt region carrying the requested SplitMode. Because it is a structural node (not a whole-InCore-scope flag), it is nestable: the region may appear inside a pl.range / pl.pipeline loop or an if, and sibling regions may carry different modes (multi-mode), each lowered independently. The loop variable binds the AIV lane index (equivalent to pl.tile.get_subblock_idx()).

A top-level for aiv_id in pl.split_aiv(...) is wrapped in an enclosing InCore scope so OutlineIncoreScopes can outline it — unless the region is already in a core context (an open pl.at(level=CORE_GROUP) scope, or a function declared pl.FunctionType.InCore), where it is emitted in place. A region is CORE_GROUP-level, so authoring one directly in an InCore function is rejected by the AivSplitValid verifier; write it in a plain @pl.function / @pl.jit (Opaque) body instead.

Two dispatch modes:

  • Data-parallel (UP_DOWN / LEFT_RIGHT): the region's vector compute is halved on the split axis (rows / cols), so each lane processes one half of every tile.
  • Task-parallel (NONE): no halving — both lanes run the full body for disjoint work the author dispatches via aiv_id (e.g. an aiv_id-strided loop). Use this when the tiles cannot be halved (unit dims) or a reduction must stay full-width. pl.aiv_shard / pl.aic_gather are valid here and preserve the shape: with no split axis they carry the one meaning that still applies — this value crosses the AIC/AIV boundary.

Manual mode. Opening even ONE region changes the contract for the whole enclosing function: the regions become authoritative for where vector work runs, and the AivSplitValid verifier enforces the division:

========================== ================= ========================== op inside a region outside every region ========================== ================= ========================== vector compute AIV rejected pl.load / pl.store AIV allowed (compiler-inserted) cube compute (matmul) rejected AIC pl.aiv_shard / pl.aic_gather the boundary rejected pld.system.notify pinned to AIV duplicated onto BOTH lanes ========================== ================= ==========================

So write one region per vector phase, with cube ops and barriers between them. A phase that must stay full width goes in its own for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE): — whose meaning is exactly "both lanes run the full body", which is what an un-regioned vector phase in such a function already did. A function with NO region is unaffected and keeps its previous behaviour.

Name every tile that crosses a region edge. Manual mode hands the AIC/AIV boundary to the author, so the compiler stops placing it: a cube-produced value read on the vector lane inside a region must arrive through pl.aiv_shard, and a value defined in a region and read on the cube lane outside it must leave through pl.aic_gather. Both crossings lower fine without the op — which is the point: an unnamed boundary is one nobody chose, so the verifier rejects it and names the value and the reader::

mm = pl.matmul(q, k)                     # cube, outside every region
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    v = pl.exp(pl.aiv_shard(mm))         # C->V: named
    kv = pl.aic_gather(v)                # V->C: named
out = pl.matmul(kv, w)                   # cube again, outside

Every crossing in one function must agree on split-vs-no-split. All the pl.aiv_shard / pl.aic_gather calls in a function ride ONE logical cross-core pipe -- one initialize_pipe per side, both directions on the same pipe -- and pto-isa carries no-split as a parameter of that pipe's TYPE (TPipe<..., IsNoSplit, ...>), selecting a different handshake protocol. A pipe is therefore split or un-split for its whole lifetime, so a NONE region that crosses the boundary cannot sit beside a data-parallel one that also crosses it::

for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    a = pl.exp(pl.aiv_shard(mm0))        # crossing, no split
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
    b = pl.exp(pl.aiv_shard(mm1))        # crossing, split   -> rejected

Two different split axes are fine -- the axis is a per-transfer choice, only split-vs-no-split belongs to the pipe -- so UP_DOWN beside LEFT_RIGHT is accepted. A region carrying no crossing is free to use any mode: the NONE region that only pins a pld.system.notify to the vector lane never touches the pipe. When two phases genuinely need different transports, put them in separate pl.at(level=pl.Level.CORE_GROUP) scopes -- each outlines into its own function, and so gets its own pipe. Diagnosed by AivSplitValid at OutlineIncoreScopes; without it the program reaches ptoas, which rejects it at pto.initialize_l2g2l_pipe.

Gather only a lane-uniform value out of a NONE region. The ISA requires both AIV sub-lanes to take part in a no-split handshake and they share one destination slot with no per-lane offset, and nothing arbitrates between them: both lanes push, so if they hold different values the cube receives an unspecified one of the two. Guarding the production of the value does not help — lane 1 still reaches the push and still sends its own tile. So a gather out of a NONE region is well-defined only for a lane-uniform value; if the lanes must contribute different data, use a data-parallel region instead. Not diagnosed.

Put cross-rank comm ops in a region. A region also decides placement for ops that have no lane of their own. pld.system.notify is core-agnostic by ISA, so in a kernel that mixes cube and vector work it is emitted on BOTH the AIC and the AIV lane — and the cube copy can publish the signal before the vector lane's put has landed the data it releases. Putting the op in a region pins it to the vector lane::

for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    pld.tensor.put(dst=win, peer=peer, src=out, ...)
    pld.system.notify(target=sig, peer=peer, offsets=[0, 0], value=1,
                      op=pld.NotifyOp.AtomicAdd)

A region is not "run once" — shard once-only side effects by aiv_id. A mode=NONE body runs on BOTH AIV sub-lanes, so the snippet above sends TWO notifies to the same peer. AtomicAdd accumulates, so that peer's counter reads 2 for one arrival, a pld.system.wait on it is released early, and the rank races ahead on data that has not landed. Shard the op, or guard it to one lane::

# sharded: each lane takes different peers
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    for owner in pl.range(aiv_id, NUM_PEERS, 2):
        pld.system.notify(target=sig, peer=owner, ...)

# guarded: lane 0 only
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    if aiv_id == 0:
        pld.system.notify(target=sig, peer=peer, ...)

This is NOT diagnosed. The compiler guarantees only that a region's comm ops stay off the cube lane; the sharded and un-sharded forms are the same single statement in the AIV body, differing only in whether aiv_id reaches the call's arguments.

GM traffic is outside all of this. The crossing rules above govern tile values. A GM tensor belongs to no lane — pld.tensor.put takes one by signature — so no boundary op can express a crossing through it, and AIC and AIV run asynchronously. ExpandMixedKernel fences one narrow shape by itself: a cube tile.store whose GM tensor a vector tile.load reads back, when the two share an origin and the load is in or under the store's body. Everything else — a comm op reading the buffer, a consumer in a sibling body — needs explicit producer-side cache publication and a GM fence, a cross-core pl.system.syncall, and consumer-side invalidation. That sequence remains the author's responsibility.

The region survives parse -> SSA -> ResolveBackendOpLayouts as a structural node (printer emits for aiv_id in pl.split_aiv(...): so parse->print->parse is a fixpoint), then is lowered by LowerAutoVectorSplit and erased by ExpandMixedKernel; it never reaches codegen.

Parameters:

Name Type Description Default
n int

The AIV sub-core count. Positional; hardware-fixed at 2 (the two AIV lanes of one AICore). Any other value is rejected by the parser.

required
mode SplitMode

Required dispatch mode, no silent default. pl.SplitMode.NONE (task-parallel, no halving), pl.SplitMode.UP_DOWN or pl.SplitMode.LEFT_RIGHT (data-parallel halving).

required

Returns:

Type Description
SplitAivContext

Loop iterator for the explicit AIV-split region.

Examples:

>>> for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
...     offset = aiv_id * 128
...     t = pl.load(a, [offset, 0], [128, 128])
...     out = pl.store(t, [offset, 0], out)

A second, full-width vector phase after a cube op needs its own region (manual mode) — mode=NONE keeps it un-halved:

>>> for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
...     ...  # phase 1, dispatched per lane via aiv_id
>>> mm = pl.matmul(q, k)  # cube work stays outside every region
>>> for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
...     out = pl.add(mm, bias)  # phase 2, full width on both lanes

split(mode, *, slot_num=None)

Create a Split optimization entry.

Parameters:

Name Type Description Default
mode SplitMode

Split mode. May be SplitMode.NONE, SplitMode.UP_DOWN, or SplitMode.LEFT_RIGHT.

required
slot_num int | None

Deprecated — use pl.cross_core_slot(slot_num=N). Must be positive when set. Emits a DeprecationWarning.

None

Returns:

Type Description
Split

Split instance for use in pl.at(..., optimizations=[...]).

Raises:

Type Description
ValueError

If slot_num is set but not positive.

cross_core_slot(*, slot_num)

Create a CrossCoreSlot optimization entry.

Parameters:

Name Type Description Default
slot_num int

Cross-core pipe slot count (ring depth). Must be positive.

required

Returns:

Type Description
CrossCoreSlot

CrossCoreSlot instance for use in pl.at(..., optimizations=[...]).

Raises:

Type Description
ValueError

If slot_num is not a positive integer.

Examples:

>>> # Deepen the auto-inserted ring so the producing core can run further ahead
>>> with pl.at(level=pl.Level.CORE_GROUP,
...            optimizations=[pl.cross_core_slot(slot_num=4)]):
...     ...

add(lhs, rhs)

add(lhs: Tensor, rhs: Tensor | int | float | Scalar) -> Tensor
add(lhs: Tile, rhs: Tile | int | float | Scalar) -> Tile
add(lhs: Scalar, rhs: Scalar | int | float) -> Scalar

Element-wise addition, dispatched by input type.

sub(lhs, rhs)

sub(lhs: Tensor, rhs: Tensor | int | float | Scalar) -> Tensor
sub(lhs: Tile, rhs: Tile | int | float | Scalar) -> Tile
sub(lhs: Scalar, rhs: Scalar | int | float) -> Scalar

Element-wise subtraction, dispatched by input type.

mul(lhs, rhs)

mul(lhs: Tensor, rhs: Tensor | int | float | Scalar) -> Tensor
mul(lhs: Tile, rhs: Tile | int | float | Scalar) -> Tile
mul(lhs: Scalar, rhs: Scalar | int | float) -> Scalar

Element-wise multiplication, dispatched by input type.

div(lhs, rhs, high_precision=False)

div(lhs: Tensor, rhs: Tensor | int | float | Scalar | _ir_core.Expr, high_precision: bool = False) -> Tensor
div(lhs: Tile, rhs: Tile | int | float | Scalar | _ir_core.Expr, high_precision: bool = False) -> Tile
div(lhs: Scalar, rhs: Scalar | int | float, high_precision: bool = False) -> Scalar

Element-wise division, dispatched by input type.

A scalar rhs against a Tile dispatches to tile.divs.

Parameters:

Name Type Description Default
high_precision bool

Select PTOAS's high-precision divide. Available for Tensor/Tensor and Tile/Tile only -- a scalar divisor has no high-precision form, so passing it there raises rather than silently falling back.

False

part_add(lhs, rhs)

part_add(lhs: Tensor, rhs: Tensor) -> Tensor
part_add(lhs: Tile, rhs: Tile) -> Tile

Partial element-wise add, dispatched by input type.

part_mul(lhs, rhs)

part_mul(lhs: Tensor, rhs: Tensor) -> Tensor
part_mul(lhs: Tile, rhs: Tile) -> Tile

Partial element-wise multiply, dispatched by input type.

part_max(lhs, rhs)

part_max(lhs: Tensor, rhs: Tensor) -> Tensor
part_max(lhs: Tile, rhs: Tile) -> Tile

Partial element-wise max, dispatched by input type.

part_min(lhs, rhs)

part_min(lhs: Tensor, rhs: Tensor) -> Tensor
part_min(lhs: Tile, rhs: Tile) -> Tile

Partial element-wise min, dispatched by input type.

maximum(lhs, rhs)

maximum(lhs: Tensor, rhs: Tensor | int | float | Scalar) -> Tensor
maximum(lhs: Tile, rhs: Tile | int | float | Scalar) -> Tile

Element-wise maximum, dispatched by input type.

exp(input)

Element-wise exponential, dispatched by input type.

log(input, high_precision=False)

Element-wise natural logarithm, dispatched by input type.

Parameters:

Name Type Description Default
input T

Input tensor or tile.

required
high_precision bool

Select PTOAS's high-precision logarithm mode.

False

cast(input, target_type, mode='round', *, saturation_mode=None)

cast(input: Tensor, target_type: int | DataType, mode: str | int = 'round', *, saturation_mode: str | int | None = None) -> Tensor
cast(input: Tile, target_type: int | DataType, mode: str | int = 'round', *, saturation_mode: str | int | None = None) -> Tile
cast(input: Scalar, target_type: int | DataType, mode: str | int = 'round', *, saturation_mode: str | int | None = None) -> Scalar

Type casting, dispatched by input type.

Parameters:

Name Type Description Default
input Tensor | Tile | Scalar

Value to convert.

required
target_type int | DataType

Destination dtype.

required
mode str | int

Rounding mode, as a name or its int code -- "none" (0), "rint" (1), "round" (2, the default), "floor" (3), "ceil" (4), "trunc" (5), "odd" (6). A Scalar input supports the default only and raises for any other mode.

'round'
saturation_mode str | int | None

Destination saturation for a Tensor or Tile input, as a name or its int code -- "off" (0) or "on" (1). "on" clamps a rounded value that falls outside the destination range to that range; "off" keeps the target's non-saturating conversion, including its overflow and non-finite behavior. Defaults to "on" for an integer destination: nothing standard fixes what an overflowing conversion to an integer produces, clamping is the safer of the two to get by accident, and it is what the hardware converts natively. A float destination keeps the target's own IEEE behavior (an out-of-range narrowing yields an infinity) unless you ask otherwise. When the cast lowers to a chain of native conversions the mode applies to the final hop. A Scalar input does not support this option, so passing it one is an error rather than a silent no-op.

None
Example

quantized = pl.cast(rounded_fp16, pl.INT8, mode="trunc", saturation_mode="on")

concat(src0, src1)

Column-wise concatenation, dispatched by input type.

reshape(input, shape)

Reshape operation, dispatched by input type.

A reshape is a zero-copy view, so it never widens the valid region: the result holds real data in exactly the cells the input did, re-expressed in shape. Because a valid region is an origin-anchored box, not every input region survives a repartition — reshaping a region that no box of shape can describe is rejected rather than silently rounded up to fully valid. Reshapes that only add or drop fully-valid unit axes always work, as does a region occupying a contiguous prefix of the buffer.

reinterpret_view(data, dtype, *, shape=None)

Reinterpret the same bytes with a different dtype.

Parameters:

Name Type Description Default
data T

Input tensor or tile.

required
dtype DataType

Target element dtype, which must differ from the source dtype.

required
shape Sequence[IntLike] | None

Optional byte-equivalent target shape. When omitted, the physically contiguous dimension is scaled according to the source/target dtype byte ratio.

None

Returns:

Type Description
T

A zero-copy view of the same kind as data.

transpose(input, axis1, axis2)

Transpose operation, dispatched by input type.

Parameters:

Name Type Description Default
input T

Value to transpose.

required
axis1 int

First axis to exchange. Must be a compile-time constant; negative indexing is supported.

required
axis2 int

Second axis to exchange. Must differ from axis1 after negative indexing is resolved -- naming the same axis twice is rejected, not treated as a no-op.

required

slice(input, shape, offset, valid_shape=None, drop_dims=None, pad_value=None, clamp=False)

Slice operation, dispatched by input type.

The slice is never valid where the source is not: the source's valid region, shifted by offset and cut to the window, bounds the result.

drop_dims lists axes to erase from the result type (numpy-style rank reduction); each must be a static unit dim of shape that is still fully valid after that intersection. None / [] drops nothing.

pad_value sets the padding mode for elements outside the effective valid region, on either path. None carries the source's mode through. Accepts PadValue.zero / PadValue.max / PadValue.min, or the literal sugars 0, math.inf, -math.inf (same spelling as fillpad). It only bites when the valid region is smaller than shape — which an explicit valid_shape, a partially-valid source, or (Tensor-only) clamp=True can each bring about; passing it otherwise warns.

clamp sanctions a window that runs off the end of the source: by default the slice asserts offset + shape stays inside the source and is rejected when that provably fails, whereas clamp=True lets the window overhang and cuts the valid region back to the source edge. It is only available on a Tensor — an on-chip tile window has nothing that could clamp it.

matmul(lhs, rhs, out_dtype=None, a_trans=False, b_trans=False, c_matrix_nz=False)

matmul(lhs: Tensor, rhs: Tensor, out_dtype: int | DataType | None = ..., a_trans: bool = ..., b_trans: bool = ..., c_matrix_nz: bool = ...) -> Tensor
matmul(lhs: Tile, rhs: Tile, out_dtype: DataType | None = ..., a_trans: Literal[False] = ..., b_trans: Literal[False] = ..., c_matrix_nz: Literal[False] = ...) -> Tile

Matrix multiplication, dispatched by input type.

a_trans / b_trans / c_matrix_nz are Tensor-only: a tensor value carries no layout, so a flag is the only place the information can live. At tile level transposition is a type property, so passing any of them with a Tile operand raises rather than being dropped.

A transpose flag swaps its own operand's two trailing axes, so that operand must be at least 2D — a_trans with a 1D lhs (or b_trans with a 1D rhs) raises rather than being ignored. On the mixed mat-vec / vec-mat forms the flag applies to the matrix side: a lhs stored [K, M] with a_trans=True against a [K] rhs deduces [M].

out_dtype is likewise Tensor-only. tile.matmul's result dtype is fixed by the Cube accumulator (FP32 for float operands, INT32 for int), so the Tile path accepts out_dtype only when it already agrees with that deduction and raises otherwise.

The Tensor path allows one conversion on top of that deduction, because the result drains L0C through the FIXPIPE, whose unscaled writeback narrows FP32 -> FP16 / FP32 -> BF16. So float operands accept FP32, FP16 or BF16, and int operands accept only INT32. Every rejected pair is a scale-bearing conversion the FIXPIPE can only do with quantization parameters this call has nowhere to carry — INT32 -> FP32 is a dequantization, FP32 -> INT8 a quantization, INT32 -> INT8 a requantization. Requesting one (notably out_dtype=FP32 on INT8 operands) raises here rather than lowering to a backend type error. Convert explicitly with pl.cast instead.

For Tensor inputs with rank > 2 on either operand, the call is lowered to tile.batch_matmul (with batch broadcasting) by ConvertTensorToTileOps and then unrolled to per-batch tile.matmul by FlattenTileNdTo2D. Use this entry point (rather than pl.batch_matmul) for tensor-level ND matmul.

batch_matmul(lhs, rhs)

Tile-only batched matrix multiplication.

Tensor batched matmul is handled by pl.matmul / pl.tensor.matmul: when any operand has rank > 2, ConvertTensorToTileOps automatically dispatches to tile.batch_matmul (and FlattenTileNdTo2D later unrolls it). Use this op only when you are working at the tile level.

row_max(input, tmp_tile=None)

row_max(input: Tensor, tmp_tile: None = ...) -> Tensor
row_max(input: Tile, tmp_tile: Tile) -> Tile

Row-wise max reduction, dispatched by input type.

For Tile inputs, tmp_tile is required and must have the same dtype and rank as the input, with every dimension at least as large as the input dimension. Tensor inputs must omit it — the scratch tile is allocated during Tensor-to-Tile lowering — and passing one raises.

row_sum(input, tmp_tile=None)

row_sum(input: Tensor, tmp_tile: None = ...) -> Tensor
row_sum(input: Tile, tmp_tile: Tile) -> Tile

Row-wise sum reduction, dispatched by input type.

For Tile inputs, tmp_tile is required and must have the same dtype and rank as the input, with every dimension at least as large as the input dimension. Tensor inputs must omit it — the scratch tile is allocated during Tensor-to-Tile lowering — and passing one raises.

row_min(input, tmp_tile=None)

row_min(input: Tensor, tmp_tile: None = ...) -> Tensor
row_min(input: Tile, tmp_tile: Tile) -> Tile

Row-wise min reduction, dispatched by input type.

For Tile inputs, tmp_tile is required and must have the same dtype and rank as the input, with every dimension at least as large as the input dimension. Tensor inputs must omit it — the scratch tile is allocated during Tensor-to-Tile lowering — and passing one raises.

row_prod(input, tmp_tile=None)

row_prod(input: Tensor, tmp_tile: None = ...) -> Tensor
row_prod(input: Tile, tmp_tile: Tile) -> Tile

Row-wise product reduction, dispatched by input type.

For Tile inputs, tmp_tile is required and must have the same dtype and rank as the input, with every dimension at least as large as the input dimension. Tensor inputs must omit it — the scratch tile is allocated during Tensor-to-Tile lowering — and passing one raises.

col_sum(input, tmp_tile=None)

col_sum(input: Tensor, tmp_tile: None = ...) -> Tensor
col_sum(input: Tile, tmp_tile: Tile | None = ...) -> Tile

Column-wise sum reduction, dispatched by input type.

For Tile inputs, passing tmp_tile activates the binary-tree reduction path; omitting it uses the sequential path. Tensor inputs must omit it: the tensor-to-tile conversion always lowers to the sequential path and allocates its own scratch, so a tmp_tile there could not select the requested strategy and raises instead.

col_max(input)

Column-wise max reduction, dispatched by input type.

For Tensor inputs, the tensor-to-tile conversion lowers to tile.col_max.

col_min(input)

Column-wise min reduction, dispatched by input type.

For Tensor inputs, the tensor-to-tile conversion lowers to tile.col_min.

col_prod(input)

Column-wise product reduction, dispatched by input type.

For Tensor inputs, the tensor-to-tile conversion lowers to tile.col_prod.

row_argmax(input, tmp_tile=None)

row_argmax(input: Tensor, tmp_tile: None = ...) -> Tensor
row_argmax(input: Tile, tmp_tile: Tile) -> Tile

Row-wise argmax (per-row max index, int32), dispatched by input type.

For Tile inputs, tmp_tile is required with exactly the same shape and dtype. Tensor inputs must omit it — the conversion injects the scratch tile — and passing one raises.

row_argmin(input, tmp_tile=None)

row_argmin(input: Tensor, tmp_tile: None = ...) -> Tensor
row_argmin(input: Tile, tmp_tile: Tile) -> Tile

Row-wise argmin (per-row min index, int32), dispatched by input type.

For Tile inputs, tmp_tile is required with exactly the same shape and dtype. Tensor inputs must omit it — the conversion injects the scratch tile — and passing one raises.

col_argmax(input, tmp_tile=None)

col_argmax(input: Tensor, tmp_tile: None = ...) -> Tensor
col_argmax(input: Tile, tmp_tile: Tile) -> Tile

Column-wise argmax (per-column max index, int32), dispatched by input type.

For Tile inputs, tmp_tile is required (unlike col_max) and must have exactly the same shape and dtype as the input: the pto-isa kernel reads the column count from the tmp/src extent, so an oversized scratch walks past the valid columns. Tensor inputs must omit it — the conversion injects the tmp tile — and passing one raises.

col_argmin(input, tmp_tile=None)

col_argmin(input: Tensor, tmp_tile: None = ...) -> Tensor
col_argmin(input: Tile, tmp_tile: Tile) -> Tile

Column-wise argmin (per-column min index, int32), dispatched by input type.

For Tile inputs, tmp_tile is required (unlike col_min) and must have exactly the same shape and dtype as the input: the pto-isa kernel reads the column count from the tmp/src extent, so an oversized scratch walks past the valid columns. Tensor inputs must omit it — the conversion injects the tmp tile — and passing one raises.

row_expand(lhs, rhs)

Row-wise expansion, dispatched by input type.

row_expand_add(lhs, rhs, tmp=None)

row_expand_add(lhs: Tensor, rhs: Tensor) -> Tensor
row_expand_add(lhs: Tile, rhs: Tile, tmp: Tile | None = None) -> Tile

Row-wise broadcast addition; tmp is available only for Tile inputs.

row_expand_sub(lhs, rhs)

Row-wise broadcast subtraction, dispatched by input type.

row_expand_mul(lhs, rhs)

Row-wise broadcast multiplication, dispatched by input type.

row_expand_div(lhs, rhs)

Row-wise broadcast division, dispatched by input type.

row_expand_max(lhs, rhs)

Row-wise broadcast maximum, dispatched by input type.

row_expand_min(lhs, rhs)

Row-wise broadcast minimum, dispatched by input type.

row_expand_expdif(lhs, rhs)

Row-wise exp-diff (exp(lhs - rhs) with per-row scalar), dispatched by input type.

col_expand(lhs, rhs)

Column-wise expansion, dispatched by input type.

col_expand_mul(lhs, rhs)

Column-wise broadcast multiplication, dispatched by input type.

col_expand_div(lhs, rhs)

Column-wise broadcast division, dispatched by input type.

col_expand_sub(lhs, rhs)

Column-wise broadcast subtraction, dispatched by input type.

col_expand_add(lhs, rhs)

Column-wise broadcast addition, dispatched by input type.

col_expand_max(lhs, rhs)

Column-wise broadcast maximum, dispatched by input type.

col_expand_min(lhs, rhs)

Column-wise broadcast minimum, dispatched by input type.

col_expand_expdif(lhs, rhs)

Column-wise exp-diff (exp(lhs - rhs) with per-column scalar), dispatched by input type.

expand_clone(src, target)

Clone and expand input to target shape.

Parameters:

Name Type Description Default
src Tensor

Source tensor

required
target Tensor

Target tensor defining output shape

required

Returns:

Type Description
Tensor

Tensor wrapping the expand_clone operation

expands(target, scalar)

Expand scalar to target shape, dispatched by target type.

Note the argument order: the value being broadcast is the second argument. target supplies the shape and dtype; it is not read.

Parameters:

Name Type Description Default
target Tensor | Tile

Value whose shape the scalar is broadcast to.

required
scalar int | float | Scalar

Value to broadcast into every element.

required

neg(input)

Element-wise negation, dispatched by input type.

abs(input)

Element-wise absolute value, dispatched by input type.

recip(input, high_precision=False)

Element-wise reciprocal (1/x), dispatched by input type.

Parameters:

Name Type Description Default
input T

Input tensor or tile

required
high_precision bool

Whether to select PTOAS's high-precision reciprocal mode (FP16/FP32 only)

False

read(src, offset)

Read a scalar value at given indices, dispatched by source type.

Parameters:

Name Type Description Default
src Tensor | Tile

Source tensor (global memory) or tile (unified buffer)

required
offset IntLike | Sequence[IntLike]

A single index expression (for 1-D flat access) or index list (one per dimension) into the source

required

Returns:

Type Description
Scalar

Scalar wrapping the read value

write(dst, offset, value)

Write a scalar value to a tensor or tile at given indices.

Parameters:

Name Type Description Default
dst Tensor | Tile

Destination tensor (global memory) or tile (unified buffer)

required
offset IntLike | Sequence[IntLike]

A single index expression (for 1-D flat access) or index list (one per dimension) into the destination

required
value Scalar

Scalar value to write

required

Returns:

Type Description
Expr

Underlying tensor.write / tile.write call expression. Direct

Expr

callers ignore it; the DSL parser surfaces it as an EvalStmt.

create_tile = create module-attribute

fillpad(value, pad_value=PadValue.zero)

Fill invalid elements, dispatched by input type.

pad_value accepts the PadValue enum or the literal sugars 0, math.inf, -math.inf. Other values raise — the hardware only supports the three padding modes.

fillpad_expand(value, shape, pad_value=PadValue.zero)

Copy a smaller source into a larger destination, padding the rest.

Dispatched by input type. The destination shape may be larger than the source in either dimension; the source's valid region is copied into the top-left of the destination and every other element is filled with pad_value (PadValue enum or the literal sugars 0, math.inf, -math.inf).

load(tensor, offsets, shapes, valid_shape=None, target_memory=None, clamp=False, cache=None)

Copy data from tensor to unified buffer (tile).

Only the valid extent is read, so the tile may be larger than the region that exists in the source. The tile's valid region is the source's valid region, shifted by offsets and cut to the tile — a load never reports source bytes that do not exist as real data.

Parameters:

Name Type Description Default
tensor Tensor

Source tensor

required
offsets Sequence[IntLike]

Offsets in each dimension. Always in the source tensor's coordinate system. Each element must be an integer scalar — a float (n / 2 rather than n // 2) or a nested sequence is rejected.

required
shapes Sequence[IntLike]

Shape of the region to load in each dimension. Always in the source tensor's coordinate system. Each element must be an integer scalar, on the same terms as offsets.

required
valid_shape Sequence[IntLike] | None

Valid shape of the tile in each dimension. When provided, sets TileView.valid_shape in the output TileType. When omitted, shapes is used as valid_shape. Uses the same coordinate convention as shapes. Narrows the tile; cannot widen it past what the source has. Each element must be an integer scalar — one extent per dimension, not a nested [start, extent] pair.

None
target_memory MemorySpace | None

Target memory space (MemorySpace.Vec or MemorySpace.Mat). None (the default) leaves an ordinary load unset for the compiler to place. MX-layout tensors default to MemorySpace.Mat so they can be passed directly to matmul_mx and placed from its operand position.

None
clamp bool

Sanction a read that runs off the end of the source. By default a load asserts offsets + valid_shape stays inside the source and is rejected when that provably fails; clamp=True cuts the request back to the source edge instead.

False
cache CachePolicy | None

GM cache-access policy for this read. None (the default) states no policy, leaving any scope-level declaration to apply. CachePolicy.BYPASS declares a streaming read — it asserts the bytes have no reuse worth caching and that nothing writes them while the kernel runs; coherency is the author's contract (see pl.set_cache_policy for the full contract). An explicit value here always wins over a scope-level pl.set_cache_policy declaration for the same tensor, in both directions: cache=CachePolicy.DEFAULT opts this one read back into the cache inside a bypassing scope. PTOAS has no L2-bypass path yet (https://github.com/hw-native-sys/PTOAS/issues/1356), so a BYPASS request warns and compiles as an ordinary cached access today.

None

Returns:

Type Description
Tile

Tile wrapping the load operation

Example

2D load

tile = load(tensor, offsets=[0, 0], shapes=[32, 32])

streaming read, no cache reuse expected

tile = load(tensor, [0, 0], [32, 32], cache=pl.CachePolicy.BYPASS)

store(tile, offsets, output_tensor, shapes=None, *, atomic=AtomicType.None_, st_phase=STPhase.Unspecified)

Copy data from tile back to tensor.

Parameters:

Name Type Description Default
tile Tile

Source tile

required
offsets Sequence[IntLike]

Offsets in each dimension. Each element must be an integer scalar — a float or a nested sequence is rejected.

required
output_tensor _TensorT

Output tensor

required
shapes Sequence[IntLike] | None

Optional ND partition shape. Injected by FlattenTileNdTo2D for ND tensors. Each element must be an integer scalar, on the same terms as offsets.

None
atomic AtomicType

Combine mode for the global-memory write. AtomicType.None_ (default) overwrites; AtomicType.Add atomically adds the tile into existing GM contents — used for split-K accumulation, where several cores accumulate partial products into one output.

NOTE: atomic-add accumulation order across cores is not fixed, so floating-point results are non-deterministic. The destination must be zero-initialised before the kernel runs. Supported tile dtypes: fp32 / bf16 / fp16 / int32 / int16 / int8. bf16 atomic-add is available on the Ascend910B (A2/A3) profile; it is not supported on A5, where an fp32 accumulator + cast is required instead.

None_
st_phase STPhase

Consumer-side unit-flag phase. A producer that finishes with acc_phase=pl.AccPhase.Final must be consumed by a store with st_phase=pl.STPhase.Final so the unit flag is cleared.

Unspecified

Returns:

Type Description
_TensorT

Tensor wrapping the store operation

Example

2D store

result = store(tile, [0, 0], tensor)

3D store

result = store(tile, [0, 0, 0], tensor)

atomic-add store (split-K)

result = store(partial, [0, 0], out, atomic=pl.AtomicType.Add)

clear the unit flag after a final phased accumulation

result = store(acc, [0, 0], out, st_phase=pl.STPhase.Final)

move(tile, target_memory, blayout=None, slayout=None)

Move tile between memory levels.

Parameters:

Name Type Description Default
tile Tile

Input tile

required
target_memory MemorySpace

Target memory space (MemorySpace.Vec, .Mat, .Left, .Right, .LeftScale, .RightScale)

required
blayout TileLayout | None

Optional block layout for the destination tile

None
slayout TileLayout | None

Optional scatter layout for the destination tile

None

Returns:

Type Description
Tile

Tile wrapping the move operation

gatherb(src, offset, *, output_dtype=None)

Gather 32-byte blocks from src by UINT32 byte offsets.

Each offset selects one 32-byte source block. One offset column expands to 32 / sizeof(output_dtype) output elements. output_dtype defaults to src.dtype and may select another supported byte interpretation. A sliced source must have a byte address that PyPTO can prove is 32-byte aligned; dynamic column offsets are rejected conservatively.

Parameters:

Name Type Description Default
src Tile

Source tile to gather blocks from.

required
offset Tile

UINT32 tile of byte offsets into src -- not element indices.

required
output_dtype int | DataType | None

Byte interpretation of the result. Defaults to src.dtype.

None

Returns:

Type Description
Tile

Tile wrapping the gatherb operation.

mgather(mem, idx, coalesce='row', *, gather_oob='undefined', target_memory=MemorySpace.Vec, scratch=None, valid_shape=None)

mgather(mem: Tensor, idx: Tile, coalesce: str | int = ..., *, gather_oob: str | int = ..., target_memory: Literal[MemorySpace.Vec] = ..., scratch: None = ..., valid_shape: None = ...) -> Tile
mgather(mem: Tensor, idx: Tensor, coalesce: Literal['row', 0] = ..., *, gather_oob: str | int = ..., target_memory: Literal[MemorySpace.Mat], scratch: None = ..., valid_shape: Sequence[int] | None = ...) -> Tile
mgather(mem: Tensor, idx: Tensor, coalesce: Literal['elem', 1], *, gather_oob: str | int = ..., target_memory: Literal[MemorySpace.Mat], scratch: Tensor, valid_shape: Sequence[int] | None = ...) -> Tile

Gather-load rows or elements from a GM tensor into a fresh Vec or Mat tile.

Vec output uses a 2D INT32 index tile. Mat output uses a GM INT32 index tensor and produces canonical NZ layout; its element mode additionally requires a same-dtype GM scratch tensor.

Parameters:

Name Type Description Default
mem Tensor

Source tensor in GM.

required
idx Tile | Tensor

Two-dimensional INT32 index tile for Vec output, or GM tensor for Mat output.

required
coalesce str | int

"row"/0 for row gather or "elem"/1 for flat element gather. Integer values support printed-IR round trips.

'row'
gather_oob str | int

Out-of-bounds handling: "undefined", "clamp", "wrap", "zero", or the corresponding integer 0..3.

'undefined'
target_memory MemorySpace

MemorySpace.Vec (default) or MemorySpace.Mat. This selects the operator variant, not merely a placement: the two take a different idx type and produce a different output shape and view, so it cannot be left for the compiler to infer.

Vec
scratch Tensor | None

Same-dtype GM workspace required by Mat element gather and forbidden by the other forms.

None
valid_shape Sequence[int] | None

Optional two-dimensional written region for Mat output. Vec output derives its valid region from the index tile.

None

mscatter(src, idx, output_tensor)

Scatter-store tile elements into a tensor at per-element indices.

Semantics: output_tensor[idx[i, j]] = src[i, j]

Maps to the PTOAS pto.mscatter instruction.

Parameters:

Name Type Description Default
src Tile

Source tile (FP16, FP32, INT16, or INT32)

required
idx Tile

Index tile (INT32, same rank as src)

required
output_tensor _TensorT

Output tensor to scatter into (same dtype as src)

required

Returns:

Type Description
_TensorT

Tensor wrapping the mscatter operation

Example

result = pl.tile.mscatter(src_tile, idx_tile, out_tensor)

sqrt(input)

Element-wise square root, dispatched by input type.

rsqrt(input, high_precision=False)

rsqrt(input: Tensor, high_precision: bool = ...) -> Tensor
rsqrt(input: Tile, high_precision: Literal[False] = ...) -> Tile

Element-wise reciprocal square root, dispatched by input type.

high_precision is Tensor-only: the compiler allocates the scratch tile during Tensor-to-Tile lowering. tile.rsqrt carries no such attribute — precision is selected purely by passing that scratch tile — so tile callers use pl.tile.rsqrt(tile, tmp) directly and passing high_precision=True with a Tile raises rather than silently yielding the low-precision path.

relu(tile)

Element-wise ReLU activation (max(0, x)).

Parameters:

Name Type Description Default
tile Tile

Input tile

required

Returns:

Type Description
Tile

Tile wrapping the relu operation

matmul_acc(acc, lhs, rhs, a_trans=False, b_trans=False, init_cond=None)

matmul_acc(acc: Tensor, lhs: Tensor, rhs: Tensor, a_trans: bool = ..., b_trans: bool = ..., init_cond: BoolLike | None = ...) -> Tensor
matmul_acc(acc: Tile, lhs: Tile, rhs: Tile, a_trans: Literal[False] = ..., b_trans: Literal[False] = ..., init_cond: BoolLike | None = ...) -> Tile

Matrix multiplication with accumulation, dispatched by input type.

a_trans / b_trans are Tensor-only for the same reason as in matmul — at tile level transposition is a type property, not an op flag — so passing either with Tile operands raises rather than being dropped.

init_cond makes the accumulator's initial value conditional: on the steps where it holds, acc is overwritten with lhs @ rhs rather than accumulated into, which is the split-K k == 0 idiom. It applies to 2D operands only.

For Tensor inputs with rank > 2 on any of acc/lhs/rhs, the call is lowered to tile.batch_matmul_acc (with batch broadcasting on lhs/rhs vs the fixed acc batch) by ConvertTensorToTileOps and then unrolled to per-batch tile.matmul_acc by FlattenTileNdTo2D.

matmul_bias(lhs, rhs, bias)

Matrix multiplication with bias add: C = lhs @ rhs + bias.

Parameters:

Name Type Description Default
lhs Tile

Left-hand side tile [M, K]

required
rhs Tile

Right-hand side tile [K, N]

required
bias Tile

Bias tile [1, N] with the accumulator dtype (FP32 for floating-point matrix operands, INT32 for integer matrix operands)

required

Returns:

Type Description
Tile

Tile wrapping the matmul_bias operation

matmul_mx(lhs, lhs_scale, rhs, rhs_scale)

MX block-scale matrix multiplication.

Scales use logical [M, K/32] / [K/32, N] shapes. Both data tiles passed to this operation must be FP8E4M3FN. For the supported FP4 x FP8 input form, explicitly cast the FP4 lhs to FP8E4M3FN before calling this operation; native FP4 x FP4 is not supported.

Parameters:

Name Type Description Default
lhs Tile

Left-hand side data tile (FP8E4M3FN)

required
lhs_scale Tile

Left-hand side scale tile (FP8E8M0)

required
rhs Tile

Right-hand side data tile (FP8E4M3FN)

required
rhs_scale Tile

Right-hand side scale tile (FP8E8M0)

required

Returns:

Type Description
Tile

Tile wrapping the matmul_mx operation

matmul_mx_acc(acc, lhs, lhs_scale, rhs, rhs_scale)

MX block-scale matmul with accumulation.

Data operands follow matmul_mx: an FP4 lhs must first be cast to FP8E4M3FN, and the operation itself receives two FP8E4M3FN tiles.

Parameters:

Name Type Description Default
acc Tile

Accumulator tile

required
lhs Tile

Left-hand side data tile (FP8E4M3FN)

required
lhs_scale Tile

Left-hand side scale tile (FP8E8M0)

required
rhs Tile

Right-hand side data tile (FP8E4M3FN)

required
rhs_scale Tile

Right-hand side scale tile (FP8E8M0)

required

Returns:

Type Description
Tile

Tile wrapping the matmul_mx_acc operation

matmul_mx_bias(lhs, lhs_scale, rhs, rhs_scale, bias)

MX block-scale matmul with bias.

Data operands follow matmul_mx: an FP4 lhs must first be cast to FP8E4M3FN, and the operation itself receives two FP8E4M3FN tiles.

Parameters:

Name Type Description Default
lhs Tile

Left-hand side data tile (FP8E4M3FN)

required
lhs_scale Tile

Left-hand side scale tile (FP8E8M0)

required
rhs Tile

Right-hand side data tile (FP8E4M3FN)

required
rhs_scale Tile

Right-hand side scale tile (FP8E8M0)

required
bias Tile

Bias tile

required

Returns:

Type Description
Tile

Tile wrapping the matmul_mx_bias operation

gemv(lhs, rhs, acc_phase=AccPhase.Unspecified)

General Matrix-Vector multiplication: C[1,N] = A[1,K] @ B[K,N].

lhs must have exactly one physical and logical row. The rhs logical K must cover the lhs logical K. Inputs must use the same INT8, FP16, BF16, or FP32 dtype; the output is INT32 for INT8 inputs and FP32 otherwise.

Parameters:

Name Type Description Default
lhs Tile

Row vector tile [1, K]

required
rhs Tile

Right-hand side tile [K, N]

required
acc_phase AccPhase

Producer-side unit-flag phase. Use pl.AccPhase.Partial for intermediate chunks and pl.AccPhase.Final for the last chunk.

Unspecified

Returns:

Type Description
Tile

Tile wrapping the gemv operation

gemv_acc(acc, lhs, rhs, acc_phase=AccPhase.Unspecified, *, init_cond=None)

GEMV with accumulation: C[1,N] += A[1,K] @ B[K,N].

acc must use the GEMV output dtype. The logical K extents and lhs/rhs dtype requirements are identical to gemv.

init_cond makes the accumulator's initial value conditional, exactly as in matmul_acc — GEMV is a matmul whose M is 1, run on the same cube MAD, so it carries the same predicate bit. On the steps where it holds, acc is overwritten with lhs @ rhs rather than accumulated into, which removes the peeled first step from split-K::

for k0 in pl.pipeline(0, K, K_TILE):
    a = pl.load(vec, [0, k0], [1, K_TILE], target_memory=pl.MemorySpace.Mat)
    b = pl.load(mat, [k0, 0], [K_TILE, N], target_memory=pl.MemorySpace.Mat)
    acc = pl.tile.gemv_acc(acc, a, b, init_cond=(k0 == 0))

A literal True / False selects one form at compile time; a runtime predicate lowers to a branch over the two, with no phi on the accumulator.

init_cond is keyword-only because acc_phase already owns the fourth positional slot.

Parameters:

Name Type Description Default
acc Tile

Accumulator tile [1, N]

required
lhs Tile

Row vector tile [1, K]

required
rhs Tile

Right-hand side tile [K, N]

required
acc_phase AccPhase

Producer-side unit-flag phase. Use pl.AccPhase.Partial for intermediate chunks and pl.AccPhase.Final for the last chunk.

Unspecified
init_cond BoolLike | None

Optional predicate selecting overwrite over accumulate

None

Returns:

Type Description
Tile

Tile wrapping the gemv_acc operation

gemv_bias(lhs, rhs, bias, acc_phase=AccPhase.Unspecified)

GEMV with bias add: C[1,N] = A[1,K] @ B[K,N] + bias[1,N].

bias must use the GEMV output dtype and its valid shape must cover the logical output shape [1, N]. The logical K extents and lhs/rhs dtype requirements are identical to gemv.

Parameters:

Name Type Description Default
lhs Tile

Row vector tile [1, K]

required
rhs Tile

Right-hand side tile [K, N]

required
bias Tile

Bias tile [1, N] with the accumulator dtype (FP32 for floating-point matrix operands, INT32 for integer matrix operands)

required
acc_phase AccPhase

Producer-side unit-flag phase. Use pl.AccPhase.Partial for intermediate chunks and pl.AccPhase.Final for the last chunk.

Unspecified

Returns:

Type Description
Tile

Tile wrapping the gemv_bias operation

minimum(lhs, rhs)

minimum(lhs: Tensor, rhs: Tensor | int | float | Scalar) -> Tensor
minimum(lhs: Tile, rhs: Tile | int | float | Scalar) -> Tile

Element-wise minimum, dispatched by input type.

min(lhs, rhs)

Scalar min of two values.

Tile reductions are direction-specific — use row_min (collapses the last axis) or col_min (collapses axis 0).

Parameters:

Name Type Description Default
lhs Scalar | int | Expr

First scalar operand

required
rhs Scalar | int | Expr

Second scalar operand

required

Returns:

Type Description
Scalar

Scalar wrapping the min operation

max(lhs, rhs)

Scalar max of two values.

Tile reductions are direction-specific — use row_max (collapses the last axis) or col_max (collapses axis 0).

Parameters:

Name Type Description Default
lhs Scalar | int | Expr

First scalar operand

required
rhs Scalar | int | Expr

Second scalar operand

required

Returns:

Type Description
Scalar

Scalar wrapping the max operation

cmp(lhs, rhs, cmp_type=0)

cmp(lhs: Tensor, rhs: Tensor | int | float | Scalar, cmp_type: int = 0) -> Tensor
cmp(lhs: Tile, rhs: Tile | int | float | Scalar, cmp_type: int = 0) -> Tile

Element-wise comparison, dispatched by input type.

Comparison type codes: 0=eq, 1=ne, 2=lt, 3=le, 4=gt, 5=ge. For Tile inputs with a scalar rhs, dispatches to tile.cmps automatically.

cmps(lhs, rhs, cmp_type=0)

Element-wise comparison of tile and scalar.

Parameters:

Name Type Description Default
lhs Tile

Tile

required
rhs int | float | Expr | Scalar

Scalar value

required
cmp_type int

Comparison type (EQ=0, NE=1, LT=2, LE=3, GT=4, GE=5)

0

Returns:

Type Description
Tile

Tile wrapping a packed predicate mask. Use tile.sel with an explicit tmp tile to materialize values.

set_validshape(input, valid_rows, valid_cols)

set_validshape(input: Tensor, valid_rows: IntLike, valid_cols: IntLike) -> Tensor
set_validshape(input: Tile, valid_rows: IntLike, valid_cols: IntLike) -> Tile

Update valid-shape metadata without data movement, dispatched by input type.

.. note:: Prefer expressing the extent at its source where possible — pl.load(..., valid_shape=...) or a slice's valid_shape=. A tile view (slice / reshape result) is rejected: it carries its valid extent in its type, so there are no runtime operands to update.

rem(lhs, rhs, tmp, high_precision=False)

Element-wise remainder (modulo) of two tiles.

Computes lhs % rhs element-wise. Maps to the TREM hardware intrinsic. On A2/A3, every INT32 input element must be in [-2**24, 2**24].

Parameters:

Name Type Description Default
lhs Tile

Left-hand side tile

required
rhs Tile

Right-hand side tile

required
tmp Tile

Same-dtype scratch tile whose physical and valid capacity provably provides two rows and covers the dividend columns on A2/A3. It must not overlap either source there.

required
high_precision bool

Whether to select PTOAS's high-precision TREM mode. This mode is defined only for FP32 and is ignored by A2/A3 hardware.

False

Returns:

Type Description
Tile

Tile wrapping the rem operation

rems(lhs, rhs, tmp)

Element-wise remainder (modulo) of tile and scalar.

Computes lhs % rhs element-wise. Maps to the TREMS hardware intrinsic. On A2/A3, source valid extents must be provably positive and every INT32 source element and scalar must be in [-2**24, 2**24].

Parameters:

Name Type Description Default
lhs Tile

Tile

required
rhs int | float | Expr | Scalar

Scalar value

required
tmp Tile

Same-dtype scratch tile whose physical and valid capacity provably provides one row and covers the dividend columns on A2/A3. It must not overlap the dividend there.

required

Returns:

Type Description
Tile

Tile wrapping the rems operation

fmod(lhs, rhs, high_precision=False)

fmod(lhs: Tensor, rhs: Tensor | int | float | Scalar, high_precision: Literal[False] = False) -> Tensor
fmod(lhs: Tile, rhs: Tile | int | float | Scalar, high_precision: bool = False) -> Tile

Element-wise truncating remainder, dispatched by input type.

Matches torch.fmod (the remainder takes the sign of the dividend). high_precision is available only for the tile-tile form.

fmods(lhs, rhs)

fmods(lhs: Tensor, rhs: int | float | Scalar) -> Tensor
fmods(lhs: Tile, rhs: int | float | Scalar) -> Tile

Element-wise truncating remainder with a scalar, dispatched by input type.

and_(lhs, rhs)

and_(lhs: Tensor, rhs: Tensor | int | Scalar) -> Tensor
and_(lhs: Tile, rhs: Tile | int | Scalar) -> Tile

Element-wise bitwise AND, dispatched by input type.

ands(lhs, rhs)

ands(lhs: Tensor, rhs: int | Scalar) -> Tensor
ands(lhs: Tile, rhs: int | Scalar) -> Tile

Element-wise bitwise AND with a scalar, dispatched by input type.

or_(lhs, rhs)

or_(lhs: Tensor, rhs: Tensor | int | Scalar) -> Tensor
or_(lhs: Tile, rhs: Tile | int | Scalar) -> Tile

Element-wise bitwise OR, dispatched by input type.

ors(lhs, rhs)

ors(lhs: Tensor, rhs: int | Scalar) -> Tensor
ors(lhs: Tile, rhs: int | Scalar) -> Tile

Element-wise bitwise OR with a scalar, dispatched by input type.

xor(lhs, rhs, tmp=None)

xor(lhs: Tensor, rhs: Tensor | int | Scalar) -> Tensor
xor(lhs: Tile, rhs: Tile | int | Scalar, tmp: Tile) -> Tile

Element-wise bitwise XOR, dispatched by input type.

pto.txor needs a scratch buffer. Tile buffer lifetimes are user-managed, so the tile path takes it as tmp; the tensor path omits it because ConvertTensorToTileOps allocates it — the same asymmetry rsqrt carries.

xors(lhs, rhs, tmp=None)

xors(lhs: Tensor, rhs: int | Scalar) -> Tensor
xors(lhs: Tile, rhs: int | Scalar, tmp: Tile) -> Tile

Element-wise bitwise XOR with a scalar, dispatched by input type.

See xor for why only the tile path takes tmp.

shl(lhs, rhs)

shl(lhs: Tensor, rhs: Tensor | int | Scalar) -> Tensor
shl(lhs: Tile, rhs: Tile | int | Scalar) -> Tile

Element-wise bitwise left shift, dispatched by input type.

shls(lhs, rhs)

shls(lhs: Tensor, rhs: int | Scalar) -> Tensor
shls(lhs: Tile, rhs: int | Scalar) -> Tile

Element-wise bitwise left shift by a scalar, dispatched by input type.

shr(lhs, rhs)

shr(lhs: Tensor, rhs: Tensor | int | Scalar) -> Tensor
shr(lhs: Tile, rhs: Tile | int | Scalar) -> Tile

Element-wise bitwise right shift, dispatched by input type.

shrs(lhs, rhs)

shrs(lhs: Tensor, rhs: int | Scalar) -> Tensor
shrs(lhs: Tile, rhs: int | Scalar) -> Tile

Element-wise bitwise right shift by a scalar, dispatched by input type.

maximums(lhs, rhs)

Element-wise maximum of tile and scalar.

Computes max(lhs, rhs) element-wise. Maps to the TMAXS hardware intrinsic.

Parameters:

Name Type Description Default
lhs Tile

Tile

required
rhs int | float | Expr | Scalar

Scalar value

required

Returns:

Type Description
Tile

Tile wrapping the maximums operation

minimums(lhs, rhs)

Element-wise minimum of tile and scalar.

Computes min(lhs, rhs) element-wise. Maps to the TMINS hardware intrinsic.

Parameters:

Name Type Description Default
lhs Tile

Tile

required
rhs int | float | Expr | Scalar

Scalar value

required

Returns:

Type Description
Tile

Tile wrapping the minimums operation

prelu(tile, slope, tmp)

Element-wise parametric ReLU of a tile.

Computes prelu(tile, slope) element-wise. Maps to the TPRELU hardware intrinsic.

Parameters:

Name Type Description Default
tile Tile

Input tile

required
slope Tile

Slope tile used for negative values

required
tmp Tile

Temporary tile required by the hardware

required

Returns:

Type Description
Tile

Tile wrapping the prelu operation

not_(input)

Element-wise bitwise NOT, dispatched by input type (int16/uint16 only).

addc(lhs, rhs, rhs2)

Element-wise carry addition of three tiles.

Computes src0 + src1 + carry element-wise. Maps to TADDC.

Parameters:

Name Type Description Default
lhs Tile

First source tile

required
rhs Tile

Second source tile

required
rhs2 Tile

Per-element carry-in tile, normally containing 0 or 1

required

Returns:

Type Description
Tile

Tile wrapping the addc operation

subc(lhs, rhs, rhs2)

Element-wise carry subtraction of three tiles.

Computes src0 - src1 + carry element-wise. Maps to TSUBC.

Parameters:

Name Type Description Default
lhs Tile

Minuend tile

required
rhs Tile

Subtrahend tile

required
rhs2 Tile

Per-element carry-in tile, normally containing 0 or 1

required

Returns:

Type Description
Tile

Tile wrapping the subc operation

addsc(lhs, rhs, rhs2)

Element-wise scalar carry addition.

Computes src0 + scalar + carry element-wise. Maps to TADDSC.

Parameters:

Name Type Description Default
lhs Tile

Source tile

required
rhs int | float | Expr | Scalar

Scalar addend with the same dtype as lhs

required
rhs2 Tile

Per-element carry-in tile, normally containing 0 or 1

required

Returns:

Type Description
Tile

Tile wrapping the addsc operation

subsc(lhs, rhs, rhs2)

Element-wise scalar carry subtraction.

Computes src0 - scalar + carry element-wise. Maps to TSUBSC.

Parameters:

Name Type Description Default
lhs Tile

Minuend tile

required
rhs int | float | Expr | Scalar

Scalar subtrahend with the same dtype as lhs

required
rhs2 Tile

Per-element carry-in tile, normally containing 0 or 1

required

Returns:

Type Description
Tile

Tile wrapping the subsc operation

lrelu(tile, slope)

Element-wise leaky ReLU with scalar slope.

Computes max(tile, slope * tile) element-wise. Maps to the TLRELU hardware intrinsic.

Parameters:

Name Type Description Default
tile Tile

Input tile

required
slope int | float | Expr | Scalar

Scalar slope for negative values

required

Returns:

Type Description
Tile

Tile wrapping the lrelu operation

sel(mask, lhs, rhs, tmp)

Per-element selection between two tiles using a predicate mask tile.

For each element (i, j): dst[i,j] = lhs[i,j] if mask[i,j] is true, else rhs[i,j]. Maps to the TSEL hardware intrinsic. The mask encoding is target-defined.

Parameters:

Name Type Description Default
mask Tile

Predicate mask tile; encoding is target-defined

required
lhs Tile

Source tile 0, selected where mask is true

required
rhs Tile

Source tile 1, selected where mask is false

required
tmp Tile

Scratch tile required by TSEL (UINT32 [1, 16] on A2/A3; unread ABI placeholder on A5)

required

Returns:

Type Description
Tile

Tile wrapping the sel operation

sels(mask, src, tmp, scalar)

Per-element selection between a source tile and a scalar.

For each element (i, j): dst[i,j] = src[i,j] if mask[i,j] is true, else scalar. Maps to the TSELS hardware intrinsic.

Parameters:

Name Type Description Default
mask Tile

Predicate mask tile; encoding is target-defined

required
src Tile

Source tile, selected where mask is true

required
tmp Tile

Scratch tile required by TSELS

required
scalar int | float | Expr | Scalar

Scalar value, selected where mask is false

required

Returns:

Type Description
Tile

Tile wrapping the sels operation

tri(diagonal, shape, valid_shape=None, dtype=DataType.INT32, upper=False)

Generate a lower- or upper-triangular mask tile.

upper=False writes one where j <= i + diagonal; upper=True writes one where j >= i + diagonal. Only the optional valid region is written.

Parameters:

Name Type Description Default
diagonal int | Scalar

Offset of the boundary from the main diagonal, in columns. 0 includes the diagonal; positive shifts it right, negative left. May be a runtime Scalar.

required
shape Sequence[int]

Shape of the destination tile (static).

required
valid_shape Sequence[int] | None

Optional written region (each dim <= shape). Elements outside it are not written, so their value is whatever the freshly allocated tile holds. Defaults to the full shape.

None
dtype DataType

Destination dtype. Defaults to INT32.

INT32
upper bool

Select the upper triangle instead of the lower.

False

Returns:

Type Description
Tile

A tile holding 1 inside the selected triangle and 0 outside it.

quant_mx(src, *, group_axis, dtype=DataType.FP8E4M3FN)

MXFP8 block-32 dynamic quantization with PTOAS grpAxis selection.

group_axis is required and must be 0 or 1 (same meaning as PTOAS #pto<mx_group_axis axis*>). Axis 1 is the A-side [M, K] path; axis 0 is the B-side [N, K] path (LowerCompositeOps transposes to [K, N] before quantization). GM TensorLayout.MX_* is only for tensor / tensor.view annotations, not this tile op.

This release supports MXFP8 only (dtype=FP8E4M3FN). MXFP4 is out of scope for this PR.

Parameters:

Name Type Description Default
src Tile

Source tile (FP16/FP32/BF16, 2D). group_axis=1: [M, K] with M%16==0, K%64==0. group_axis=0: [N, K] with N%32==0, K%64==0.

required
group_axis int

PTOAS grouping axis — 1 (A-side) or 0 (B-side).

required
dtype DataType

Must be FP8E4M3FN (default).

FP8E4M3FN

Returns:

Type Description
Tile

group_axis=1: (quant[M,K], scale[M,K/32]) row/row ZZ scale.

Tile

group_axis=0: (quant[K,N], scale[K/32,N]) — data transposed to

tuple[Tile, Tile]

Cube RHS layout; scale col/col NN.

Note

On Ascend950, quant_mx and matmul_mx may share one InCore mixed task. The compiler carries both generated results over V2C; the FP8E8M0 scale keeps its logical MX scale layout.

tmov_x2zz(src, tmp, *, group_axis=1, dst_rows=None, dst_cols=None)

Exponent X-to-ZZ layout conversion (A5-only).

tmp is a write-only workspace. Axis1 requires capacity 64 + ceil(rows/16) * cols bytes (typically 32-byte-aligned by the caller) and dst_rows/dst_cols for the ZZ [M, G] shape. Axis0 (TMovDnTo2Zz) still requires a Vec tmp operand; use a minimal 32-byte-aligned pad.

AUTO = -1 module-attribute

tpush_to_aiv(tile, *, split, lane_stride=None, id=None, span=None)

Push tile data from AIC to AIV via cross-core pipe.

The Vector side receives it with tpop_from_aic and releases the slot with tfree_to_aic; split and id must match across all three.

Parameters:

Name Type Description Default
tile Tile

Tile to send. Its Cube-side buffer stays live until the consumer frees the slot.

required
split int

pto-isa split code (0=none, 1/2=up-down/left-right, 3/4=the same axes over an odd extent). Selects the axis along which the two AIV lanes divide the tile, and how their extents relate; 0 sends it whole.

required
lane_stride int | None

Partition stride carried when a ragged boundary was balanced across the two AIV lanes; omit for the default box partition.

None
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

tpush_to_aic(tile, *, split, id=None, span=None)

Push tile data from AIV to AIC via cross-core pipe.

The Cube side receives it with tpop_from_aiv and releases the slot with tfree_to_aiv; split and id must match across all three.

Parameters:

Name Type Description Default
tile Tile

Tile to send. Its Vector-side buffer stays live until the consumer frees the slot.

required
split int

Split mode (0=none, 1=up-down, 2=left-right). Selects the axis along which the two AIV lanes divide the tile; 0 sends it whole.

required
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

aiv_shard(x, span=None)

Bring a cube-produced operand onto the AIV lane (AIC -> AIV crossing).

In a data-parallel region (UP_DOWN / LEFT_RIGHT) the crossing also halves the operand along the split axis, so each lane gets one half; in a task-parallel mode=NONE region there is no split axis, so it crosses and preserves the shape. Either way, writing it is how a C->V crossing into a region is named — the AivSplitValid verifier rejects an unnamed one.

The split mode is inherited from the enclosing for aiv_id in pl.split_aiv(mode=...) scope — it is not passed here. This wrapper therefore only resolves inside a parsed kernel, where the parser intercepts the call and fills the inherited mode. Calling it eagerly (outside a parsed program) raises, since there is no scope to read the mode from.

The operand may be a Tile (legacy @pl.program form -> tile.aiv_shard) or a high-level Tensor (@pl.jit / pl.spmd form -> tensor.aiv_shard, lowered 1:1 to tile.aiv_shard at ConvertTensorToTileOps). The Tensor form is region-only. Distributed tensors are not supported.

Parameters:

Name Type Description Default
x _SplitOperandT

Input operand (2D Tile or Tensor)

required
span Span | None

Optional source span

None

Returns:

Type Description
_SplitOperandT

Operand of the same kind: the split axis halved in a data-parallel region,

_SplitOperandT

the shape unchanged in a mode=NONE one.

aic_gather(x, span=None)

Hand a vector-produced operand to the cube (AIV -> AIC crossing).

Inverse of aiv_shard: in a data-parallel region it rejoins the two lanes' halves along the split axis, and in a task-parallel mode=NONE region it crosses and preserves the shape. It is how a V->C crossing out of a region is named; an unnamed one is rejected by the AivSplitValid verifier.

Out of a mode=NONE region the two lanes share one destination slot with no per-lane offset and nothing arbitrates between them: both push, so when they hold different values the cube receives an unspecified one of the two. Guarding the production of the value does not help — lane 1 still reaches the push and still sends its own tile. Gather only a value both lanes agree on; if they must contribute different data, use a data-parallel region.

Like aiv_shard, the split mode is inherited from the enclosing for aiv_id in pl.split_aiv(mode=...) scope and must not be passed here. Calling it eagerly (outside a parsed program) raises, since there is no scope to read the mode from.

The operand may be a Tile (legacy @pl.program form -> tile.aic_gather) or a high-level Tensor (@pl.jit / pl.spmd form -> tensor.aic_gather, lowered 1:1 to tile.aic_gather at ConvertTensorToTileOps). The Tensor form is region-only. Distributed tensors are not supported.

Parameters:

Name Type Description Default
x _SplitOperandT

Input operand (2D Tile or Tensor)

required
span Span | None

Optional source span

None

Returns:

Type Description
_SplitOperandT

Operand of the same kind: the split axis doubled in a data-parallel region,

_SplitOperandT

the shape unchanged in a mode=NONE one.

tpop_from_aic(*, shape=None, dtype=None, split=0, lane_stride=None, id=None, span=None)

Pop tile data from AIC cross-core pipe into AIV.

Parameters:

Name Type Description Default
shape list[int] | None

Shape of the tile to receive

None
dtype DataType | None

Data type of the tile to receive

None
split int

pto-isa split code (0=none, 1/2=up-down/left-right, 3/4=the same axes over an odd extent)

0
lane_stride int | None

Partition stride carried when a ragged boundary was balanced across the two AIV lanes; omit for the box partition

None
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

tpop_from_aiv(*, shape=None, dtype=None, split=0, id=None, span=None)

Pop tile data from AIV cross-core pipe into AIC.

Parameters:

Name Type Description Default
shape list[int] | None

Shape of the tile to receive

None
dtype DataType | None

Data type of the tile to receive

None
split int

Split mode (0=none, 1=up-down, 2=left-right)

0
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

aic_initialize_pipe(c2v_consumer_buf=0, v2c_consumer_buf=0, *, dir_mask, slot_size, slot_num=None, local_slot_num=None, id=None, span=None)

Initialize cross-core pipe on AIC side.

Parameters:

Name Type Description Default
c2v_consumer_buf PipeBufOperand

C2V consumer buffer base (Expr, int, or DSL Scalar; default 0)

0
v2c_consumer_buf PipeBufOperand

V2C consumer buffer base (Expr, int, or DSL Scalar; default 0)

0
dir_mask int

Direction mask for pipe

required
slot_size int

Size of each pipe slot

required
slot_num int | None

Optional ring-buffer slot count. Omit to let PTOAS pick its default (8 unidirectional, 4 per direction bidirectional).

None
local_slot_num int | None

Optional local slot count (a2/a3 only, must be <= slot_num). On a3 the reserved/imported buffer is sized slot_size * local_slot_num; on a5 it is slot_size * slot_num.

None
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

aiv_initialize_pipe(c2v_consumer_buf=0, v2c_consumer_buf=0, *, dir_mask, slot_size, slot_num=None, local_slot_num=None, id=None, span=None)

Initialize cross-core pipe on AIV side.

Parameters:

Name Type Description Default
c2v_consumer_buf PipeBufOperand

C2V consumer buffer base (Expr, int, or DSL Scalar; default 0)

0
v2c_consumer_buf PipeBufOperand

V2C consumer buffer base (Expr, int, or DSL Scalar; default 0)

0
dir_mask int

Direction mask for pipe

required
slot_size int

Size of each pipe slot

required
slot_num int | None

Optional ring-buffer slot count. Omit to let PTOAS pick its default (8 unidirectional, 4 per direction bidirectional).

None
local_slot_num int | None

Optional local slot count (a2/a3 only, must be <= slot_num). On a3 the reserved/imported buffer is sized slot_size * local_slot_num; on a5 it is slot_size * slot_num.

None
id int | None

Optional frontend pipe id. Omit to use PTOAS default id 0.

None
span Span | None

Optional source span

None

reserve_buffer(*, name, size, base=AUTO, span=None)

Reserve a named buffer for cross-core communication.

Parameters:

Name Type Description Default
name str

Buffer name for cross-core reference.

required
size int

Buffer size in bytes.

required
base int

Base address in local SRAM. Use AUTO (-1) to let the compiler pick a non-conflicting address, or an explicit integer for manual kernels.

AUTO
span Span | None

Optional source span.

None

Returns:

Type Description
Scalar

pl.Scalar[pl.INT32] wrapping the system.reserve_buffer IR call (PTO ... -> i32).

import_peer_buffer(*, name, peer_func, span=None)

Import a buffer from a peer function in the same group.

Parameters:

Name Type Description Default
name str

Buffer name to import (must match peer's reserve_buffer name).

required
peer_func str

Name of the peer function that owns the buffer.

required
span Span | None

Optional source span.

None

Returns:

Type Description
Scalar

pl.Scalar[pl.INT32] wrapping the system.import_peer_buffer IR call (PTO ... -> i32).

tfree_to_aic(tile, span=None, *, split=None, id=None)

Release ring buffer slot back to AIC producer.

Call this once the tile from tpop_from_aic has been consumed. Until it runs, the slot stays occupied and the producer blocks once the ring fills.

Parameters:

Name Type Description Default
tile Tile

The tile returned by the matching tpop_from_aic.

required
span Span | None

Optional source span

None
split int | None

Leave None. The StampTfreeSplit pass always takes this from the originating tpop, so a value passed here cannot override it.

None
id int | None

Optional frontend pipe id, inherited from the originating tpop when omitted. Supplying one that disagrees with the tpop is rejected.

None

tfree_to_aiv(tile, span=None, *, split=None, id=None)

Release ring buffer slot back to AIV producer.

Call this once the tile from tpop_from_aiv has been consumed. Until it runs, the slot stays occupied and the producer blocks once the ring fills.

Parameters:

Name Type Description Default
tile Tile

The tile returned by the matching tpop_from_aiv.

required
span Span | None

Optional source span

None
split int | None

Leave None. The StampTfreeSplit pass always takes this from the originating tpop, so a value passed here cannot override it.

None
id int | None

Optional frontend pipe id, inherited from the originating tpop when omitted. Supplying one that disagrees with the tpop is rejected.

None

assemble(target, source, offset, *, atomic=AtomicType.None_)

assemble(target: Tensor, source: Tensor, offset: Sequence[IntLike], *, atomic: AtomicType = ...) -> Tensor
assemble(target: Tile, source: Tile, offset: Sequence[IntLike], *, atomic: Literal[AtomicType.None_] = ...) -> Tile

Write source into target at offset, dispatched by target type.

atomic is Tensor-only: the combine lowers to an atomic-add store into global memory, which a tile-to-tile assemble has no destination for. Passing the documented default keeps working on both paths; any other value with a Tile target raises.

cos(input)

Element-wise cosine (input in radians), dispatched by input type. FP32 only.

gather_row(dst, src, dst_offset, src_offset, shapes, transpose=False, *, valid_shape=None)

Gather one GM row into a sub-region of an on-chip accumulator (DPS).

Dispatched on dst — the destination accumulator. src is a Tensor on both paths: this op always reads from global memory, so it is the destination, not the source, that names the level.

mrgsort(src0, src1=None, src2=None, src3=None, tmp=None, *, exhausted=False, block_len=None)

Merge sort — format1 (single-list) or format2 (2-4 way), dispatched by input type.

tmp is Tile-only: at tensor level the scratch buffer is synthesized during Tensor-to-Tile lowering, so passing one is rejected rather than silently dropped. The tile path's format2 requires it, and pl.tile.mrgsort raises with the per-format guidance when it is missing.

exhausted is keyword-only here. The tile wrapper also accepts it as a sixth positional argument; that spelling stays available as pl.tile.mrgsort(...).

scatter_update(input, *args, **kwargs)

Update rows at positions given by a 2D index, dispatched by input type.

Accepts the same flexible call shapes as either level's wrapper — the positional/keyword forms are identical on both, so the arguments are forwarded unchanged.

sin(input)

Element-wise sine (input in radians), dispatched by input type. FP32 only.

sort32(src, idx)

Sort fixed 32-element blocks, permuting idx alongside src.

Dispatched by input type. Returns 8-byte value-index pairs; the last dimension is 2x the input width for FP32 and 4x for FP16.

create_tensor = create module-attribute

dim(tensor, axis)

Extract a shape dimension from a tensor as a scalar value.

Parameters:

Name Type Description Default
tensor Tensor

Input tensor

required
axis int | ConstInt

Dimension index (supports negative indexing). Accepts either a Python int or a ConstInt (parser-shape).

required

Returns:

Type Description
Scalar

Scalar wrapping the dim operation (INT64)

full(shape, dtype, value)

Create a tensor of specified shape filled with a constant value.

Parameters:

Name Type Description Default
shape Sequence[IntLike]

Shape of the tensor

required
dtype DataType

Data type of tensor elements

required
value int | float

Filling scalar value (int or float)

required

Returns:

Type Description
Tensor

Tensor wrapping the full operation

ScopeMode

Bases: Enum

Dependency-tracking mode of a runtime scope (SIMPLER_SCOPE).

  • AUTO: OverlapMap auto-dependency tracking is on (SIMPLER_SCOPE()).
  • MANUAL: auto tracking is off; the user declares every edge via pl.submit(..., deps=[...]) (SIMPLER_SCOPE(ScopeMode::MANUAL)).

AUTO = 0 class-attribute instance-attribute

MANUAL = 1 class-attribute instance-attribute

manual_scope

Alias for pl.scope(mode=pl.ScopeMode.MANUAL).

Turns OverlapMap auto dep-tracking off for a region.

Inside this block, the simpler runtime skips OverlapMap lookup and insert for every kernel submit, so the user takes full responsibility for declaring task-to-task ordering edges. manual_scope is the coarsest-grained of the runtime's auto-dep-tracking opt-outs; finer granularities are available and compose with auto scope:

  • with pl.manual_scope(): — this construct; whole-region opt-out.
  • pl.create_tensor(..., manual_dep=True) — opt out a single tensor for its entire lifetime (any task referencing it skips OverlapMap).
  • pl.no_dep(arg) at a kernel-call arg position — opt out a single tensor for a single task only.

Manual dependency edges (pl.submit(..., deps=[...])) are orthogonal to all of the above: the runtime adds them on top of whatever auto-tracked deps remain (final fanin = auto ∪ explicit), so deps= works in auto scope too. Use manual_scope only when you want full ownership of the dep graph; otherwise stay in auto scope and use pl.submit(..., deps=) as a precision tool that patches the edges auto cannot infer.

Usage::

# Full-manual region.
with pl.manual_scope():
    scratch, tid = pl.submit(self.stage1, x, scratch)
    out, _       = pl.submit(self.stage2, scratch, out, deps=[tid])

# Auto scope with explicit-edge patching (no manual_scope needed).
a, a_tid = pl.submit(self.k1, x)
b, _     = pl.submit(self.k2, x, deps=[a_tid])
Restrictions
  • Must appear inside an Orchestration function (not InCore).
  • Cannot be nested inside another manual_scope (runtime forbids).

submit(*args, **kwargs)

Submit a kernel and capture its producer TaskId.

pl.submit is a parser construct, not a runtime function — the DSL parser intercepts result, tid = pl.submit(self.kernel, *args, deps=[...]) syntactically and never actually calls this body. It is defined only so the name resolves (for imports / linters).

Surface form (must be unpacked as a 2-tuple)::

out, tid    = pl.submit(self.stage1, x, scratch, deps=[prev_tid])
(a, b), tid = pl.submit(self.multi_out_kernel, x)
out, tid    = pl.submit(self.stage1, x, scratch, deps=[prev_tid], dumps=[x])
out, tid    = pl.submit(self.stage1, x, scratch, allow_early_resolve=True)

The kernel-side ir.Submit natively returns Tuple[<kernel return>, TASK_ID]; element 0 is the tensor result(s), element 1 is the producer TaskId (Scalar[TASK_ID]). The optional deps=[...] kwarg lists TaskId scalars / arrays this submit must wait on. The single-LHS form res = pl.submit(...) is also accepted and binds the whole flat tuple to one variable.

pl.submit and its deps= kwarg work in both auto and manual scope: Arg::set_dependencies is orthogonal to OverlapMap auto-tracking (final fanin = auto ∪ explicit). In auto scope, use deps=[...] as a precision tool to patch edges the runtime cannot infer; in pl.manual_scope(), use it to declare every edge.

In a distributed HOST orchestrator, the returned TaskId is backed by the L3 runtime's opaque TaskHandle. Each explicit deps=[producer] entry lowers to TaskArgs.add_dep_wait(producer) and therefore adds ordering without extending the producer's resource lifetime, on top of the automatically maintained per-rank communication ordering. This HOST form accepts individual TaskIds and requires every callee argument, including Out/InOut tensors, because L3 does not allocate output tensors.

The optional dumps=[...] kwarg is the submit-side selective tensor dump surface (symmetric with deps=): it lists tensor arguments of this submit to mark for dump (simpler#844), so an enabled dump pipeline filters down to just those bindings. Each entry must be a tensor passed positionally to the submitted kernel. dumps= is the explicit dump surface on a submit; the declarative pl.dump_tag(t) statement feeds the same dump_vars set. These marks only take effect under partial dump (RunConfig.enable_dump_args == 1); they are a no-op when dump is off (0) and irrelevant under full dump (2), which captures every binding regardless.

The optional allow_early_resolve=True kwarg (default False) opts this task in as a speculative early-dispatch producer (simpler#1065): the scheduler may pre-stage this task's consumers onto idle cores before it completes, releasing them with a doorbell the instant it finishes. It is a producer-side hint — a consumer only pre-stages once all of its producers are flagged (or already complete). It lowers to Arg::set_allow_early_resolve(true) in orchestration codegen and is a pure scheduling optimisation (no effect on results). Pays off on critical paths built from many short tasks; harmless otherwise.

The optional predicate=(...) kwarg attaches a dispatch predicate the scheduler evaluates at this task's dispatch point — after its dependencies are satisfied, so the value is current without an orchestration-time wait. When the comparison is false the task is retired inline (never dispatched to a core) while its fanin/fanout still settle, so downstream consumers still unlock. Use it to skip work whose need is only known at runtime (e.g. an MoE expert with an empty row count)::

out, tid = pl.spmd_submit(self.expert_ffn, tokens, out, core_num=N,
                          deps=[gather_tid],
                          predicate=(row_count[e] > 0))

The comparison is matched syntactically, never evaluated: in this position row_count[e] > 0 is a declarative spec handed to the scheduler, not a tensor.read plus a compare. Only tensor[indices] OP int-literal is expressible (one comparison; == != > < >= <=), mirroring the runtime's single-comparison predicate — no chained comparisons, arithmetic, or boolean combination. Reduce anything richer to a single gate value in a prior kernel and predicate on that.

Contract: the operand tensor's producer must be one of this submit's deps= (the parser enforces this where statically provable), so the dispatch-point read observes the current value.

The return annotation is Any (not NoReturn) because the parser intercepts the call and binds a 2-tuple to the LHS — downstream code that does out, tid = pl.submit(...) would not type-check under NoReturn.

spmd_submit(*args, **kwargs)

Launch a kernel as an SPMD task and capture its producer TaskId.

pl.spmd_submit is a parser construct, not a runtime function — the DSL parser intercepts result, tid = pl.spmd_submit(self.kernel, *args, core_num=N, sync_start=..., deps=[...]) syntactically and never actually calls this body. It is defined only so the name resolves (imports / linters).

It is the SPMD sibling of submit: a single orchestration task that the runtime fans out across core_num logical blocks (each kernel reads its block index via pl.tile.get_block_idx()). Like submit it returns one producer TaskId, so the whole dispatch can be named as a dependency of later tasks.

Surface form (must be unpacked as a 2-tuple)::

out, tid = pl.spmd_submit(self.incore_kernel, x, y, core_num=8)
out, tid = pl.spmd_submit(self.kernel, x, core_num=8, sync_start=True,
                          deps=[prev_tid])

core_num is a required keyword argument (a positive integer expression) — the positional slots are the kernel's own arguments. sync_start (default False) requires all blocks to launch atomically. deps=[...], allow_early_resolve=True, timing_slot=N and predicate=(...) work exactly as on submit (note: a sync_start task cannot itself be block-by-block pre-staged, but it can still be flagged to let its consumers pre-stage). The callee may be an InCore / AIC / AIV kernel or a co-scheduled Group.

Like submit, it works in both auto and manual scope; its primary use is explicit dependency wiring inside pl.manual_scope().

no_dep(tensor)

Mark a kernel-call argument as no-dependency (caller-site override).

This is a parser-recognized marker — at runtime it returns the wrapped tensor unchanged. The parser detects pl.no_dep(t) at kernel call arg positions and threads the per-arg ArgDirection.NoDep override into the IR Call's attrs. DeriveCallDirections then overwrites the auto-derived direction at that slot to NoDep.

Effect at runtime: the simpler runtime skips both the OverlapMap dependency lookup and the producer insert for this argument. The marker is legal regardless of whether the callee declares the param as In (read) or Out / InOut (write): the caller is asserting out-of-band that there is no RaW / WaW / WaR conflict on the slot — for example, paged-attention writes whose target offset is data-dependent (so the compiler cannot prove disjointness) but are guaranteed disjoint by the runtime allocation protocol.

Only valid as a direct argument to a kernel call::

# Read-side override: shared_input is read-only at the callee, but
# auto-dep would otherwise serialise sibling fan-out reads.
result = self.my_kernel(pl.no_dep(shared_input), output)

# Write-side override: the callee writes into k_cache via
# ``pl.assemble`` at a data-dependent offset; the user knows the
# offsets are disjoint across the parallel fan-out.
self.rope_kv_cache(q_proj, pl.no_dep(k_cache), pl.no_dep(v_cache))

Outside a kernel call argument list it is a no-op (returns tensor unchanged); the parser only injects the override at recognized call sites.

Parameters:

Name Type Description Default
tensor Tensor

The tensor to pass through. Must be a Tensor value.

required

Returns:

Type Description
Tensor

The tensor unchanged. The marker is consumed at parse time.

dump_tag(tensor)

Mark a tensor for selective dump within the enclosing orchestration.

Declarative per-tensor dump marker (simpler#844 selective tensor dump). Writing pl.dump_tag(q) as a standalone statement records q so that every subsequent kernel dispatch consuming that exact value gets q merged into the dispatch's dump_vars — whether that dispatch lowers to a plain ir.Call (the typical @pl.jit / tensor-op path) or an ir.Submit. The runtime then marks those Arg slots via Arg::dump(...).

This is the declarative counterpart to the explicit dumps=[...] kwarg on pl.submit(...) / pl.at(...) — both feed the same dump_vars set (mirroring how a scope's auto-inferred and explicit deps= edges both feed manual_dep_edges). Use dumps= when you want to list the targets explicitly at a single task launch; use pl.dump_tag when one declaration should stick across every subsequent consumer.

Use this to keep tensor dump viable on large workloads (e.g. paged-attention 64bat/8192ctx) where full dump (enable_dump_args=2) saturates the host-side dump collector (~42 MB/s drain rate) by dumping every binding, eventually triggering a STARS op-timeout kill on the AICPU side. Run partial dump (enable_dump_args=1) and tag only the tensors of interest, so the runtime filters out large bindings (1 GB kv-cache, output buffers, etc.) from the collector queue.

Semantics:

  • Forward-sticky over the orch scope — one pl.dump_tag(q) statement affects all subsequent kernel calls in the same orch that consume the tagged value. Tracked by Var identity, never by name: reassigning q (e.g. q = self.foo(q)) produces a new value that the prior tag does not cover — re-tag it if needed.
  • Only effective under partial dump (RunConfig.enable_dump_args == 1) — selective dump filters within the partial pipeline. A no-op when dump is off (0); under full dump (2) every binding is captured, so the tag has nothing to narrow.
  • Consumed at parse time — recorded into the consuming dispatch's dump_vars and emits no IR statement of its own.

Valid as a standalone statement inside an Orchestration function or an Inline helper (@pl.jit.inline / FunctionType.Inline) that the orchestration inlines. The parser rejects pl.dump_tag written in any other function type (AIV / AIC / Mix kernel bodies) with a clear error::

@pl.function(type=pl.FunctionType.Orchestration)
def orch(self, q: pl.Tensor[...], k_cache: pl.Tensor[...], out: pl.Out[...]):
    pl.dump_tag(q)           # mark q for selective dump
    pl.dump_tag(out)         # mark out for selective dump
    s = self.qk_matmul(q, k_cache, scratch)  # q is dumped here
    out = self.pv_matmul(s, k_cache, out)    # out is dumped here

For a single task launch, list the targets explicitly with the dumps=[...] kwarg: out, tid = pl.submit(self.qk_matmul, q, k_cache, scratch, dumps=[q]) or with pl.at(..., dumps=[q]) as tid:.

Inside an Inline helper, the recorded dump_vars ride on the inline body's kernel calls; the InlineFunctions pass splices those calls into the caller and the mutator substitutes the caller's arg for each inline parameter, so tags on both inline parameters and inline body-local pl.create_tensor(...) results take effect at the inlined call sites. No tag migration is needed; multi-level inlining works at the pass's fixpoint.

Limitations (MVP):

A tag fires only when the tagged Var reaches a kernel dispatch as a static, whole-tensor Arg. Values that never reach such an Arg are silently not dumped:

  • Dynamic-offset reads — a value read only through a data-dependent offset (q_flat[runtime_row : runtime_row + N, ...]) lowers to a gather / dynamic-address load, not a whole-tensor Arg, so the tag does not attach. Stage it through a buffer read with static, compile-time tiled offsets and tag that buffer.
  • Orch-level pl.assemble buffers — y = pl.assemble(y, tile, off) lowers to a pure name-alias and emits no kernel dispatch, so there is no Arg to mark. Use a static in-place slice store y[off_slice] = tile and tag y, or dump the producer kernels' output Args.
  • Orch-tier scalar reads — a tensor consumed only by orchestration-level pl.read(...) (e.g. block tables read to compute offsets) never enters a device kernel as a Tensor Arg; the MVP runtime path covers per-task device Args only.
  • Distributed L3+ programs: only chip-level orchestration tasks honour the tag; HOST-tier Python SubWorker tensors are not covered by the runtime's selective dump path.

Parameters:

Name Type Description Default
tensor Tensor

The tensor to mark. Must be a Tensor value bound in the enclosing orchestration's scope.

required

Returns:

Type Description
Tensor

The tensor unchanged. The marker is consumed at parse time.

set_cache_policy(tensor, policy)

Declare the GM cache-access policy for every read of tensor in this scope.

The scope-level surface of the cache policy, for tensor programming — where the GM reads are implicit and there is no pl.load call to annotate. The per-access counterpart is the cache= kwarg on pl.load, which names one read instead of all of them. pl.slice deliberately takes no cache=: a slice computes an address descriptor, it moves no data.

Write it as a standalone statement directly inside a pl.at(...) / pl.spmd(...) scope body. Anywhere else — nested in an if / for inside the scope, or outside a scope altogether — the parser rejects it, because the declaration attaches to the scope, and a conditionally-executed declaration would be a promise the compiler cannot check.

Semantics:

  • A contract, not a hint. CachePolicy.BYPASS asserts two things about tensor: that the kernel streams it with no reuse worth caching, and that nothing writes those bytes while the kernel runs. Mixing a cached write and a bypassing read of the same bytes is a coherency bug the compiler cannot detect, so coherency is the author's contract. This is why the policy is never a default and never inferred. Declaring BYPASS on a tensor the scope itself writes is rejected at outlining.
  • Tracked by Var identity, never by name. The declaration names the binding live at the scope, so rebinding the name afterwards (b = self.foo(b)) yields a new value the declaration does not cover.
  • Consumed at parse time. It emits no IR statement of its own: the parser records it on the enclosing scope, the scope outliner resolves it to the outlined kernel's parameters, and ConvertTensorToTileOps turns it into a cache kwarg on each tile.load that reads a declared parameter.
  • Explicit wins. An explicit pl.load(..., cache=...) overrides the scope declaration for that one access, in both directions — so cache=pl.CachePolicy.DEFAULT opts a single read back into the cache inside a bypassing scope.

Current status: PTOAS has no L2-bypass path yet (https://github.com/hw-native-sys/PTOAS/issues/1356). The declaration is carried all the way to codegen, but codegen emits a warning and compiles it as an ordinary cached access, so generated code is unchanged today. Writing the declaration now is what makes the kernel pick the bypass up for free once that lands.

Parameters:

Name Type Description Default
tensor Tensor

The tensor whose reads the policy applies to. Must be a Tensor value bound outside the scope and read inside it.

required
policy CachePolicy

CachePolicy.BYPASS to declare a streaming, non-cached read; CachePolicy.DEFAULT for an ordinary cached read.

required

Returns:

Type Description
None

Nothing. The marker is consumed at parse time and produces no value.

Example

with pl.at(level=pl.Level.CORE_GROUP, name_hint="mm"): ... pl.set_cache_policy(b, pl.CachePolicy.BYPASS) ... c = pl.matmul(a, b, out_dtype=pl.FP32) ... out = pl.assemble(out, c, [0, 0])

scatter(input, dim=None, index=None, src=None, *, mask_pattern=None, dst=None)

scatter(input: Tensor, dim: int, index: Tensor, src: Tensor) -> Tensor
scatter(input: Tensor, *, mask_pattern: int, dst: Tensor) -> Tensor

Scatter elements of src into input (tensor-level) — index or mask form.

The tensor layer exposes a single unified scatter. Based on the arguments you pass, it lowers to one of two tile-level ops:

Index form (dim + index + src) → pl.tile.scatter — the column-wise inverse of gather, so index has the same shape as src (just like gather's index matches its output)::

output = input
output[b, index[b, k]] = src[b, k]   # for all b, k

MVP: rank-2 input with ``dim == -1``. ``src``/``index`` are ``[rows, K]``;
``input``/output are ``[rows, S]`` with ``K <= S``. ``index`` element
width must match ``input``: 4-byte input → INT32, 2-byte → INT16,
1-byte → INT16.

Mask form (mask_pattern=<int> + dst) → pl.tile.scatter_mask: Writes each row of input into the columns of dst selected by the hardware mask pattern. dst.cols equals input.cols * stride (stride = 2 for P0101/P1010, 4 for P0001..P1000, 1 for P1111). Unlike the gather mask form (a real pto.tgather ISA op on A2/A3 and A5), mask-pattern scatter is not a distinct pto-isa instruction — PyPTO emits it as a pto.tscatter mask-form construct for A2/A3 / CPU-sim style lowering paths.

Parameters:

Name Type Description Default
input Tensor

Base tensor (FP16/FP32/BF16/INT8/INT16/INT32, 2D).

required
dim int | None

(index form) Axis along which to scatter. MVP accepts -1.

None
index Tensor | None

(index form) Per-element destination column indices, same shape as src (INT16/INT32).

None
src Tensor | None

(index form) Source values tensor (same dtype as input).

None
mask_pattern int | None

(mask form, keyword-only) Mask pattern selector (1-7). 1=P0101, 2=P1010, 3=P0001, 4=P0010, 5=P0100, 6=P1000, 7=P1111.

None
dst Tensor | None

(mask form, keyword-only) Destination tensor; dst.cols == input.cols * stride.

None

Returns:

Type Description
Tensor

Tensor representing the post-scatter result.

Examples:

out = scatter(input, dim=-1, index=idx, src=src_vals) out = scatter(input, mask_pattern=pl.tile.MaskPattern.P0101, dst=dst)

arange(start, shape, dtype=DataType.INT32, descending=False)

Generate a contiguous integer sequence into a tensor.

Equivalent to numpy.arange / torch.arange. Lowers to tile.cipto.tci.

Note

pto.tci only populates the first row. Leading dimensions must be 1 — prefer shapes of the form [1, N].

Parameters:

Name Type Description Default
start int | Scalar

Starting integer (plain int or Scalar). Must match dtype.

required
shape Sequence[IntLike]

Destination tensor shape (innermost dim != 1).

required
dtype DataType

Destination dtype. One of {INT16, INT32}. Defaults to INT32.

INT32
descending bool

If True, generate a descending sequence.

False

Returns:

Type Description
Tensor

Tensor wrapping the ci operation.

gather(input, dim=None, index=None, *, mask_pattern=None, output_dtype=None, kvalue=None, cmp_mode=None, out_cols=None, offset=0, count_dtype=None)

gather(input: Tensor, dim: int, index: Tensor) -> Tensor
gather(input: Tensor | Tile, dim: Tensor | Tile) -> Tensor
gather(input: Tensor | Tile, *, index: Tensor | Tile) -> Tensor
gather(input: Tensor, *, mask_pattern: int, output_dtype: int | DataType | None = None) -> Tensor
gather(input: Tensor, *, kvalue: int | Scalar | Expr, cmp_mode: str | int, out_cols: int, offset: int = 0, count_dtype: int | DataType | None = None) -> tuple[Tensor, Tensor]

Gather elements of input — flat / axis / mask / compare form.

The tensor layer exposes a single unified gather. Based on the arguments you pass, it selects the matching form and lowering:

Flat form (index, no dim)::

output = input.reshape(-1)[index]

Also accepts ``pl.gather(input, index)``. Runtime indices are 2D INT32;
shape and valid shape follow ``index``, dtype follows ``input``
(FP16/FP32/INT16/INT32). Indices must address valid source elements;
negative indexing and bounds checking are unsupported.
Contiguous ND GM sources use [`pl.tile.mgather`][pypto.language.tile.mgather];
static 2D unboxed row-major Vec sources use
[`pl.tile.gather`][pypto.language.tile.gather] with managed packing/scratch.
On-chip source rows must be 32-byte aligned unless there is only one row.
GM operands accept local distributed windows; tile indices must be
unboxed row-major Vec. Physical index columns must be positive static
multiples of 16 for FP16/INT16, or 8 for FP32/INT32, including single-row
tiles. Pad physical storage and use ``set_validshape`` for narrower
valid regions, which need not be aligned.

Axis form (dim + index) → pl.tile.gather, for example dim=1::

output[b, k] = input[b, index[b, k]]

Lowering supports rank-2/rank-3 inputs and any axis, including negative axes.
``index`` must be an INT32 tensor, or INT16 when ``input`` is a 16-bit
dtype (FP16/INT16; INT16 indices require A5). Its rank matches ``input``;
non-gather extents cannot exceed the source. Output shape follows ``index``
and dtype follows ``input``.

Mask form (mask_pattern=<int>) → pl.tile.gather_mask: Selects columns of each row by a fixed hardware mask pattern. Last-dim shrinks by 2 (P0101/P1010) or 4 (P0001..P1000), or stays the same for P1111.

Compare form (kvalue + cmp_mode + out_cols) → pl.tile.gather_compare: Scalar threshold compare (applied to every row). Returns (dst, cdst) — gathered indices [rows, out_cols] INT32 and per-row match counts [1, rows] count_dtype.

Parameters:

Name Type Description Default
input Tensor | Tile

Source tensor (FP16/FP32/INT16/INT32); flat form also accepts a Tile.

required
dim int | Tensor | Tile | None

Axis to gather along; omit for flat indexing. A tensor/tile in this positional slot is interpreted as the flat index.

None
index Tensor | Tile | None

Flat form: 2D INT32 tensor/tile. Axis form: tensor with the same rank as input and the index dtype constraints above.

None
mask_pattern int | None

(mask form, keyword-only) Mask pattern selector (1-7). 1=P0101, 2=P1010, 3=P0001, 4=P0010, 5=P0100, 6=P1000, 7=P1111.

None
output_dtype int | DataType | None

(mask form, keyword-only) Optional output dtype with the same bit width as input.dtype (e.g. FP32 → UINT32 for sort32 index bits).

None
kvalue int | Scalar | Expr | None

(compare form, keyword-only) Scalar threshold (dtype must match input).

None
cmp_mode str | int | None

(compare form, keyword-only) "eq"/"ne"/"lt"/"le"/ "gt"/"ge" or int 0..5.

None
out_cols int | None

(compare form, keyword-only) Output column count per row.

None
offset int

(compare form, keyword-only) Starting index offset (default 0).

0
count_dtype int | DataType | None

(compare form, keyword-only) Per-row count dtype, INT32 or UINT32 (defaults to INT32).

None

Returns:

Type Description
Tensor | tuple[Tensor, Tensor]

Tensor (index/mask form) or (dst, cdst) tuple (compare form).

Examples:

out = gather(input, index=flat_idx) out = gather(input, flat_idx) out = gather(input, dim=-1, index=idx) out = gather(input, mask_pattern=1) out = gather(input, mask_pattern=pl.tile.MaskPattern.P1010, output_dtype=pl.UINT32) dst, cdst = gather(input, kvalue=kv, cmp_mode="eq", out_cols=8)

paged_gather(src, indices, block_table, block_size, size, max_indices, *, space=MemorySpace.Mat, col_off=0, is_trans=False, is_b_matrix=False)

Paged gather directly into an on-chip buffer (L1 by default, or UB).

Gathers scattered rows of a 2D paged KV pool src selected by indices, translated through a paged block_table, directly into an L1 (space=Mat, default) or UB (space=Vec) tile — so a subsequent matmul reads from L1 without a GM round-trip. The lowering is a fully-scalar per-row GM -> on-chip load loop on the Cube core (the bulk KV never touches UB).

Physical row per logical index idx::

phys = block_table[idx // block_size] * block_size + idx % block_size

Parameters:

Name Type Description Default
src Tensor

Paged KV pool in GM (2D; FP16/BF16/FP32/INT8).

required
indices Tensor

Logical row indices (INT32; 1D [n] or 2D [1, n]).

required
block_table Tensor

Page table mapping logical block -> physical block (INT32).

required
block_size int

Number of tokens per page block.

required
size int

Number of elements gathered per row (<= src columns).

required
max_indices int

Static upper bound on gathered rows; sizes the on-chip tile.

required
space MemorySpace

Destination space — MemorySpace.Mat (L1, default) or MemorySpace.Vec (UB).

Mat
col_off int

Column start offset within each src row (default 0).

0
is_trans bool

Transpose for matmul B-operand layout (requires space=Mat).

False
is_b_matrix bool

Hint that the result feeds matmul as the B matrix.

False

Returns:

Type Description
Tensor

Tensor of shape [max_indices, size] (or transposed) in the chosen space.

Examples:

kv_l1 = pl.paged_gather(kv_pool, topk_idx, block_table, block_size=128, size=head_dim, max_indices=256) out = pl.matmul(q, kv_l1)

random(key0, key1, counter0, counter1, counter2, counter3, shape, dtype=DataType.UINT32, rounds=10)

Generate counter-based pseudo-random values into a tensor.

Implements a counter-based (Philox/ChaCha-style) RNG. Each element is derived deterministically from the 64-bit key (key0, key1) and 128-bit counter (counter0..counter3) plus the element position, so the same seeds always reproduce the same tensor. Lowers to tile.randompto.trandom.

Parameters:

Name Type Description Default
key0 int | Scalar

Low INT32 key word (plain int or Scalar).

required
key1 int | Scalar

High INT32 key word (plain int or Scalar).

required
counter0 int | Scalar

First INT32 counter word.

required
counter1 int | Scalar

Second INT32 counter word.

required
counter2 int | Scalar

Third INT32 counter word.

required
counter3 int | Scalar

Fourth INT32 counter word.

required
shape Sequence[IntLike]

Destination tensor shape (static).

required
dtype DataType

Destination dtype. One of {INT32, UINT32}. Defaults to UINT32.

UINT32
rounds int

Cipher round count, 7 or 10. Defaults to 10.

10

Returns:

Type Description
Tensor

Tensor wrapping the random operation.

create_l1(shape, dtype, transpose=False)

Create an on-chip (L1/Mat) accumulator for a kernel-driven paged gather.

Companion of gather_row. Returns a tensor-typed value that composes with pl.matmul / softmax but lowers to an L1 (MemorySpace.Mat) tile, so a kernel can build a matmul operand directly on-chip — no GM round-trip.

Parameters:

Name Type Description Default
shape Sequence[IntLike]

Accumulator shape (static dims). For a matmul B-operand use [rows, size]; for a b_trans B-operand filled by transposing gathers use the transposed shape [size, rows] with transpose=True.

required
dtype DataType

Element dtype (matches the gathered src pool).

required
transpose bool

Allocate the transposed Mat (ZN) layout — required when filling with pl.gather_row(..., transpose=True) (a DN2ZN per-row load) and consuming as a pl.matmul(..., b_trans=True) B-operand.

False

Returns:

Type Description
Tensor

Tensor of shape shape backed by L1.

Examples:

kv = pl.create_l1([ATTN_K_TILE, HEAD_DIM], pl.BF16) for r in pl.range(ATTN_K_TILE): kv = pl.gather_row(kv, kv_pool, [r, 0], [phys, 0], [1, HEAD_DIM]) oi = pl.matmul(probs, kv)

get_block_idx()

Get the current block index (tensor-scope alias of pl.tile.get_block_idx).

Lowers to tile.get_block_idx in ConvertTensorToTileOps.

Returns:

Type Description
Scalar

Scalar wrapping the tensor.get_block_idx operation (INDEX type)

Example

block_idx = pl.tensor.get_block_idx()

get_block_num()

Get the total number of blocks in the current SPMD task.

Tensor-scope alias of pl.tile.get_block_num; lowers to tile.get_block_num in ConvertTensorToTileOps.

Returns:

Type Description
Scalar

Scalar wrapping the tensor.get_block_num operation (INDEX type)

Example

block_idx = pl.tensor.get_block_idx() block_num = pl.tensor.get_block_num()

get_subblock_idx()

Get the current sub-block (vector core) index (tensor-scope alias of pl.tile.get_subblock_idx).

Lowers to tile.get_subblock_idx in ConvertTensorToTileOps.

Returns:

Type Description
Scalar

Scalar wrapping the tensor.get_subblock_idx operation (INDEX type)

FunctionType

Bases: Enum

Function type classification.

Categorizes functions by their execution context and purpose: - Opaque: Unspecified (default) - Orchestration: Runs on host/AICPU for control flow and dependency analysis - InCore: Sub-graph on specific AICore (unspecialized) - AIC: Cube core kernel (specialized InCore) - AIV: Vector core kernel (specialized InCore) - Group: Co-scheduled group of AIC + AIV kernels - Spmd: SPMD data-parallel dispatch - Inline: Whole-body substitution at every call site by the InlineFunctions pass (eliminated before any other pass runs) - Graph: A callable orchestration fragment. Its body is orchestration code, but each call site is a single task launch that the host_build_graph runtime records once and replays thereafter.

Opaque = ... class-attribute instance-attribute

Unspecified function type (default).

Orchestration = ... class-attribute instance-attribute

Host/AICPU control and coordination.

InCore = ... class-attribute instance-attribute

AICore sub-graph execution (unspecialized).

AIC = ... class-attribute instance-attribute

Cube core kernel (specialized InCore).

AIV = ... class-attribute instance-attribute

Vector core kernel (specialized InCore).

Group = ... class-attribute instance-attribute

Co-scheduled group of AIC + AIV kernels.

Spmd = ... class-attribute instance-attribute

SPMD data-parallel dispatch.

Inline = ... class-attribute instance-attribute

Whole-body substitution at every call site.

Graph = ... class-attribute instance-attribute

Recordable/replayable orchestration fragment.

ForKind

Bases: Enum

For loop kind classification.

Distinguishes sequential, parallel, unroll, and pipeline for loops: - Sequential: Standard sequential for loop (default) - Parallel: Parallel for loop - Unroll: Compile-time unrolled for loop - Pipeline: Software-pipelined loop (transient marker; stripped by CanonicalizeIOOrder)

Sequential = ... class-attribute instance-attribute

Standard sequential for loop (default).

Parallel = ... class-attribute instance-attribute

Parallel for loop.

Unroll = ... class-attribute instance-attribute

Compile-time unrolled for loop.

Pipeline = ... class-attribute instance-attribute

Software-pipelined loop — lowered by LowerPipelineLoops. The kind persists as a marker through CanonicalizeIOOrder (which demotes it to Sequential on exit) and must not survive past that pass.

AccPhase

Bases: IntEnum

Producer-side unit-flag phase for GEMV accumulator operations.

Stored as int in op kwargs. The values match PTO-ISA's AccPhase ABI.

Unspecified = 0 class-attribute instance-attribute

Do not use the unit-flag protocol.

Partial = 2 class-attribute instance-attribute

Check the unit flag without setting it.

Final = 3 class-attribute instance-attribute

Check and set the unit flag.

AtomicType

Bases: IntEnum

Combine mode for global-memory writes — pld.tensor.put (TPUT) and tile.store (TSTORE).

Stored as int in op kwargs; the C++ deducer validates the int falls within this enum's range.

None_ = 0 class-attribute instance-attribute

Plain store — overwrite the destination.

Add = 1 class-attribute instance-attribute

Atomically add the source data into the destination.

CachePolicy

Bases: IntEnum

GM cache-access policy declared for a tensor read.

A semantic contract the author states, never a hint the compiler invents. Stored as int in the tile.load cache kwarg.

DEFAULT = 0 class-attribute instance-attribute

Ordinary cached GM access.

BYPASS = 1 class-attribute instance-attribute

Streaming access declared to bypass the cache.

Asserts this tensor has no reuse worth caching and that nothing writes those bytes while the kernel runs — mixing a cached write and a bypassing read of the same bytes is a coherency bug.

KernelType

Bases: Enum

Which generated kernel an op belongs to.

A mixed InCore function is expanded into an AIC kernel and an AIV kernel. This says which of the two a cross-core sync op lands in; MIX means both take part, which only a barrier can ask for.

Two neighbouring enums mean different things:

  • FunctionType.AIC / .AIV classify a function, and this classifies an op inside one. A function already declared FunctionType.AIV needs no KernelType on its ops -- there is only one kernel to land in.
  • ir.CoreType labels one physical core in the SoC inventory, which is what Backend::GetCoreCount counts. That is hardware; this is not, and MIX has no CoreType counterpart at all.

Members carry no wire value: each op spells the same kernel differently in its IR attr (system.syncall writes "aic_only" where system.sync_set writes "aic"), so the lowering tables are explicit.

AIC = auto() class-attribute instance-attribute

AIV = auto() class-attribute instance-attribute

MIX = auto() class-attribute instance-attribute

SyncAllMode

Bases: Enum

Barrier implementation selected by system.syncall.

  • HARD: FFTS barrier with no operands; requires full-core occupancy.
  • SOFT: GM-polling barrier; works at partial occupancy.

HARD = 'hard' class-attribute instance-attribute

SOFT = 'soft' class-attribute instance-attribute

Level

Bases: Enum

Hierarchy level in the Linqu machine model.

Levels map bottom-up from individual cores (Level 0) to the global coordinator. Alias values resolve to the same level as their primary name.

AIV = ... class-attribute instance-attribute

Single AIV (Vector) core.

AIC = ... class-attribute instance-attribute

Single AIC (Cube) core.

CORE_GROUP = ... class-attribute instance-attribute

Core-group (e.g. 1 AIC + 2 AIV).

CHIP_DIE = ... class-attribute instance-attribute

Chip die (optional in single-die models).

CHIP = ... class-attribute instance-attribute

Chip (UMA).

HOST = ... class-attribute instance-attribute

Host (single OS instance).

CLUSTER_0 = ... class-attribute instance-attribute

Cluster-level-0 (pod).

CLUSTER_1 = ... class-attribute instance-attribute

Cluster-level-1 (supernode).

CLUSTER_2 = ... class-attribute instance-attribute

Cluster-level-2 (cross-rack).

GLOBAL = ... class-attribute instance-attribute

Global coordinator.

L2CACHE = ... class-attribute instance-attribute

Alias for CHIP_DIE.

PROCESSOR = ... class-attribute instance-attribute

Alias for CHIP.

UMA = ... class-attribute instance-attribute

Alias for CHIP.

NODE = ... class-attribute instance-attribute

Alias for HOST.

POD = ... class-attribute instance-attribute

Alias for CLUSTER_0.

CLOS1 = ... class-attribute instance-attribute

Alias for CLUSTER_1.

CLOS2 = ... class-attribute instance-attribute

Alias for CLUSTER_2.

MemRef

Bases: MemRef

DSL-level memory reference accepting pl.Ptr bases and Scalar offsets.

Identical to ir.MemRef at runtime. The overloads only widen what pyright accepts so that printed IR — which uses pl.Ptr-annotated base variables and Scalar arithmetic byte offsets — type-checks cleanly when re-loaded as a @pl.program.

Called with no offset and size, it declares an allocation of your own rather than describing an existing one: size and address are left for InitMemRef to derive, and the compiler's opportunistic reuse never packs anything else into it. The allocation takes the name of the variable it is bound to, so the name is written once::

scratch = pl.MemRef()

t0: pl.Tile[[64, 64], pl.FP32, scratch, pl.Mem.Vec] = pl.load(x, [0, 0], [64, 64])
t1: pl.Tile[[64, 64], pl.FP32, scratch, pl.Mem.Vec] = pl.exp(t0)

Pass slots=N for N equally-sized slots of one allocation, then pick one by subscript. The slots are contiguous and identically sized, so rotating through them is a ping-pong the packer cannot collapse::

l0c = pl.MemRef(slots=2)

ping: pl.Tile[[64, 64], pl.FP32, l0c[0], pl.Mem.Acc] = pl.tile.matmul(q, b0)
pong: pl.Tile[[64, 64], pl.FP32, l0c[1], pl.Mem.Acc] = pl.tile.matmul(q, b1)

The subscript is an ordinary index expression, so it may be a runtime value — a rotation needs no unrolling::

for i in pl.range(N):
    t: pl.Tile[[64, 64], pl.FP32, l0c[i % 2], pl.Mem.Acc] = pl.tile.matmul(q, b)

A constant index folds into a static address; a runtime one becomes the tile's address at run time. Nothing here is tied to a particular loop form — the index is just an expression, so pl.range, a pl.pipeline sub-index, or a function parameter all work.

Reference it by variable, so a misspelling is a NameError rather than a second allocation. Since the variable is the name, one declaration may not be reached through two names (b = a) and two declarations may not claim one name; both are rejected. pl.MemRef("other") names it explicitly, overriding the variable — that is the form the IR printer emits, so a dumped program reparses without a surrounding Python scope.

Co-liveness is checked per slot: two tiles on different slots are meant to be live together (that is the ping-pong), and only two tiles landing on the same slot can corrupt each other. A runtime index has no static slot to attribute a tile to, so the check is skipped there — the rotation is yours to get right — while isolation from every other allocation still holds. Tiles sharing one declared allocation must also agree on memory space.

Declaring an allocation inside a pl.pipeline(stage=2) body is rejected — the cloned stages would make a tile co-live with itself. Declaring slots and asking the compiler to multi-buffer are alternatives, not layers: to hand-manage a level, drive it with pl.range and give it its own slots.

Note: pl.MemRef(...) calls inside a @pl.program body are resolved by the parser (parser/type_resolver.py), not dispatched through this __init__. A Scalar byte offset is therefore only ever seen by pyright; it never reaches the underlying ir.MemRef constructor.

Role

Bases: Enum

Function role at L3-L7 hierarchy levels.

Distinguishes orchestrators (which build task DAGs and submit work) from sub-workers (which execute concrete compute or data tasks dispatched by an orchestrator at the same level).

Orchestrator = ... class-attribute instance-attribute

Builds DAG, submits tasks, never computes directly.

SubWorker = ... class-attribute instance-attribute

Executes compute/data tasks dispatched by the orchestrator at the same level.

SplitMode

Bases: Enum

Split mode for cross-core data transfer.

NONE = 0 class-attribute instance-attribute

No split.

UP_DOWN = 1 class-attribute instance-attribute

Split vertically (height halved).

LEFT_RIGHT = 2 class-attribute instance-attribute

Split horizontally (width halved).

STPhase

Bases: IntEnum

Consumer-side unit-flag phase for tile.store.

Stored as int in op kwargs. Supported values use PTO-ISA's STPhase ABI.

Unspecified = 0 class-attribute instance-attribute

Do not use the unit-flag protocol.

Final = 3 class-attribute instance-attribute

Check and clear the unit flag.

Mem = MemorySpace module-attribute

MemRefType

Opaque sentinel type for MemRef-typed variables in printed IR.

The C++ printer emits pl.MemRefType as the annotation for a bare MemRef variable (the SSA-edge type of a MemRef expression). This class exists solely so that the printed Python source is valid Python that the text-parser can exec().

Note: this is not the type of a tile.alloc / tensor.alloc result — those produce a base Ptr (PtrType); see alloc.

MemorySpace

Bases: Enum

Memory space enumeration.

DDR = ... class-attribute instance-attribute

DDR memory (off-chip).

Vec = ... class-attribute instance-attribute

Vector/unified buffer (on-chip).

Mat = ... class-attribute instance-attribute

Matrix/L1 buffer.

Left = ... class-attribute instance-attribute

Left matrix operand buffer.

Right = ... class-attribute instance-attribute

Right matrix operand buffer.

Acc = ... class-attribute instance-attribute

Accumulator buffer.

Bias = ... class-attribute instance-attribute

Bias buffer.

LeftScale = ... class-attribute instance-attribute

L0A-side MX block-scale buffer (A5).

RightScale = ... class-attribute instance-attribute

L0B-side MX block-scale buffer (A5).

ScalarLocal = ... class-attribute instance-attribute

On-core scalar register file / C stack (ArrayType).

PipeType

Bases: IntEnum

Pipeline type enumeration for hardware execution units.

MTE1 = ... class-attribute instance-attribute

MTE2 = ... class-attribute instance-attribute

MTE3 = ... class-attribute instance-attribute

M = ... class-attribute instance-attribute

V = ... class-attribute instance-attribute

S = ... class-attribute instance-attribute

FIX = ... class-attribute instance-attribute

ALL = ... class-attribute instance-attribute

Ptr

DSL wrapper for an ir.PtrType-valued expression.

Construct without arguments to obtain an annotation-only placeholder (buf: pl.Ptr). Construct with expr= to wrap an IR Call returned by an allocation op.

unwrap()

Return the wrapped ir.Expr.

Raises:

Type Description
RuntimeError

If the instance was constructed without an expr (annotation-only — never returned by a wrapper).

PrefetchAsyncContext

Bases: _OpaqueHandle

Handle to an asynchronous GM->L2 prefetch context.

Produced by pl.prefetch.make_context; consumed by pl.prefetch.async_prefetch and pl.prefetch.session.

AsyncEvent

Bases: _OpaqueHandle

Handle to an in-flight asynchronous prefetch completion event.

Produced by pl.prefetch.async_prefetch; consumed by pl.prefetch.wait together with the matching AsyncSession.

AsyncSession

Bases: _OpaqueHandle

Handle to the asynchronous session an AsyncEvent belongs to.

Produced by pl.prefetch.session; consumed by pl.prefetch.wait.

TensorLayout

Bases: Enum

Tensor layout type enumeration.

ND = ... class-attribute instance-attribute

ND layout.

DN = ... class-attribute instance-attribute

DN layout.

NZ = ... class-attribute instance-attribute

NZ layout.

MX_A_ZZ = ... class-attribute instance-attribute

MX Left/A scale GM pack (ZZ).

MX_B_NN = ... class-attribute instance-attribute

MX Right/B scale GM pack (NN).

TensorView

TensorView factory: accepts Expr or int in stride/valid_shape.

TileLayout

Bases: Enum

Tile layout enumeration (shared by blayout and slayout).

none_box = ... class-attribute instance-attribute

No layout constraint.

row_major = ... class-attribute instance-attribute

Row-major layout.

col_major = ... class-attribute instance-attribute

Column-major layout.

PadValue

Bases: Enum

Tile pad mode enumeration.

null = ... class-attribute instance-attribute

No padding.

zero = ... class-attribute instance-attribute

Zero padding.

max = ... class-attribute instance-attribute

Max value padding.

min = ... class-attribute instance-attribute

Min value padding.

CompactMode

Bases: Enum

Partial-tile compact mode enumeration.

null = ... class-attribute instance-attribute

Ordinary non-compact layout.

normal = ... class-attribute instance-attribute

Compact valid-region layout.

TileView

TileView factory: accepts Expr or int in valid_shape/stride.

ND = TensorLayout.ND module-attribute

DN = TensorLayout.DN module-attribute

NZ = TensorLayout.NZ module-attribute

MX_A_ZZ = TensorLayout.MX_A_ZZ module-attribute

MX_B_NN = TensorLayout.MX_B_NN module-attribute

FP4 = DataType.FP4 module-attribute

FP8E4M3FN = DataType.FP8E4M3FN module-attribute

FP8E5M2 = DataType.FP8E5M2 module-attribute

FP8E8M0 = DataType.FP8E8M0 module-attribute

FP16 = DataType.FP16 module-attribute

FP32 = DataType.FP32 module-attribute

BF16 = DataType.BF16 module-attribute

HF4 = DataType.HF4 module-attribute

HF8 = DataType.HF8 module-attribute

INT4 = DataType.INT4 module-attribute

INT8 = DataType.INT8 module-attribute

INT16 = DataType.INT16 module-attribute

INT32 = DataType.INT32 module-attribute

INT64 = DataType.INT64 module-attribute

UINT4 = DataType.UINT4 module-attribute

UINT8 = DataType.UINT8 module-attribute

UINT16 = DataType.UINT16 module-attribute

UINT32 = DataType.UINT32 module-attribute

UINT64 = DataType.UINT64 module-attribute

BOOL = DataType.BOOL module-attribute

INDEX = DataType.INDEX module-attribute

TASK_ID = DataType.TASK_ID module-attribute

TaskId = Scalar[TASK_ID] module-attribute