triton.language.atomic_add

triton.language.atomic_add(pointer, val, mask=None, sem=None, scope=None, _semantic=None)

Performs an atomic add at the memory location specified by pointer.

Return the data stored at pointer before the atomic operation.

参数:

  • pointer (Block of dtype=triton.PointerDType) -- The memory locations to operate on

  • val (Block of dtype=pointer.dtype.element_ty) -- The values with which to perform the atomic operation

  • sem (str, optional) -- Specifies the memory semantics for the operation. Acceptable values are "acquire", "release", "acq_rel" (stands for "ACQUIRE_RELEASE"), and "relaxed". If not provided, the function defaults to using "acq_rel" semantics.

  • scope (str, optional) -- Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are "gpu" (default), "cta" (cooperative thread array, thread block), or "sys" (stands for "SYSTEM"). The default value is "gpu".

This function can also be called as a member function on tensor, as x.atomic_add(...) instead of atomic_add(x, ...).

Example

import triton
import triton.language as tl
import torch


@triton.jit
def atomic_add(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr):
    xoffset = tl.program_id(0) * BLOCK_SIZE
    xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:]
    yindex = tl.arange(0, BLOCK_SIZE)[:]
    xmask = xindex < n_elements
    x0 = xindex
    x1 = yindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tl.atomic_add(out_ptr0 + (x1), tmp0, xmask)
    tl.store(out_ptr1 + (x1), tmp1, xmask)


def test_atomic_add():
    dtype, shape, ncore = ['int32', (32, 32), 2]

    block_size = shape[0] * shape[1] / ncore
    split_size = shape[0] // ncore
    x0_value = 3
    x0 = torch.full(shape, x0_value, dtype=eval(f'torch.{dtype}')).npu()
    x1 = torch.full((split_size, shape[1]), 2, dtype=eval(f'torch.{dtype}')).npu()
    y = torch.full((split_size, shape[1]), -10, dtype=eval(f'torch.{dtype}')).npu()

    y_ref = x1 + 0
    x1_ref = x1 + ncore * x0_value

    n_elements = shape[0] * shape[1]
    atomic_add[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1])
    assert torch.equal(x1, x1_ref)


if __name__ == "__main__":
    test_atomic_add()

Special Restrictions

  • DataType: Ascend does not support fp64 (hardware limitation).

  • sem: Ascend does not support "acquire","release","relaxed"

  • scope: Ascend does not support "cta","sys"