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:

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.

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.

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.

Parameters:

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()

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

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.

triton.language.load

Contents

triton.language.load#