triton.language.atomic_min

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

Performs an atomic min 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_min(...) instead of atomic_min(x, ...).

Example

import triton
import triton.language as tl
import torch


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


def test_atomic_min():
    dtype = 'int32'
    shape = (3, 1)

    x0 = torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)).npu()
    x1 = torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)).npu()

    x1_ref = torch.minimum(x0, x1)

    n_elements = shape[0] * shape[1]
    triton_atomic_min[shape[0], 1, 1](x0, x1, n_elements, BLOCK_SIZE=shape[1])
    assert torch.equal(x1, x1_ref)


if __name__ == "__main__":
    test_atomic_min()

Special Restrictions

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

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