asc2.atomic_add

asc2.atomic_add(src: LocalTensor, dst: GlobalTensor, offsets: Iterable[PlainValue | int]) → None

Atomically add local tensor elements to a global tensor at specified offsets.

Performs an atomic read-modify-write operation, adding each element of src to the corresponding element in dst at the given offsets.

The supported data types for the inputs are: int16, int32, float16, bfloat16, float32.

Parameters:

  • src – The source local tensor whose elements will be added

  • dst – The destination global tensor

  • offsets – The offsets into the global tensor for each dimension

Raises:

  • TypeError – If src is not a LocalTensor, dst is not a GlobalTensor, or offsets contains non-integer values

  • RuntimeError – If src or dst dtype is not supported, or offsets is empty

Examples

Atomically add local tensor elements to a global tensor:

src = asc2.copy_in(x_gm, [0], [128])
asc2.atomic_add(src, out_gm, [0])