triton.language.store

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

将数据张量存储到由 pointer 定义的内存位置。

  1. 如果 pointer 是单个元素指针，则存储一个标量。在这种情况下

    • mask 也必须是标量，并且

    • boundary_check 和 padding_option 必须为空。

  2. 如果 pointer 是 N 维指针张量，则存储一个 N 维块。在这种情况下

    • mask 会隐式广播到 pointer.shape，并且

    • boundary_check 必须为空。

  3. 如果 pointer 是由 make_block_ptr 定义的块指针，则存储一个数据块。在这种情况下

    • mask 必须为 None，并且

    • 可以指定 boundary_check 来控制越界访问的行为。

value 会隐式广播到 pointer.shape 并类型转换为 pointer.dtype.element_ty。

参数:

  • pointer (triton.PointerType，或 dtype=triton.PointerType 的块) – 存储 value 元素的内存位置

  • value (Block) – 要存储的元素张量

  • mask (triton.int1 块，可选) – 如果 mask[idx] 为 false，则不在 pointer[idx] 处存储 value[idx]

  • boundary_check (整数元组，可选) – 整数元组，指示应执行边界检查的维度

  • cache_modifier (字符串，可选，应为 {“”, “.wb”, “.cg”, “.cs”, “.wt”} 之一，其中 “.wb” 代表缓存写回所有一致性级别，“.cg” 代表缓存全局，“.cs” 代表缓存流，“.wt” 代表缓存直写，详情请参阅 缓存操作符。) – 更改 NVIDIA PTX 中的缓存选项

  • eviction_policy (字符串，可选，应为 {"", "evict_first", "evict_last"} 之一) – 更改 NVIDIA PTX 中的逐出策略

此函数也可以作为 tensor 的成员函数调用，例如 x.store(...) 而不是 store(x, ...)。