Tritonの言語機能、コンパイラ設計思想、一般的なインストール方法、およびデバッグのヒントについては、upstream Triton READMEおよび公式Tritonドキュメントを参照してください。本ドキュメントではPPUバックエンドの拡張機能とupstream Tritonに対する最適化に焦点を当てています。

1. 概要

本リポジトリはTritonのPPUフォークであり、T-Head Semiconductorが独自開発したPPUをサポートします。

  • 対象ハードウェア: T-Head PPU。2世代のTensor Core — PPU0010 (MMAv1) および PPU0015 (MMAv2) をカバーします。
  • ランタイム: T-Head SAIL SDK (PPU SDK) 環境が必要です。
  • プログラミングインターフェース: upstream Tritonと全く同じPythonプログラミングインターフェースでカーネルを作成できます。PPUバックエンドがPPUハードウェアを対象としたコンパイルと実行を透過的に処理します。

2. 主要機能と最適化

PPUバックエンドはTritonの標準的な複数ステージのコンパイルフロー (ttir → ttgir → llir → hgbin) に従い、各ステージでPPU固有のパスを挿入します。コア実装はthird_party/ppu/backend/compiler.pyにあります。

PPUにおける主な最適化:

  • AIU非同期データ移動: グローバルメモリから共有メモリ (TSM) へのデータロードは、Compute Unit内のAIアクセラレータであるAIUによる非同期コピーに自動的に変換されます。ソフトウェアパイプラインと組み合わせることで、ループ内のマルチバッファリングを実現し、タイルのプリフェッチとMMA演算をオーバーラップさせてグローバルメモリアクセスレイテンシを隠蔽し、演算効率を向上させます。

  • スウィズルされた共有メモリレイアウト: Tritonはワープ数、タイル形状、および要素のビット幅からタイリング方式とスウィズルエンコーディングを自動的に導出します。これにより、複数のワープが同時にアクセスする場合の共有メモリバンク競合を排除し、効果的なバンド幅を保証するとともに、TSM内のデータレイアウトをMMAオペランドレイアウトに合わせることで、余分なレイアウト変換オーバーヘッドを削減します。

  • Tensor Coreアクセラレーションと低精度サポート: Tritonはtl.dotをPPUのMMAハードウェア命令にコンパイルし、タイル分割の粒度を自動的に導出してTensor Coreアクセラレーションを最大限に活用します。FP8 (E5M2 / E4M3)、FP16、BF16を含む低精度／混合精度サポートを提供します。

3. AIU利用ガイド

PPU TritonはTriton言語に新しいaiu_load APIを追加し、Compute Unit内のAIアクセラレータであるAIUを介してグローバルメモリから共有メモリへデータを移動し、データ転送を高速化します。現在、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.
    """

BLOCK_M x BLOCK_KテンソルをM x Kグローバルデータテンソルから共有メモリへロードする必要がある場合はtl.aiu_loadを使用してください。

aiu_loadの制約:

  • テンソルは32バイト境界にアラインメントされている必要があります。
  • BLOCK_Mは16の倍数である必要があります。
  • BLOCK_K次元のデータは32バイトの倍数である必要があります。
  • ブロックテンソルはMおよびK次元の両方で連続している必要があります。
  • サポートされるデータ型:
    • PPU0010: b16
    • PPU0015: b16, b8

以下はtl.aiu_loadを使用したmatmulカーネルの例です:

@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など (詳細はupstream Triton READMEを参照)。
export PPU_SDK=/usr/local/PPU_SDK           # 実際のインストールパスに合わせて調整してください
export LD_LIBRARY_PATH=$PPU_SDK/lib:$LD_LIBRARY_PATH

4.2 ソースからのインストール

# リポジトリルートで
pip install -r python/requirements.txt      # ビルド時依存関係
pip install -e .                            # 編集可能モードでコンパイルおよびインストール

またはvirtualenvを使用する場合:

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

ビルドオプション (カスタムLLVM、ccache、MAX_JOBSによるメモリ使用量制限など) はupstream Tritonと同一です。「Building with a custom LLVM」および「Tips for building」をupstream Triton READMEで参照してください。

4.3 インストールの検証

ソースからビルドした後、python -c "import triton; print(triton.__version__)"でインポートとバージョンを確認してください。バージョン番号が正しく表示されれば、Tritonが正常にインストールされ、C++/MLIR拡張が正しくロードされています。

その後、AIUセクションの例を実行してエンドツーエンドの検証を行うことができます。

4.4 環境変数

一般的な環境変数についてはupstream Triton READMEを参照してください。以下はPPU固有の追加環境変数です:

  • PPU_LLC_OPTIONS: ppu-llcに渡す追加オプション。
  • DISABLE_PPU_LLC_OPT: ppu-llc最適化を無効化。
  • TRITON_LIBDEVICE_PATH: PPU libdeviceパスを指定。
  • TRITON_DUMP_COMPILE_LOG: コンパイルログをエクスポート。