triton.language.store

triton.language.store(pointer, value, mask=None, boundary_check=(), cache_modifier='', eviction_policy='', _semantic=None)

Store a tensor of data into memory locations defined by pointer.

  1. If pointer is a single element pointer, a scalar is stored. In this case:

    • mask must also be scalar, and

    • boundary_check and padding_option must be empty.

  2. If pointer is an N-dimensional tensor of pointers, an N-dimensional block is stored. In this case:

    • mask is implicitly broadcast to pointer.shape, and

    • boundary_check must be empty.

  3. If pointer is a block pointer defined by make_block_ptr, a block of data is stored. In this case:

    • mask must be None, and

    • boundary_check can be specified to control the behavior of out-of-bound access.

value is implicitly broadcast to pointer.shape and typecast to pointer.dtype.element_ty.

参数:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) -- The memory location where the elements of value are stored

  • value (Block) -- The tensor of elements to be stored

  • mask (Block of triton.int1, optional) -- If mask[idx] is false, do not store value[idx] at pointer[idx]

  • boundary_check (tuple of ints, optional) -- tuple of integers, indicating the dimensions which should do the boundary check

  • cache_modifier (str, optional, should be one of {"", ".wb", ".cg", ".cs", ".wt"}, where ".wb" stands for cache write-back all coherent levels, ".cg" stands for cache global, ".cs" stands for cache streaming, ".wt" stands for cache write-through, see cache operator for more details.) -- changes cache option in NVIDIA PTX

  • eviction_policy -- changes eviction policy in NVIDIA PTX

Example

import torch
import triton
import triton.language as tl


@triton.jit
def kernel(output_ptr, x_ptr, N: tl.constexpr):
    """
    Test tl.load/tl.store with a 1D tensor pointer:
    1. Load data via pointer + offset
    2. Store loaded data to output
    """
    pid = tl.program_id(0)
    offset = pid * N + tl.arange(0, N)[:]
    data = tl.load(x_ptr + offset)
    tl.store(output_ptr + offset, data)


def test_load_store():
    N = 128
    total = 4 * N
    x = torch.arange(total, dtype=torch.float32).npu()
    output = torch.zeros_like(x)

    # Run Triton kernel
    kernel[(4, )](output, x, N)

    assert torch.allclose(output.cpu(), x.cpu())


if __name__ == "__main__":
    test_load_store()

Special Restrictions

  • DataType: Ascend A2/A3 does not support uint16/uint32/uint64/fp64, Ascend 950 does not support fp64 (hardware limitation).

  • cache_modifier: Ascend A5 SIMT-only accepts the Triton upstream strings ('', '.wb', '.cg', '.cs', '.wt'); the backend folds them to a two-level cache / uncache mapping ('.wb' -> cache, '.cg' / '.cs' / '.wt' -> uncache). Effective on global memory; ignored on shared memory. A2/A3 and non-SIMT-only paths ignore the modifier.

  • eviction_policy: has no effect on Ascend.

  • volatile: not accepted (the Triton upstream tl.store signature does not expose a volatile= keyword).

  • Compatibility issues with branch and loop statements: Complex pointer and mask calculations involving branches or loops may cause compilation failures.