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 会被隐式广播到 pointer.shape，

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

    • boundary_check 和 padding_option 必须为空。

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

    • mask 和 other 必须是 None，并且

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

参数:

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

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

  • other (Block, 可选) – 如果 mask[idx] 为 false，则返回 other[idx]

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

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

  • cache_modifier (字符串，可选，应为 {“”, “.ca”, “.cg”, “.cv”} 之一，其中“.ca”表示在所有级别缓存，“.cg”表示在全局级别缓存（在 L2 及以下缓存，而不是 L1），“.cv”表示不缓存并重新获取。有关更多详细信息，请参阅 缓存操作符。) – 更改 NVIDIA PTX 中的缓存选项

  • eviction_policy (字符串, 可选) – 更改 NVIDIA PTX 中的逐出策略

  • volatile (布尔值, 可选) – 更改 NVIDIA PTX 中的 volatile 选项