helion.language.split#

helion.language.split(tensor)[source]#

Split the last dimension of a tensor with size two into two separate tensors.

Parameters:: tensor (Tensor) – The input tensor whose last dimension has length two.
Return type:: tuple[Tensor, Tensor]
Returns:: A tuple (lo, hi) where each tensor has the same shape as tensor without its last dimension.

PyTorch alternate forms

The Triton and CuTe backends support torch.chunk(x, 2, dim) for two equal chunks and torch.unbind(x, dim) when the selected dimension has size two. Tensor method forms (x.chunk(...) and x.unbind(...)), including saved bound methods, are also supported. Both operations accept positive and negative axes and return a tuple of two tensors. chunk keeps the input rank; unbind removes the selected axis. They lower through reshape/permute and hl.split.

For an accumulator of shape [tile_m, 128]:

left, right = torch.chunk(acc, 2, dim=-1)  # each is [tile_m, 64]
grouped = acc.reshape(tile_m, 2, 64).permute(0, 2, 1)
left, right = grouped.unbind(dim=-1)  # same contiguous halves
left, right = hl.split(grouped)  # equivalent lowering

Note

These PyTorch alternate forms are available inside device loops on the Triton and CuTe backends. The axis and split size must be known at compile time; use specialize() before the loop when needed. Other tile dimensions can remain symbolic. Because tile dimensions are padded to powers of two, chunk requires a power-of-two split size of at least two. Other chunk counts, uneven chunks, and unbinding dimensions of other sizes raise an unsupported-configuration error. Host-side calls retain normal PyTorch behavior.

On Triton, lowerings that permute rank-compacted tile tensors are unsupported. This includes chunk even on the trailing axis. Unbinding an already trailing size-two axis remains supported with flattened tiles. Direct permute calls have the same restriction: the autotuner skips configurations that compact the input rank. Disable flatten_loops for the affected tile axes to use these permutations.

On CuTe, each thread holds one element of a tile: a split of a loaded tile re-reads each half from memory, while a split of a computed tile exchanges the halves through shared memory, which is refused unless both halves are proved to be staged in the same lane-loop iteration (the halves of a computed row wider than a warp are not). The halves can feed element-wise ops, hl.join, stores of full slices and a torch.unbind of their torch.stack; other consumers, such as a reduction over a half or an operand distributed by another block (bias[tile, :]), are refused.