triton.language.store

triton.language.store(pointer, value, mask=None, boundary_check=(), cache_modifier='', eviction_policy='')
Store a tensor of data into memory locations defined by pointer:

  1. pointer could be a single element pointer, then a scalar will be stored - mask must be scalar too - boundary_check and padding_option must be empty

  2. pointer could be element-wise tensor of pointers, in which case: - mask is implicitly broadcast to pointer.shape - boundary_check must be empty

  3. or pointer could be a block pointer defined by make_block_ptr, in which case: - mask must be None - 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.

Parameters:

  • 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) – changes cache option in NVIDIA PTX

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