helion.language.StackTensor

class helion.language.StackTensor(tensor_like: torch.Tensor, dev_ptrs: torch.Tensor)[source]

This class should not be instantiated directly. It is the result of hl.stacktensor_like(…). It presents a batch of tensors of the same properties (shape, dtype and stride) but reside at different memory locations virtually stacked together.

StackTensor provides a way to perform parallel memory accesses to multiple tensors with a single subscription.

Core Concept:

Instead of performing separate memory operations on each tensor individually, StackTensor allows you to broadcast a single memory operation (hl.load, hl.store, hl.atomic_add, hl.signal, hl.wait etc.) to multiple tensor buffers in parallel. This is particularly useful for batch processing scenarios where the same operation needs to be applied to multiple tensors.

Memory Operation Behavior:

  • Loads: When you index into a StackTensor (e.g., stack_tensor[i]), it performs the same indexing operation on all underlying tensor buffers and returns a new tensor where the results are stacked according to the shape of dev_ptrs.

  • Stores: When you assign to a StackTensor (e.g., stack_tensor[i] = value), the value tensor is “unstacked” - each slice of the value tensor is written to the respective underlying tensor buffer. This is the reverse operation of loading. (e.g. value[j] is writtent to tensor_j[i]).

Shape Semantics:

The StackTensor’s shape is dev_ptrs.shape + tensor_like.shape, where:

  • dev_ptrs.shape becomes the stacking dimensions

  • tensor_like.shape represents the shape of each individual tensor

__init__()

Methods

__init__()

count(value, /)

Return number of occurrences of value.

index(value[, start, stop])

Return first index of value.

new_empty(*args, **kwargs)

Attributes

dev_ptrs

A tensor containing device pointers (memory buffer addresses) to the actual tensors in device memory.

device

dtype

shape

tensor_like

A template host tensor that defines the shape, dtype, and other properties for all tensors in the stack group.
tensor_like: Tensor

A template host tensor that defines the shape, dtype, and other properties for all tensors in the stack group. Must be a Host tensor (created outside of the device loop).

dev_ptrs: Tensor

A tensor containing device pointers (memory buffer addresses) to the actual tensors in device memory. Must be of dtype torch.uint64.

property dtype: dtype
property device: device
property shape: Size
new_empty(*args, **kwargs)[source]
Parameters:
Return type:

Tensor