triton.language.atomic_and#

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

Performs an atomic logical and at the memory location specified by pointer.

Return the data stored at pointer before the atomic operation.

Parameters:

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_and(...) instead of atomic_and(x, ...).

Example

import triton
import triton.language as tl
import torch


@triton.jit
def atomic_and(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_and(out_ptr0 + (x1), tmp0, xmask)
    tl.store(out_ptr1 + (x1), tmp1, xmask)


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

    block_size = shape[0] * shape[1] // ncore
    split_size = shape[0] // ncore

    val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu()

    pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu()
    pointer_old = torch.full_like(pointer, -10).npu()
    pointer_ref = pointer.clone()

    for i in range(ncore - 1):
        pointer_ref &= val[(i * split_size):((i + 1) * split_size)]

    pointer_ref_last = pointer_ref.clone()
    pointer_ref &= val[((ncore - 1) * split_size):(ncore * split_size)]

    n_elements = shape[0] * shape[1]
    atomic_and[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1])
    assert torch.equal(pointer, pointer_ref)


if __name__ == "__main__":
    test_atomic_and()

DataType Support

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	×	×	×	×	×	√	×	√	×	×	×	×	×	×	×
Ascend 950	×	×	×	×	√	√	√	√	×	×	×	×	×	×	×

Special Restrictions

sem: Ascend does not support “acquire”,”release”,”relaxed”
scope: Ascend does not support “cta”,”sys”

triton.language.atomic_and

Contents

triton.language.atomic_and#