有關 Triton's 語言功能、編譯器設計理念、一般安裝與除錯提示，請參閱 上游 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 程式設計介面撰寫 kernel；PPU 後端會透明處理針對 PPU 硬體的編譯與執行。

2. 核心功能與最佳化

PPU 後端遵循 Triton 的標準多階段編譯流程：ttir → ttgir → llir → hgbin，並在各階段插入 PPU 專屬的 pass。核心實作位於 third_party/ppu/backend/compiler.py。

PPU 上的主要最佳化：

  • AIU 非同步資料移動：從全域記憶體載入到共享記憶體 (TSM) 的資料載入會自動轉換為由 AIU (compute Unit 中的 AI 加速器) 執行的非同步複製。結合軟體流水線，可在迴圈內實現多重緩衝，將 tile 預取與 MMA 運算重疊，以隱藏全域記憶體存取延遲並提升運算效率。

  • 交錯共享記憶體佈局：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 為 Triton 語言新增 aiu_load API，可透過 AIU (compute Unit 中的 AI 加速器) 將資料從全域記憶體移至共享記憶體以加速資料傳輸。目前提供 tl.aiu_load、make_block_ptr 與 make_tensor_descriptor 介面，使用者可輕鬆整合至現有 Triton kernel。

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-byte 對齊。
  • BLOCK_M 必須是 16 的倍數。
  • BLOCK_K 維度的資料必須是 32 bytes 的倍數。
  • 區塊張量在 M 與 K 維度上必須是連續的。
  • 支援的資料型別：
    • PPU0010: b16
    • PPU0015: b16, b8

以下是使用 tl.aiu_load 的矩陣乘法 kernel 範例：

@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 kernel 中使用 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

或使用 virtualenv：

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 中的「Building with a custom LLVM」與「Tips for building」。

4.3 驗證安裝

從原始碼建置完成後，請使用 python -c "import triton; print(triton.__version__)" 檢查 import 與版本。若版本號正確印出，則表示 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：匯出編譯記錄。