Memory and synchronization

How kernels address memory with masks, pointer tiles, and tensor descriptors, and when their results become visible to the CPU and to other frameworks.

Masks

A mask is a boolean tile, usually a comparison such as offs < n. tl.load reads only where the mask is true and returns other elsewhere. tl.store and the atomics write only where it's true.

Note

The GPU doesn't check bounds on pointer accesses. To find an out-of-bounds access, run the kernel in the interpreter, which raises IndexError with the kernel's line.

Pointers and tensor descriptors

Enceladus gives you two ways to address memory:

Approach	How it works	Use it for
Pointer tiles	`x_ptr + offs` gives a tile of pointers. You compute every address and mask every access.	Elementwise ops, reductions, gathers, strided or transposed access
Tensor descriptors	`tl.make_tensor_descriptor(ptr, shape, strides, block_shape)` loads and stores whole blocks at element offsets and checks its own bounds. The innermost stride must be 1.	Matrix multiplication and attention. `tl.dot` reads descriptor loads straight from device memory.

descriptors

a = tl.make_tensor_descriptor(a_ptr, [M, K], [stride_am, 1], [BM, BK])
tile = a.load([pid_m * BM, k])          # zeros where the block leaves the array
c.store([pid_m * BM, pid_n * BN], acc)  # skips elements outside the array

Streams and synchronization

Launches go to a stream, an ordered queue of GPU work like a CUDA stream. Kernels on the stream run in launch order, so each kernel sees the writes of the kernels before it. The stream submits a Metal command buffer every 64 launches, or sooner when the CPU needs results.

To wait for all launched work, call enceladus.synchronize(). A failed tl.device_assert or command buffer surfaces at the next synchronization as enceladus.DeviceAssertionError or enceladus.MetalError. The stream isn't thread-safe, so launch kernels from one thread at a time.

When a launch waits depends on the arrays that you pass, as the following table shows:

Array arguments	Where the kernel runs	When results are visible
`enceladus.Tensor` only	Enceladus's stream, asynchronously	After a sync. `Tensor.numpy()`, `tolist()`, `print()`, and `np.asarray()` sync for you.
NumPy arrays, alone or with tensors	Enceladus's stream; the launch waits	When the launch returns. With `enceladus.async_numpy(True)`, after `enceladus.synchronize()`.
PyTorch MPS tensors only	PyTorch's MPS stream, through `torch.mps.compile_shader`	In order with surrounding PyTorch operations
MLX arrays	Enceladus's stream; inputs are evaluated first and the launch waits	When the launch returns
PyTorch tensors mixed with other kinds	Enceladus's stream, between syncs of both streams	When the launch returns

A synchronized launch costs about 70-100 µs. To batch many launches over NumPy arrays, turn the wait off and synchronize once:

async_numpy.py

x = np.zeros(10_000, np.float32)
enceladus.async_numpy(True)
for _ in range(100):
inc_kernel[(enceladus.cdiv(x.size, 1024),)](x, x.size, BLOCK=1024)
enceladus.synchronize()  # required before reading x
enceladus.async_numpy(False)
assert (x == 100).all()