triton.language.load

triton.language.load(pointer, mask=None, other=None, boundary_check=(), padding_option='', cache_modifier='', eviction_policy='', volatile=False, _semantic=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"}, where ".ca" stands for cache at all levels, ".cg" stands for cache at global level (cache in L2 and below, not L1), and ".cv" means don’t cache and fetch again. see cache operator for more details.) -- changes cache option in NVIDIA PTX

  • eviction_policy (str, optional) -- changes eviction policy in NVIDIA PTX

  • volatile (bool, optional) -- changes volatile option 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 ('', '.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.