关于 Triton 的语言特性、编译器设计理念、常规安装和调试技巧，请参阅 上游 Triton README 和 官方 Triton 文档。本文档重点介绍 PPU 后端相对于上游 Triton 的扩展和优化。

1. 概述

本仓库是 Triton 的PPU 分支，为 T-Head Semiconductor 独立开发的 PPU 提供支持。

  • 目标硬件：T-Head PPU，涵盖两代 Tensor Core — PPU0010 (MMAv1) 和 PPU0015 (MMAv2)。
  • 运行时：需要 T-Head SAIL SDK（PPU SDK）环境。
  • 编程接口：使用与上游 Triton 完全相同的 Python 编程接口编写内核；PPU 后端会透明地处理针对 PPU 硬件的编译和执行。

2. 核心特性与优化

PPU 后端遵循 Triton 的标准多阶段编译流程：ttir → ttgir → llir → hgbin，在每个阶段插入 PPU 特定 pass。核心实现位于 third_party/ppu/backend/compiler.py。

PPU 上的关键优化：

  • AIU 异步数据移动：从全局内存到共享内存（TSM）的数据加载会被自动转换为由 AIU（计算单元中的 AI 加速器）执行的异步复制。结合软件流水线，这在循环内实现了多缓冲，将 tile 预取与 MMA 计算重叠，以隐藏全局内存访问延迟并提高计算效率。

  • Swizzled 共享内存布局：Triton 会根据 warp 数量、tile 形状和元素位宽自动推导 tiling 方案和 swizzle 编码。一方面，这消除了共享内存 bank 冲突，保证多 warp 并发访问下的有效带宽；另一方面，它使 TSM 中的数据布局与 MMA 操作数布局对齐，减少额外的布局转换开销。

  • Tensor Core 加速与低精度支持：Triton 将 tl.dot 编译为 PPU 的 MMA 硬件指令，并自动推导 tile 划分粒度以充分利用 Tensor Core 加速。它提供低精度/混合精度支持，包括 FP8 (E5M2 / E4M3)、FP16 和 BF16。

3. AIU 使用指南

PPU Triton 通过新的 aiu_load API 扩展了 Triton 语言，该 API 通过 AIU（计算单元中的 AI 加速器）将数据从全局内存移动到共享内存以加速数据传输。它目前提供 tl.aiu_load、make_block_ptr 和 make_tensor_descriptor 接口，用户可以轻松将其集成到现有的 Triton 内核中。

3.1 tl.aiu_load

def aiu_load(pointer, offsets=(0, 0), block_shape=(0, 0), shape=(0, 0), dtype=void, order=(1, 0), _semantic=None):
    """
    Return a tensor of data whose values are loaded from memory at location defined by `pointer`:

        (1) If `pointer` is a single element pointer, offsets&block_shape&shape should be provided.  In
            this case:
            - `pointer` is the start address of a memory tensor
            - `offsets` is a 2D value indicate the offset of the start pointer of the loading tile to the memory tensor
            - `block_shape` is a 2D value indicate the tile size loaded by aiu_load.
            - `shape` is a 2D value indicate the memory tensor size.
            - `order` is the order of the original data format

        (2) If `pointer` is a block pointer defined by `make_block_ptr`, a tensor is loaded.
    """

当需要从 M x K 全局数据张量加载 BLOCK_M x BLOCK_K 张量到共享内存时，使用 tl.aiu_load。

aiu_load 的约束：

  • 张量必须 32 字节对齐。
  • BLOCK_M 必须是 16 的倍数。
  • BLOCK_K 维度上的数据必须是 32 字节的倍数。
  • 块张量在 M 和 K 维度上必须是连续的。
  • 支持的数据类型：
    • PPU0010: b16
    • PPU0015: b16, b8

以下是使用 tl.aiu_load 的矩阵乘法内核示例：

@triton.jit
def matmul_kernel_aiu(a_ptr, b_ptr, c_ptr,
                      stride_am, stride_ak,
                      stride_bk, stride_bn,
                      stride_cm, stride_cn,
                      M, N, K,
                      BLOCK_SIZE_M: tl.constexpr,
                      BLOCK_SIZE_N: tl.constexpr,
                      BLOCK_SIZE_K: tl.constexpr):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m
    offs_am = pid_m * BLOCK_SIZE_M
    offs_bn = pid_n * BLOCK_SIZE_N
    offs_k = 0
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.aiu_load(a_ptr, [offs_am, offs_k], [BLOCK_SIZE_M, BLOCK_SIZE_K], [M, K], tl.float16)
        b = tl.aiu_load(b_ptr, [offs_k, offs_bn], [BLOCK_SIZE_K, BLOCK_SIZE_N], [K, N], tl.float16)
        accumulator = tl.dot(a, b, acc=accumulator)
        offs_k += BLOCK_SIZE_K
    c = accumulator.to(tl.float16)

    # -----------------------------------------------------------
    # Write back the block of the output matrix C with masks.
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

3.2 make_block_ptr

使用 make_block_ptr 可以简化 aiu_load；它需要将 make_block_ptr / aiu_load / advance 一起使用。

def make_block_ptr(base: tensor, shape, strides, offsets, block_shape, order, _semantic=None):
    """
    Returns a pointer to a block in a parent tensor

    :param base: The base pointer to the parent tensor
    :param shape: The shape of the parent tensor
    :param strides: The strides of the parent tensor
    :param offsets: The offsets to the block
    :param block_shape: The shape of the block
    :param order: The order of the original data format
    """
  • tl.make_block_ptr：构造指向父张量中某个块的指针。
  • tl.aiu_load：接受 make_block_ptr 的结果作为输入。
  • tl.advance：推进块指针。
def advance(base, offsets, _semantic=None):
    """
    Advance a block pointer

    :param base: the block pointer to advance
    :param offsets: the offsets to advance, a tuple by dimension
    """

示例：

@triton.jit
def matmul_kernel_aiu(a_ptr, b_ptr, c_ptr,
                      stride_am, stride_ak,
                      stride_bk, stride_bn,
                      stride_cm, stride_cn,
                      M, N, K,
                      BLOCK_SIZE_M: tl.constexpr,
                      BLOCK_SIZE_N: tl.constexpr,
                      BLOCK_SIZE_K: tl.constexpr):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_am = pid_m * BLOCK_SIZE_M
    offs_bn = pid_n * BLOCK_SIZE_N
    offs_k = 0
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    a_tensor_ptr = tl.make_block_ptr(a_ptr, (M, K), (stride_am, stride_ak), (offs_am, offs_k), (BLOCK_SIZE_M, BLOCK_SIZE_K), (1, 0))
    b_tensor_ptr = tl.make_block_ptr(b_ptr, (K, N), (stride_bk, stride_bn), (offs_k, offs_bn), (BLOCK_SIZE_K, BLOCK_SIZE_N), (1, 0))

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.aiu_load(a_tensor_ptr)
        b = tl.aiu_load(b_tensor_ptr)
        accumulator = tl.dot(a, b, acc=accumulator)
        a_tensor_ptr = tl.advance(a_tensor_ptr, (0, BLOCK_SIZE_K))
        b_tensor_ptr = tl.advance(b_tensor_ptr, (BLOCK_SIZE_K, 0))

    c = accumulator.to(tl.float16)
    # -----------------------------------------------------------
    # Write back the block of the output matrix C with masks.
    offs_cm = pid_m * BLOCK_SIZE_M
    offs_cn = pid_n * BLOCK_SIZE_N
    c_tensor_ptr = tl.make_block_ptr(c_ptr, (M, N), (stride_cm, stride_cn), (offs_cm, offs_cn), (BLOCK_SIZE_M, BLOCK_SIZE_N), (1, 0))
    tl.store(c_tensor_ptr, c)

3.3 make_tensor_descriptor

PPU Triton 为 tl.make_tensor_descriptor 添加了 AIU 支持；当满足条件时，加载操作可以自动提升为 AIU 加载。

def make_tensor_descriptor(
    base: tensor,
    shape: List[tensor],
    strides: List[tensor],
    block_shape: List[constexpr],
    padding_option="zero",
    _semantic=None,
) -> tensor_descriptor:
    """Make a tensor descriptor object

    :param base: the base pointer of the tensor, must be 16-byte aligned
    :param shape: A list of non-negative integers representing the tensor shape
    :param strides: A list of tensor strides. Leading dimensions must be multiples
        of 16-byte strides and the last dimension must be contiguous.
    :param block_shape: The shape of block to be loaded/stored from global memory
    """

在 Triton 内核中使用 tl.make_tensor_descriptor 时，当满足条件时，加载操作可以自动提升为 AIU 加载。

3.4 更多示例

有关更多 AIU 用例，请参阅仓库中的测试和性能示例：

4. 构建与安装

4.1 前置条件

  • PPU SDK：安装 PPU SDK，并通过 PPU_SDK 环境变量指定其路径（默认为 /usr/local/PPU_SDK）。SDK 提供头文件、库以及 ppu-llc 和 llvm-irformatter 等工具链可执行文件。
  • 运行时驱动：安装 PPU 运行时驱动 libhggc.so，并确保可以通过 ldconfig 或 LD_LIBRARY_PATH 找到它。
  • 构建依赖：Python 3.10+、C++ 编译器、CMake、Ninja 等（请参阅 上游 Triton README）。
export PPU_SDK=/usr/local/PPU_SDK           # adjust to your actual install path
export LD_LIBRARY_PATH=$PPU_SDK/lib:$LD_LIBRARY_PATH

4.2 从源码安装

# In the repository root
pip install -r python/requirements.txt      # build-time dependencies
pip install -e .                            # compile and install in editable mode

或使用虚拟环境：

python -m venv .venv --prompt triton
source .venv/bin/activate
pip install -r python/requirements.txt
pip install -e .

构建选项（自定义 LLVM、ccache、使用 MAX_JOBS 限制内存使用等）与上游 Triton 相同；请参阅 上游 Triton README 中的“使用自定义 LLVM 构建”和“构建技巧”。

4.3 验证安装

从源码构建后，使用 python -c "import triton; print(triton.__version__)" 检查导入和版本。如果版本号正确打印，则表示 Triton 已成功安装且 C++/MLIR 扩展已正确加载。

随后可以运行 AIU 部分的示例进行端到端验证。

4.4 环境变量

有关通用环境变量，请参阅 上游 Triton README。以下是额外的 PPU 特定环境变量：

  • PPU_LLC_OPTIONS：传递给 ppu-llc 的额外选项。
  • DISABLE_PPU_LLC_OPT：禁用 ppu-llc 优化。
  • TRITON_LIBDEVICE_PATH：指定 PPU libdevice 路径。
  • TRITON_DUMP_COMPILE_LOG：导出编译日志。