triton.language.load

triton.language.load(pointer, mask=None, other=None, boundary_check=(), padding_option='', cache_modifier='', eviction_policy='', volatile=False, _semantic=None)

返回一个张量，其值是从 pointer 定义的内存位置加载的。

  1. 如果 pointer 是单元素指针，则加载标量。在这种情况下：

    • mask 和 other 也必须是标量，

    • other 会被隐式转换为 pointer.dtype.element_ty 类型，并且

    • boundary_check 和 padding_option 必须为空。

  2. 如果 pointer 是 N 维指针张量，则加载 N 维张量。在这种情况下：

    • mask 和 other 会被隐式广播（broadcast）为 pointer.shape 的形状。

    • other 会被隐式转换为 pointer.dtype.element_ty 类型，并且

    • boundary_check 和 padding_option 必须为空。

  3. 如果 pointer 是由 make_block_ptr 定义的块指针（block pointer），则加载一个张量。在这种情况下：

    • mask 和 other 必须为 None，并且

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

参数:

  • pointer (triton.PointerType，或 dtype=triton.PointerType 的块) – 指向要加载数据的指针

  • mask (triton.int1 的块，可选) – 如果 mask[idx] 为 false，则不加载地址 pointer[idx] 处的数据（使用块指针时必须为 None）

  • other (Block, 可选) – 如果 mask[idx] 为 false，则返回 other[idx]。如果 other 为 None，则掩码覆盖（masked-out）的值未定义。

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

  • padding_option – 应为 {“”, “zero”, “nan”} 之一，即越界时使用的填充值。"" 表示未定义值。

  • cache_modifier (str, 可选，应为 {“”, “.ca”, “.cg”, “.cv”} 之一，其中 “.ca” 表示在所有层级缓存，“.cg” 表示在全局层级缓存（在 L2 及以下缓存，不在 L1），“.cv” 表示不缓存并重新获取。详见 缓存操作符) – 更改 NVIDIA PTX 中的缓存选项

  • eviction_policy (str, 可选) – 更改 NVIDIA PTX 中的驱逐策略

  • volatile (bool, 可选) – 更改 NVIDIA PTX 中的 volatile（易失性）选项