pl.tensor¶
Tensor-level operators — they take and return pl.Tensor values, so they belong in
orchestration or wherever a whole tensor is the unit of work. See
Choosing a namespace.
Tensor operations for PyPTO Language DSL.
This module provides type-safe wrappers around pypto.ir.op.tensor operations that accept and return Tensor types instead of raw Expr/Call objects.
create_tensor = create
module-attribute
¶
create(shape, dtype, layout=TensorLayout.ND, manual_dep=False, init_value=None)
¶
Create a new tensor with specified shape and dtype.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
shape
|
Sequence[IntLike]
|
List of dimension sizes (int or Expr) |
required |
dtype
|
DataType
|
Data type of tensor elements |
required |
layout
|
TensorLayout
|
Tensor layout (default: ND) |
ND
|
init_value
|
int | float | None
|
Removed. Passing anything but |
None
|
manual_dep
|
bool
|
Opt this tensor out of OverlapMap auto-dep tracking for
its entire lifetime. When True, every task that reads or
writes this tensor skips OverlapMap lookup and insert, so the
runtime neither makes the task wait on prior writers nor
registers it as a producer for later readers. Creator retention
(the original This is the tensor-lifetime granularity of opting out of auto-dep tracking. The other two granularities, both orthogonal to this one, are:
All three opt-outs compose with the orthogonal explicit edges
mechanism ( Internally also used by |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the create operation |
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 |
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: reassigningq(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_varsand 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.assemblebuffers —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 storey[off_slice] = tileand tagy, 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 |
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.BYPASSasserts two things abouttensor: 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
ConvertTensorToTileOpsturns it into acachekwarg on eachtile.loadthat reads a declared parameter. - Explicit wins. An explicit
pl.load(..., cache=...)overrides the scope declaration for that one access, in both directions — socache=pl.CachePolicy.DEFAULTopts 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
|
required |
policy
|
CachePolicy
|
|
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])
read(tensor, indices)
¶
Read a scalar value from a tensor at given indices.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor |
required |
indices
|
IntLike | Sequence[IntLike]
|
A single index expression (for 1-D flat access) or a list of index expressions (one per tensor dimension) |
required |
Returns:
| Type | Description |
|---|---|
Scalar
|
Scalar wrapping the read operation |
write(tensor, indices, value)
¶
Write a scalar value into a tensor at given indices.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Destination tensor |
required |
indices
|
IntLike | Sequence[IntLike]
|
A single index expression (for 1-D flat access) or a list of index expressions (one per tensor dimension) |
required |
value
|
Scalar | Expr
|
Scalar value to write (DSL Scalar or raw Expr) |
required |
Returns:
| Type | Description |
|---|---|
Expr
|
The underlying |
Expr
|
typically ignore it; the DSL parser surfaces it as an |
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 |
required |
Returns:
| Type | Description |
|---|---|
Scalar
|
Scalar wrapping the dim operation (INT64) |
slice(tensor, shape, offset, valid_shape=None, drop_dims=None, pad_value=None, clamp=False)
¶
Create a slice of a tensor with new shape and optional valid shape.
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.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
_TensorT
|
Input tensor |
required |
shape
|
Sequence[IntLike]
|
New shape dimensions. Always full-rank — a scalar-indexed axis
contributes a unit dim here and is listed in |
required |
offset
|
Sequence[IntLike]
|
Offset dimensions for the slice |
required |
valid_shape
|
Sequence[IntLike] | None
|
Valid shape dimensions. When omitted, the source's validity under the window is used. Narrows the result; cannot widen it. |
None
|
drop_dims
|
Sequence[int | Expr] | None
|
Optional axes to erase from the result type (numpy-style rank
reduction). Each listed axis must be a static unit dim of |
None
|
pad_value
|
PadValue | int | float | None
|
Optional padding mode for out-of-valid-shape elements.
|
None
|
clamp
|
bool
|
Sanction a window that runs off the end of the source. By default
a slice asserts |
False
|
Returns:
| Type | Description |
|---|---|
_TensorT
|
Tensor wrapping the slice operation |
fillpad(tensor, pad_value=PadValue.zero)
¶
Fill invalid tensor view elements with the specified padding value.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor |
required |
pad_value
|
PadValue | int | float
|
|
zero
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the fillpad operation |
fillpad_expand(tensor, shape, pad_value=PadValue.zero)
¶
Copy a smaller source tensor into a larger destination tensor, padding the rest.
Unlike fillpad (which keeps the same shape and only fills the invalid
view region), 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.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Source tensor |
required |
shape
|
Sequence[IntLike]
|
Destination shape; each dimension must be >= the source dimension |
required |
pad_value
|
PadValue | int | float
|
|
zero
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the fillpad_expand operation (a new, larger tensor). |
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 |
ci(start, shape, dtype=DataType.INT32, descending=False)
¶
Generate a contiguous integer sequence into a tensor.
Equivalent to numpy.arange / torch.arange. Lowers to tile.ci → pto.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 |
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. |
arange = ci
module-attribute
¶
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.random → pto.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. |
matmul(lhs, rhs, out_dtype=None, a_trans=False, b_trans=False, c_matrix_nz=False)
¶
Matrix multiplication with optional transpose.
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, so a lhs stored [K, M] with a_trans=True
against a [K] rhs deduces [M].
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
Tensor
|
Right-hand side tensor |
required |
out_dtype
|
int | DataType | None
|
Output data type. Optional; inferred from the operands when
omitted. When given it must be one of the dtypes the Cube writeback
can produce: FP32/FP16/BF16 for float operands, INT32 for int
operands (see |
None
|
a_trans
|
bool
|
Whether to transpose lhs (requires a 2D+ lhs) |
False
|
b_trans
|
bool
|
Whether to transpose rhs (requires a 2D+ rhs) |
False
|
c_matrix_nz
|
bool
|
C matrix non-zero flag |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the matmul operation |
matmul_acc(acc, lhs, rhs, a_trans=False, b_trans=False, init_cond=None)
¶
Matrix multiplication with accumulation: acc += lhs @ rhs.
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. This is the split-K idiom, and it removes the need to zero
the accumulator or to peel the first K step::
for k0 in pl.pipeline(0, K, K_TILE, stage=2):
acc[t0 : t0 + R, :] = pl.matmul_acc(
acc[t0 : t0 + R, :], x_k, w_k, b_trans=True, init_cond=(k0 == 0)
)
Only 2D operands support the predicate; loop over the batch dimension
instead of passing higher-rank operands alongside init_cond.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
acc
|
Tensor
|
Accumulator tensor |
required |
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
Tensor
|
Right-hand side tensor |
required |
a_trans
|
bool
|
Whether to transpose lhs |
False
|
b_trans
|
bool
|
Whether to transpose rhs |
False
|
init_cond
|
BoolLike | None
|
Optional predicate selecting overwrite over accumulate |
None
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the matmul_acc operation |
mul(lhs, rhs)
¶
Element-wise multiplication of tensor and tensor or scalar.
Automatically selects between tensor.mul (tensor x tensor) and tensor.muls (tensor x scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar (int/float/Tensor/Scalar) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the mul operation |
muls(lhs, rhs)
¶
add(lhs, rhs)
¶
Element-wise addition of tensor and tensor or scalar.
Automatically selects between tensor.add (tensor + tensor) and tensor.adds (tensor + scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar (int/float/Tensor/Scalar) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the add operation |
adds(lhs, rhs)
¶
sub(lhs, rhs)
¶
Element-wise subtraction of tensor and tensor or scalar.
Automatically selects between tensor.sub (tensor - tensor) and tensor.subs (tensor - scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar (int/float/Tensor/Scalar) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the sub operation |
subs(lhs, rhs)
¶
div(lhs, rhs, high_precision=False)
¶
Element-wise division of tensor and tensor or scalar.
Automatically selects between tensor.div (tensor / tensor) and tensor.divs (tensor / scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar (int/float/Tensor/Scalar) |
required |
high_precision
|
bool
|
Whether to select PTOAS's high-precision division mode.
Only available when |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the div operation |
divs(lhs, rhs)
¶
part_add(lhs, rhs)
¶
part_mul(lhs, rhs)
¶
part_max(lhs, rhs)
¶
part_min(lhs, rhs)
¶
fmod(lhs, rhs)
¶
Element-wise truncating remainder of tensor and tensor or scalar.
Automatically selects between tensor.fmod (tensor, tensor) and
tensor.fmods (tensor, scalar) based on the rhs type. The result matches
torch.fmod (the remainder takes the sign of the dividend).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar (int/float/Tensor/Scalar) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the fmod operation |
fmods(lhs, rhs)
¶
maximum(lhs, rhs)
¶
Element-wise maximum of tensor and tensor or scalar.
The conversion pass handles the tensor-vs-tensor / tensor-vs-scalar
dispatch internally — there is no separate maximums front-end op.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the maximum operation |
minimum(lhs, rhs)
¶
Element-wise minimum of tensor and tensor or scalar.
The conversion pass handles the tensor-vs-tensor / tensor-vs-scalar
dispatch internally — there is no separate minimums front-end op.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the minimum operation |
cmp(lhs, rhs, cmp_type=0)
¶
Element-wise comparison of tensor and tensor or scalar (returns 0/1 tensor).
The conversion pass handles the tensor-vs-tensor / tensor-vs-scalar
dispatch internally — there is no separate cmps front-end op.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor |
required |
rhs
|
int | float | Tensor | Scalar | Expr
|
Right-hand side tensor or scalar |
required |
cmp_type
|
int
|
Comparison type code (0=eq, 1=ne, 2=lt, 3=le, 4=gt, 5=ge) |
0
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor of 0/1 with the same shape and dtype as |
and_(lhs, rhs)
¶
Element-wise bitwise AND of tensor and tensor or scalar.
Automatically selects between tensor.and (tensor & tensor) and tensor.ands (tensor & scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Tensor | Scalar | Expr
|
Right-hand side tensor or integer scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the and operation |
ands(lhs, rhs)
¶
or_(lhs, rhs)
¶
Element-wise bitwise OR of tensor and tensor or scalar.
Automatically selects between tensor.or (tensor | tensor) and tensor.ors (tensor | scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Tensor | Scalar | Expr
|
Right-hand side tensor or integer scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the or operation |
ors(lhs, rhs)
¶
xor(lhs, rhs)
¶
Element-wise bitwise XOR of tensor and tensor or scalar.
Automatically selects between tensor.xor (tensor ^ tensor) and tensor.xors (tensor ^ scalar) based on the rhs type.
Unlike pl.tile.xor, there is no tmp parameter: the scratch buffer
that pto.txor needs is allocated during Tensor-to-Tile lowering.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Tensor | Scalar | Expr
|
Right-hand side tensor or integer scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the xor operation |
xors(lhs, rhs)
¶
not_(input)
¶
shl(lhs, rhs)
¶
Element-wise bitwise left shift of tensor by tensor or scalar.
Automatically selects between tensor.shl (tensor << tensor) and tensor.shls (tensor << scalar) based on the rhs type.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Tensor | Scalar | Expr
|
Shift amount as a tensor or integer scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the shl operation |
shls(lhs, rhs)
¶
Element-wise bitwise left shift of tensor by scalar.
Note
The shift amount must be zero or positive. A negative constant is rejected when the op is built; a negative value only known at runtime is undefined behaviour on the hardware.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Expr | Scalar
|
Shift amount; must be >= 0 |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the shls operation |
shr(lhs, rhs)
¶
Element-wise bitwise right shift of tensor by tensor or scalar.
Automatically selects between tensor.shr (tensor >> tensor) and tensor.shrs (tensor >> scalar) based on the rhs type. The shift is arithmetic for signed dtypes and logical for unsigned ones, matching the tile ops and the underlying ISA.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Tensor | Scalar | Expr
|
Shift amount as a tensor or integer scalar |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the shr operation |
shrs(lhs, rhs)
¶
Element-wise bitwise right shift of tensor by scalar.
Note
The shift amount must be zero or positive. A negative constant is rejected when the op is built; a negative value only known at runtime is undefined behaviour on the hardware.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lhs
|
Tensor
|
Left-hand side tensor (integer dtype) |
required |
rhs
|
int | Expr | Scalar
|
Shift amount; must be >= 0 |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the shrs operation |
row_max(input)
¶
row_sum(input)
¶
row_min(input)
¶
row_prod(input)
¶
col_sum(input)
¶
col_max(input)
¶
col_min(input)
¶
col_prod(input)
¶
row_argmax(input)
¶
row_argmin(input)
¶
col_argmax(input)
¶
col_argmin(input)
¶
row_expand(target, row_vec)
¶
Row-wise expansion: expand row_vec [M, 1] to target shape [M, N].
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
target
|
Tensor
|
Target tensor defining output shape (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector to expand (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand operation |
row_expand_mul(tensor, row_vec)
¶
Row-wise broadcast multiplication: tensor[i,:] * row_vec[i,0].
Multiplies each row of the tensor by the corresponding row vector value,
for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_mul operation |
row_expand_div(tensor, row_vec)
¶
Row-wise broadcast division: tensor[i,:] / row_vec[i,0].
Divides each row of the tensor by the corresponding row vector value,
for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_div operation |
row_expand_add(tensor, row_vec)
¶
Row-wise broadcast addition: tensor[i,:] + row_vec[i,0].
Adds a row vector to each row of the tensor, for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_add operation |
row_expand_sub(tensor, row_vec)
¶
Row-wise broadcast subtraction: tensor[i,:] - row_vec[i,0].
Subtracts a row vector from each row of the tensor, for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_sub operation |
row_expand_max(tensor, row_vec)
¶
Row-wise broadcast maximum: max(tensor[i,:], row_vec[i,0]).
Takes the element-wise maximum of each row and the row vector value,
for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_max operation |
row_expand_min(tensor, row_vec)
¶
Row-wise broadcast minimum: min(tensor[i,:], row_vec[i,0]).
Takes the element-wise minimum of each row and the row vector value,
for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_min operation |
row_expand_expdif(tensor, row_vec)
¶
Row-wise exp-diff: exp(tensor[i,:] - row_vec[i,0]).
Computes the exponential of the per-row difference, for all i.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
row_vec
|
Tensor
|
Row vector providing per-row scalar (TensorType [M, 1]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the row_expand_expdif operation |
col_expand_mul(tensor, col_vec)
¶
Column-wise broadcast multiplication: tensor[:,j] * col_vec[0,j].
Multiplies each column of the tensor by the corresponding column vector
value, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_mul operation |
col_expand(tensor, col_vec)
¶
Column-wise expansion: expand col_vec [1, N] to target shape [M, N].
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Target tensor defining output shape (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector to expand (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand operation |
col_expand_div(tensor, col_vec)
¶
Column-wise broadcast division: tensor[:,j] / col_vec[0,j].
Divides each column of the tensor by the corresponding column vector
value, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_div operation |
col_expand_sub(tensor, col_vec)
¶
Column-wise broadcast subtraction: tensor[:,j] - col_vec[0,j].
Subtracts a column vector from each column of the tensor, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_sub operation |
col_expand_add(tensor, col_vec)
¶
Column-wise broadcast addition: tensor[:,j] + col_vec[0,j].
Adds a column vector to each column of the tensor, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_add operation |
col_expand_max(tensor, col_vec)
¶
Column-wise broadcast maximum: max(tensor[:,j], col_vec[0,j]).
Takes the element-wise maximum of each column and the column vector
value, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_max operation |
col_expand_min(tensor, col_vec)
¶
Column-wise broadcast minimum: min(tensor[:,j], col_vec[0,j]).
Takes the element-wise minimum of each column and the column vector
value, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_min operation |
col_expand_expdif(tensor, col_vec)
¶
Column-wise exp-diff: exp(tensor[:,j] - col_vec[0,j]).
Computes the exponential of the per-column difference, for all j.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (TensorType [M, N]) |
required |
col_vec
|
Tensor
|
Column vector providing per-column scalar (TensorType [1, N]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the col_expand_expdif operation |
expands(target, scalar)
¶
expand_clone(src, target)
¶
exp(input)
¶
log(input, high_precision=False)
¶
sin(input)
¶
cos(input)
¶
neg(input)
¶
abs(input)
¶
recip(input, high_precision=False)
¶
sqrt(input)
¶
rsqrt(input, high_precision=False)
¶
Element-wise reciprocal square root operation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
input
|
Tensor
|
Input tensor |
required |
high_precision
|
bool
|
If True, lower to the higher-precision PTO path. The compiler allocates a scratch buffer during tensor-to-tile conversion. |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the rsqrt operation |
cast(input, target_type, mode='round', *, saturation_mode=None)
¶
Type casting operation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
input
|
Tensor
|
Input tensor |
required |
target_type
|
int | DataType
|
Target data type |
required |
mode
|
str | int
|
Rounding mode — string name ("none", "rint", "round", "floor", "ceil", "trunc", "odd") or int (0–6) |
'round'
|
saturation_mode
|
str | int | None
|
Destination saturation — |
None
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the cast operation |
assemble(target, source, offset, *, atomic=AtomicType.None_)
¶
Write/update tensor values at specified offset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
target
|
_TensorT
|
Target tensor to update |
required |
source
|
Tensor
|
Source tensor to write |
required |
offset
|
Sequence[IntLike]
|
Offset dimensions for where to write |
required |
atomic
|
AtomicType
|
Combine mode for the write. NOTE: atomic-add accumulation order across cores is not fixed, so floating-point results are non-deterministic. The target must be zero-initialised before the kernel runs. Supported 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_
|
Returns:
| Type | Description |
|---|---|
_TensorT
|
Tensor wrapping the assemble operation |
concat(src0, src1)
¶
reshape(tensor, shape)
¶
Reshape tensor to new shape.
The valid region is carried through, never widened: the result holds real
data in exactly the cells the input did. See pl.reshape for the cases
that always map.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor |
required |
shape
|
Sequence[IntLike]
|
New shape dimensions |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the reshape operation |
Raises:
| Type | Description |
|---|---|
ValueError
|
If the element count changes, or if the input holds real
data in only part of its buffer and no origin-anchored region of
|
reinterpret_view(data, dtype, *, shape=None)
¶
Reinterpret a tensor over the same bytes with a different dtype.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
Tensor
|
Input tensor. |
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 |
|---|---|
Tensor
|
Tensor wrapping the zero-copy reinterpret-view operation. |
transpose(tensor, axis1, axis2)
¶
Transpose tensor by swapping two axes.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor |
required |
axis1
|
int
|
First axis to swap (supports negative indexing) |
required |
axis2
|
int
|
Second axis to swap (supports negative indexing) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the transpose operation |
view(tensor, shape=None, valid_shape=None, *, layout=None)
¶
Reinterpret a tensor over the same physical memory.
At least one of shape or layout must be provided: shape derives
canonical strides for the requested shape, layout derives the canonical
ND/DN layout view. The result is a zero-copy view over the same physical
memory, and its target shape must have rank at least 1.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
_TensorT
|
Source tensor. |
required |
shape
|
Sequence[IntLike] | None
|
New shape for the view. Must be product-preserving unless symbolic dimensions are present. Rank-zero views are not supported. |
None
|
valid_shape
|
Sequence[IntLike] | None
|
Explicit valid dimensions for a packed ND
leading-dimension collapse to 2D or contiguous-prefix linear
collapse to |
None
|
layout
|
TensorLayout | None
|
Target |
None
|
Returns:
| Type | Description |
|---|---|
_TensorT
|
Tensor wrapping the view operation. |
Raises:
| Type | Description |
|---|---|
ValueError
|
If the requested shape/layout is missing, unsupported, or inconsistent with the source tensor metadata. |
scatter_update(input, *args, **kwargs)
¶
Update input tensor rows at positions specified by 2D index with values from src.
Supports two rank variants:
- 2D:
input [rows, d],src [b*s, d],index [b, s] - 4D:
input [blockNum, blockSize, 1, d],src [b, s, 1, d],index [b, s]
For each (i, j), row input[index[i*s + j]] receives row src[i*s + j]
(linear layout).
Accepts the same flexible call shapes as the IR builder
pypto.ir.op.tensor.scatter_update:
scatter_update(input, dim, index, src)scatter_update(input, index, src, dim=-2)scatter_update(input, dim, index=..., src=...)
Tensor / Scalar wrappers are unwrapped before forwarding so the IR
builder receives raw Expr operands.
set_validshape(tensor, valid_rows, valid_cols)
¶
Update valid-shape metadata of a tensor without data movement.
.. note::
Prefer expressing the extent at its source where possible —
pl.load(..., valid_shape=...) or a slice's valid_shape= — and use
this to pin an extent the type deducer cannot infer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tensor
|
Tensor
|
Input tensor (must be 2D) |
required |
valid_rows
|
IntLike
|
Number of valid rows (int or Scalar[INDEX]) |
required |
valid_cols
|
IntLike
|
Number of valid columns (int or Scalar[INDEX]) |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor with updated valid_shape metadata |
sort32(src, idx)
¶
Sort fixed 32-element blocks with explicit index tensor (tensor-level).
Tensor-level counterpart of pl.tile.sort32. Sorts 32-element blocks in
src, permuting idx alongside. Returns an 8-byte value-index-pair tensor;
its last dimension is 2x the input width for FP32 and 4x for FP16.
For FP16 src: initialize idx with [0, 1, 2, ..., 31] per block. For FP32 src: initialize idx with [0, 2, 4, ..., 62] per block.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
src
|
Tensor
|
Input value tensor (FP16 or FP32) |
required |
idx
|
Tensor
|
Input index tensor with sequential offsets |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor wrapping the dtype-dependent expanded sort32 output |
mrgsort(src0, src1=None, src2=None, src3=None, *, exhausted=False, block_len=None)
¶
Merge sort — format1 (single-list) or format2 (2-4 way merge), tensor-level.
Tensor-level counterpart of pl.tile.mrgsort. The scratch tmp and
executed tiles required by the tile-level op are synthesized
automatically during conversion as local Vec tiles — users do not pass them.
Format1 usage (keyword block_len): out = mrgsort(src, block_len=64)
Format2 usage
out = mrgsort(src0, src1) # 2-way out = mrgsort(src0, src1, src2) # 3-way out = mrgsort(src0, src1, src2, src3) # 4-way out = mrgsort(src0, src1, exhausted=True)
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
src0
|
Tensor
|
For format1: input tensor with pre-sorted runs (FP16 or FP32). For format2: first sorted input tensor. |
required |
src1
|
Tensor | None
|
(format2) Second sorted input tensor. |
None
|
src2
|
Tensor | None
|
(format2, optional) Third sorted input tensor (3-way or 4-way). |
None
|
src3
|
Tensor | None
|
(format2, optional) Fourth sorted input tensor (4-way only). |
None
|
exhausted
|
bool
|
(format2) If True, marks inputs as exhausted (default: False). |
False
|
block_len
|
int | Scalar | None
|
(format1, keyword-only) Run length, must be multiple of 64. |
None
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor with merged sorted elements |
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 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 |
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 |
None
|
kvalue
|
int | Scalar | Expr | None
|
(compare form, keyword-only) Scalar threshold (dtype must match |
None
|
cmp_mode
|
str | int | None
|
(compare form, keyword-only) |
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 |
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 |
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 — |
Mat
|
col_off
|
int
|
Column start offset within each src row (default 0). |
0
|
is_trans
|
bool
|
Transpose for matmul B-operand layout (requires |
False
|
is_b_matrix
|
bool
|
Hint that the result feeds matmul as the B matrix. |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor of shape |
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)
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
|
required |
dtype
|
DataType
|
Element dtype (matches the gathered |
required |
transpose
|
bool
|
Allocate the transposed Mat (ZN) layout — required when filling
with |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor of shape |
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)
gather_row(acc, 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).
Per-row primitive for a kernel-driven paged gather into L1 — the flexible
counterpart to paged_gather: the caller computes the physical
src_offset (block-table lookup, multi-source selection, invalid clamping)
and the dst_offset slot itself, so arbitrary gather logic stays in the
kernel. DMAs src straight into acc (GM -> L1, no tmov); the
returned tensor feeds pl.matmul directly.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
acc
|
Tensor
|
On-chip accumulator from |
required |
src
|
Tensor
|
Source pool in GM. |
required |
dst_offset
|
Sequence[IntLike]
|
|
required |
src_offset
|
Sequence[IntLike]
|
|
required |
shapes
|
Sequence[IntLike]
|
GM row window |
required |
valid_shape
|
Sequence[IntLike] | None
|
How much of that window to actually transfer, defaulting to
all of it. May hold runtime |
None
|
transpose
|
bool
|
Place the GM row |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
Tensor aliasing |
scatter(input, dim=None, index=None, src=None, *, mask_pattern=None, dst=None)
¶
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 |
None
|
src
|
Tensor | None
|
(index form) Source values tensor (same dtype as |
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; |
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)
alloc(memory_space, size)
¶
Stub for the internal tensor.alloc IR operation.
This function is never called in user-written DSL code. It is emitted by the C++ python-printer after the InitMemRef pass and must be importable so that the printed source is valid Python.
The result is a base Ptr (allocation identity token): the printer
annotates the assignment target as pl.Ptr, matching the IR design
where tensor.alloc Calls carry PtrType.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
memory_space
|
MemorySpace
|
Space the allocation lives in, as resolved by InferTileMemorySpace. |
required |
size
|
int
|
Allocation size in bytes. |
required |
Returns:
| Type | Description |
|---|---|
Ptr
|
A |
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 |
Example
block_idx = pl.tensor.get_block_idx()
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 |
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 |
Example
block_idx = pl.tensor.get_block_idx() block_num = pl.tensor.get_block_num()