【Triton 教程】triton_language.load

Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境，以高效编写自定义 DNN 计算内核，并能够在现代 GPU 硬件上以最大吞吐量运行。

更多 Triton 中文文档可访问 →triton.hyper.ai/

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

返回 1 个数据张量，其值从由指针所定义的内存位置处加载：

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

• mask 和 other 必须也是标量，
• other 会隐式地转换为 pointer.dtype.element_ty 类型，
• boundary_check 和 padding_option 必须为空。

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

• mask 和 other 会被隐式地广播到 pointer.shape，
• other 会隐式地转换为 pointer.dtype.element_ty 类型，
• boundary_check 和 padding_option 必须为空。

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

• mask 和 other 必须为 None，
• 可以指定 boundary_check 和 padding_option 来控制超出越界访问的行为。

参数**：**

• pointer（triton.PointerType*，*或 dtype=triton.PointerType 的块）- 指向要加载的数据的指针。
• mask（triton.int1 的块*，*可选）- 如果 mask[idx] 为 false，则不加载 pointer[idx] 处的数据（对于块指针必须为 None）。
• other (块, 可选) - 如果 mask[idx] 为 false，则返回 other[idx]。
• boundary_check（整数元组*，*可选）- 表示应进行边界检查维度的元组。
• padding_option - 应为 {“”, “zero”, “nan”} 中的一个，越界时进行填充。
• cache_modifier**（*str，可选，*应为 {“”, “ca”, “cg”} 中的一个）- 其中「ca」表示在所有层级进行缓存，「cg」表示在全局层级缓存（在 L2 及以下缓存，不是 L1），详细信息请参见缓存操作符。）在 NVIDIA PTX 中更改缓存选项。
• eviction_policy（str, 可选） - 更改 NVIDIA PTX 中的驱逐策略。
• volatile（bool, 可选） - 更改 NVIDIA PTX 中的易失性选项。