triton.language.load

triton.language.load(pointer, mask=None, other=None, boundary_check=(), padding_option='', cache_modifier='', eviction_policy='', volatile=False, care_padding=True, _builder=None)

Return a tensor of data whose values are loaded from memory at location defined by pointer:

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

    • mask and other must also be scalars,

    • other is implicitly typecast to pointer.dtype.element_ty, and

    • boundary_check and padding_option must be empty.

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

    • mask and other are implicitly broadcast to pointer.shape,

    • other is implicitly typecast to pointer.dtype.element_ty, and

    • boundary_check and padding_option must be empty.

  3. If pointer is a block pointer defined by make_block_ptr, a tensor is loaded. In this case:

    • mask and other must be None, and

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

参数:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) -- Pointer to the data to be loaded

  • mask (Block of triton.int1, optional) -- if mask[idx] is false, do not load the data at address pointer[idx] (must be None with block pointers)

  • other (Block, optional) -- if mask[idx] is false, return other[idx]

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

  • padding_option -- should be one of {"", "zero", "nan"}, the padding value to use while out of bounds. "" means an undefined value.

  • cache_modifier (str, optional, should be one of {"", ".ca", ".cg", ".cv"}) -- changes cache option for the load operation. On Ascend A5 SIMT-only, accepts the upstream strings ('', '.ca', '.cg', '.cv'); the backend folds them to a two-level cache / uncache mapping. Effective on global memory; ignored on shared memory. A2/A3 and non-SIMT-only paths ignore the modifier.

  • eviction_policy (str, optional) -- has no effect on Ascend.

  • volatile (bool, optional) -- on Ascend A5 SIMT-only, guarantees memory ordering and prevents the compiler/hardware from hoisting or eliding the load. Effective on both global and shared memory. A2/A3 and non-SIMT-only paths ignore the modifier.

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 ('', '.ca', '.cg', '.cv'); the backend folds them to a two-level cache / uncache mapping ('.ca' -> cache, '.cg' / '.cv' -> 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: bool; on Ascend A5 SIMT-only, guarantees memory ordering and prevents the compiler/hardware from hoisting or eliding the load. Effective on both global and shared memory. A2/A3 and non-SIMT-only paths ignore the modifier.

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