helion.language.atomic_add#

helion.language.atomic_add(target, index, value, sem='relaxed')[source]#

Atomically add a value to a target tensor.

Performs an atomic read-modify-write that adds value to target[index]. This is safe for concurrent access from multiple threads/blocks.

Parameters:

target (Tensor) – Tensor to update.
index (list[object]) – Indices selecting elements to update. Can include tiles.
value (Tensor | float) – Value(s) to add (tensor or scalar).
sem (str) – Memory ordering semantics. One of "relaxed", "acquire", "release", "acq_rel". Defaults to "relaxed".

Returns:

The previous value(s) stored at target[index] before the update.

Return type:

torch.Tensor

Example

@helion.kernel def global_sum(x: torch.Tensor, result: torch.Tensor) -> torch.Tensor:

for tile in hl.tile(x.size(0)):
hl.atomic_add(result, [0], x[tile].sum())

return result

Notes