triton.language.store

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

将数据张量存储到由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。

Parameters:

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

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

  • mask (Block of triton.int1, optional) – 如果mask[idx]为false，则不将value[idx]存储到pointer[idx]

  • boundary_check (tuple of ints, optional) – 整数元组，表示应进行边界检查的维度

  • cache_modifier (str, 可选参数，可选值为 {"", ".wb", ".cg", ".cs", ".wt"}，其中 ".wb" 表示 缓存回写所有一致性级别，".cg" 表示全局缓存，".cs" 表示流式缓存，".wt" 表示透写缓存，详见cache operator获取更多信息。) - 修改NVIDIA PTX中的缓存选项

  • eviction_policy (str, optional, should be one of {"", "evict_first", "evict_last"}) - 修改NVIDIA PTX中的驱逐策略

此函数也可以作为成员函数在tensor上调用， 使用x.store(...)而非 store(x, ...)。