【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 中的易失性选项。