Language Reference

meTile traces a decorated Python function with symbolic arguments, then compiles the recorded operations into a GPU program. Because Python controls the trace, a Python if or range must depend on compile-time values. Use where for per-element selection and tile_range for a traced loop.

The snippets below go inside a kernel. For a complete program with launch code and numerical checks, start with Your First Kernel or Matrix Multiply (GEMM).

Define and launch a kernel

@metile.kernel
def my_kernel(input_ptr, output_ptr, count, BLOCK: metile.constexpr):
    ...

my_kernel[(metile.cdiv(count, 256),)](
    input_buffer, output_buffer, count, BLOCK=256
)

Buffer arguments become device pointers with their storage dtype. For example, float32 buffers use device float* and float16 buffers use device half*. Python integers and floats become 32-bit integer and float scalar arguments. Pass runtime arguments in signature order and constexpr arguments by keyword, as shown above.

The launcher recognizes the actual metile.constexpr annotation object. Avoid postponed/string annotations on kernel parameters: a file containing from __future__ import annotations does not currently preserve this identity check.

A launch grid is a tuple of program counts, not element counts or threads per group. Each program runs as a Metal threadgroup. Elementwise kernels support one, two or three grid axes; the standard GEMM path uses a two-dimensional output grid.

kernel[grid].prepare(...) launches, synchronizes, and returns a dispatcher bound to those resources and scalar values. Preparation is not a compile-only operation. If you plan to reuse a dispatcher, prefer explicit Buffer arguments; see Memory Model.

Program identity and indices

metile.program_id(axis)

Return the program’s threadgroup coordinate along axis 0, 1, or 2.

metile.arange(start, end, *, layout=None)

For integer start and end, create the half-open sequence [start, end). Its length is a compile-time tile size. For clarity, write offset + metile.arange(0, BLOCK) when the origin is dynamic.

The current overload also accepts a traced scalar start; in that form, the second argument is the tile’s length, not an endpoint.

An optional ThreadLayout assigns logical elements to physical threads and registers. See Thread Ownership and Software Staging for supported geometry.

metile.cdiv(numerator, denominator)

Ceiling division, commonly used on the host to compute a launch grid.

metile.next_power_of_2(value)

Return the smallest power of two greater than or equal to the value. Use it on host integers when selecting a tile size.

Tensor declarations

metile.tensor(pointer, *, shape, strides=None, access='readwrite', block_shape=None, address_space=None)

Declare a logical view of existing storage. Shapes and strides are tuples; strides are measured in elements and default to contiguous row-major storage. Access is "read", "write", or "readwrite". An explicit address space must match the pointer’s allocation.

Declare views near the start of the kernel, then use view.load(indices, other=0) and view.store(indices, value). These check logical coordinate bounds. The caller must still provide enough storage and keep index/offset arithmetic within signed 32-bit range.

block_shape selects supported two-dimensional matrix-tile access. It does not allocate memory or set thread/register ownership. See Tensor and Memory Contracts for the full bounds, dtype, and layout contract.

inputs = metile.tensor(input_ptr, shape=(count,), access="read")
outputs = metile.tensor(output_ptr, shape=(count,), access="write")
positions = metile.program_id(0) * BLOCK + metile.arange(0, BLOCK)
outputs.store((positions,), inputs.load((positions,)) * 2.0)

Raw memory operations

metile.load(ptr, mask=None)

Load elements from a pointer expression. Masked-off elements read zero. Without a mask, the caller must ensure every address is valid.

metile.store(ptr, value, mask=None)

Store through a pointer expression. Masked-off stores are skipped.

positions = metile.program_id(0) * BLOCK + metile.arange(0, BLOCK)
valid = positions < count
values = metile.load(input_ptr + positions, mask=valid)
metile.store(output_ptr + positions, values, mask=valid)

metile.tile_load(ptr, row_offset, col_offset, stride, shape)

Load a two-dimensional tile. stride is the row stride in elements; shape is the tile’s (rows, columns). This legacy interface does not declare full tensor bounds. New matrix kernels should use tensor(..., block_shape=...) so lowering can validate dimensions and memory layout.

metile.tile_store(ptr, row_offset, col_offset, stride, value, shape)

Store a matrix tile using the same legacy layout convention.

Accumulators and matrix multiply

metile.zeros(shape, dtype='f32')

Create a zero-valued tile, commonly a float32 matrix accumulator.

metile.dot(left, right, accumulator)

Return accumulator + left @ right. The left and right tile shapes must have compatible reduction dimensions. Supported matrix kernels accumulate in float32.

Lowering selects SIMDgroup matrix operations or Metal 4 tensor operations according to device/toolchain support and the execution schedule. The direct NAX path has additional restrictions. See Tile Operations and Hardware Mapping and Matrix Multiply (GEMM).

Control flow and scalar state

metile.tile_range(start, end, step=1, num_stages=1)

Record a GPU loop with an exclusive end. Its body runs once during tracing, producing the operations repeated at execution time. A positive compile-time integer step is the usual tiling pattern.

num_stages requests staging on a supported lowering path. It does not guarantee that memory access and computation overlap.

partial_sum = 0.0
for start in metile.tile_range(0, count, BLOCK):
    positions = start + metile.arange(0, BLOCK)
    partial_sum = partial_sum + inputs.load((positions,))
total = metile.sum(partial_sum)

metile.scalar(value, dtype=None)

Make explicit scalar SSA state for a loop-carried recurrence. This is useful when a state value is scalar per physical thread, as in Decode Attention.

Math and conversions

Arithmetic such as +, -, *, /, comparisons, and supported integer bit operations builds elementwise IR. Use & and | to combine traced boolean masks; Python and and or do not describe elementwise operations.

Function	Behavior
metile.exp(value)	Exponential
metile.fast_exp(value)	Metal’s fast exponential; accuracy differs from the regular intrinsic
metile.log(value)	Natural logarithm
metile.sqrt(value)	Square root
metile.abs(value)	Absolute value
metile.tanh(value)	Hyperbolic tangent
metile.where(condition, left, right)	Select a value per element; it is not short-circuit control flow
metile.maximum(left, right), metile.minimum(left, right)	Elementwise maximum or minimum
metile.cast(value, dtype)	Convert a scalar or tile, for example to "f32"

where evaluates its inputs before choosing one. Wrapping an unmasked load in where therefore does not make the load safe. Mask the load directly or use a tensor view.

Reductions

metile.sum(value)

metile.max(value)

metile.min(value)

Reduce a supported tile to a scalar sum, maximum, or minimum. Lowering may combine register-local work, SIMDgroup operations, and shared memory between SIMDgroups.

Choose padding values appropriate to the reduction, and mask any transformed padding that must not contribute. See Softmax for maximum and exponential-sum reductions over a partial tile. Floating-point reduction order may differ from a CPU reference.

Thread and SIMDgroup operations

These operations expose physical execution details. On the supported Apple GPU paths, a SIMDgroup contains 32 threads.

metile.thread_id()

Thread index inside the threadgroup.

metile.simd_lane_id()

Lane index inside its SIMDgroup, from 0 through 31.

metile.simd_shuffle_xor(value, mask)

Exchange a value with the lane whose index is XORed with mask.

metile.simd_broadcast(value, lane)

Broadcast a value from one lane inside the SIMDgroup.

metile.simd_sum(value)

metile.simd_max(value)

Native reduction across the current SIMDgroup, not the whole threadgroup.

metile.simdgroup_role(role, num_roles=2, num_sgs=0)

Context manager recording work for a subset of SIMDgroups. num_sgs=0 requests an even division among roles. A role region does not establish ordering or data dependencies with another role. See Fused Activations and SIMDgroup Roles.

metile.shared(size, dtype='f32')

Allocate uninitialized threadgroup memory. The size and dtype are fixed for the specialization.

metile.barrier()

Synchronize the threadgroup and order threadgroup memory accesses. All participating threads must reach the barrier. It does not synchronize separate threadgroups or provide a device-memory producer/consumer protocol.

Ownership and schedules

metile.ThreadLayout describes physical ownership of tile elements. metile.convert_layout(value, layout) preserves logical values while the compiler inserts required redistribution. See Thread Ownership and Software Staging for the supported operations and communication rules.

metile.tile_swizzle(pid_m, pid_n, pattern='auto', block_size=4)

Record a traversal for output tiles. Patterns include "auto", "linear", "diagonal", "morton", "hilbert", and grouped traversals. Unsupported panel geometries fall back to a valid traversal. See Tile Operations and Hardware Mapping for the panel-size constraints.

Pass SCHEDULE=metile.Schedule(...) at launch to require a backend, SIMDgroup geometry, staging or supported vector width. The compiler checks those requirements. Use the prepared dispatch’s explain() report to see the selected schedule and what lowering produced; see Execution Schedules and Fusion.

Host APIs and specialized integrations are listed in API Reference. They are not part of the traced kernel operations described here.