asc2.atomic_max

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

Atomically compute the maximum between local tensor elements and global tensor elements at specified offsets.

Performs an atomic read-modify-write operation, storing the maximum of each element in src and 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 containing comparison values

  • 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 compute the maximum between local tensor and global tensor elements:

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