AIシステムにおける算子のハンドオプティマイゼーション

前回の記事では、算子の計算とスケジューリングの概念を探り、効率的なスケジューリング戦略がハードウェア性能の解放と遅延の低減において重要であることを強調しました。本稿では、手書き算子のスケジューリング時に考慮すべき重要な要素を深く掘り下げ、いくつかの著名な高性能算子ライブラリを紹介します。

計算分析

算子を最適化する前に、現在のプログラムのボトルネックがどこにあるかを把握する必要があります。計算ボトルネックかメモリアクセスボトルネックかです。これらの分析には通常、RoofLineモデルが使用されます。

計算分析指標

まずいくつかの重要な指標を定義します：

  • 計算量：プログラムが1回の完全な計算を実行する際に発生する浮動小数点演算の総数、つまり時間計算量を指し、単位はFlopsです。例えば、畳み込み層の計算量は$M^{2*K^{2}*C_{in}*C_{out}}$となります。ここでMは入力特徴マップのサイズ、Kは畳み込みカーネルのサイズ、Cはチャネル数です。
  • メモリアクセス量：プログラムが1回の完全な計算を実行する際に発生するメモリ交換の総量、つまり空間計算量です。理想的な場合、プログラムのメモリアクセス量はモデルのパラメータと出力特徴マップのメモリ使用量の合計です（入力のメモリ使用量は前の算子の出力使用量として計算されます）。単位はByteで、float32型の場合は4を乗算する必要があります。畳み込み層の場合、メモリ使用量は$K^{2}*C_{in}*C_{out}+M^{2}*C_{out}$となります。
  • モデルの計算強度：計算量をメモリアクセス量で割った値が算子の計算強度であり、計算過程において、1バイトのメモリ交換あたりに何回の浮動小数点計算が行われるかを示します。計算強度が大きいほど、メモリ使用効率が高いと言えます。
  • モデルの理論性能：モデルが計算プラットフォームで達成できる毎秒浮動小数点演算回数の上限であり、RoofLineモデルがこの指標を計算する方法を提供しています。

RoofLineモデル

RoofLineモデルは、高性能計算プラットフォームの性能を評価・分析するための有効なツールです。計算量とメモリアクセス量を分析することで、特定の計算能力と帯域幅の条件下で計算タスクが達成できる理論的な性能上限を決定します。RoofLineモデルの核心は、ハードウェアリソースの制約を明らかにし、開発者や研究者が現行の計算プラットフォームの制約下でアプリケーションが実現できる理論性能上限を理解するのを助けることです。

  • 計算能力が「屋根」の高さを決定します（緑色の線分）

  • 帯域幅が「屋根の傾斜」を決定します（赤色の線分）

  • Compute-Bound：算子の計算強度が計算プラットフォームの計算強度上限を超える場合、算子は現在の計算プラットフォームで計算ボトルネックに陥ります。計算プラットフォームの計算能力を最大限に活用する観点から見ると、この時点で算子は計算プラットフォームの全計算能力を利用しています。

  • Memory-Bound：算子の計算強度が計算プラットフォームの強度上限より小さい場合、算子の性能は完全に計算プラットフォームの帯域幅上限とモデル自身の計算強度によって決定されるため、この状態をMemory-Boundと呼びます。明らかに、モデルが帯域幅ボトルネックの領域にある前提では、計算プラットフォームの帯域幅が大きいほど（屋根の傾斜が急になるほど）、またはモデルの計算強度が大きいほど、モデルの理論性能は線形的に増加します。

(3,224,224)の入力条件下で、VGG16モデルの前向き伝播計算量は15GFLOPs、メモリアクセス量は約600MBであり、計算強度は25FLOPs/Byteとなります。MobileNetの計算量は0.5GFLOPs、メモリアクセス量は74MBで、その計算強度はわずか7FLOPs/Byteです。1080Ti GPU上では、その計算能力は11.3TFLOP/s、帯域幅は484GB/sであり、したがってこのプラットフォームの最大計算強度は約24となります。

上図からわかるように、MobileNetはMemory-Bound領域にあり、1080Ti上の理論性能は3.3TFLOPsに過ぎません。VGGはCompute-Bound領域にあり、1080Tiの全計算能力を完全に利用しています。RoofLineモデルを通じて、計算量とメモリアクセス量が増加すると、性能向上はハードウェアの計算能力と帯域幅の制約を受けることが明確にわかります。この分析は計算集約型およびメモリ帯域幅集約型アプリケーションの最適化にとって極めて重要です。なぜなら、開発者が性能ボトルネックを特定し、対応する最適化戦略を立てるのを助けるからです。

さらに、RoofLineモデルはハードウェア設計とソフトウェアアルゴリズムの選択を指導するのにも使用できます。例えば、アプリケーションの性能がメモリ帯域幅によって制限されている場合、開発者はより効率的なデータ構造やアルゴリズムを使用してメモリアクセスを削減することを検討するかもしれません。同様に、ハードウェア設計者もRoofLineモデルを使用して、特定のアプリケーション性能に対する異なるハードウェア構成の影響を評価し、より合理的な設計決定を行うことができます。

注意すべきは、RoofLineモデルは理論性能の上限であり、実際の実行では計算能力と帯域幅に加えて、キャッシュサイズと速度など多くの他の要因がプログラムの実行に影響を与えることです。

最適化戦略

具体的な最適化戦略を深く掘り下げる際、主に3つの核心領域に注目します：ループ最適化、命令最適化、およびストレージ最適化です。これらの最適化戦略は、算子の計算特性とハードウェアリソースの特徴に合わせてカスタマイズされることを目的としています。本節ではこれらの最適化技術の概要を簡単に紹介し、後続の章ではより詳細な分析と議論を提供します。

  • ループ最適化

AI算子は一般的に高度に規則化された多層のネストされたループ構造を持っているため、最適化のための豊富な技術手段が提供されています。ReLU、加算（Add）、乗算（Mul）などの要素ごとの操作を例に挙げると、すべてのループ軸で反復を実行して計算を行うことができます。畳み込み（Conv）のようなより複雑な操作でも、7層のネストされたループで実装できます。しかし、これらの直感的な元の計算方法だけを使用すると、しばしば効率が低くなります。

算子のデータレイアウトとメモリアクセス特性を深く分析し、それに応じてループ構造を調整することで、不要なオーバーヘッドを大幅に削減し、ハードウェアリソースをより効率的に利用し、遅延を低減できます。一般的なループ最適化技術には、ループブロッキング（Loop Blocking）、ループアンローリング（Loop Unrolling）、ループ再配置（Loop Reordering）、ループ融合（Loop Fusion）、およびループ分割（Loop Splitting）などが含まれます。これらの最適化技術は、慎重に設計されたループ変換を通じて、計算密度を向上させると同時にデータの局所性を改善し、メモリアクセスパターンを最適化し、最終的に性能の飛躍を実現します。

  • 命令最適化

IntelのAVX-512やARMのSVEのような現代のプロセッサは、強力なベクトル処理能力を提供し、単一の命令で複数のデータポイントを同時に操作できるようにします。この方法により、命令最適化は命令の数と実行サイクルを減らし、遅延を低減し、性能を向上させることができます。さらに、特定のハードウェア向けにカスタマイズされた命令セット、例えばNVIDIAのTensor Coresは、並列処理能力をさらに強化し、ディープラーニングにおけるテンソル演算に特別な最適化を提供します。

命令最適化において、開発者はターゲットハードウェアのアーキテクチャ特性を深く理解し、それらの特性をアルゴリズム実装にどのようにマッピングするかを理解する必要があります。これには、既存のアルゴリズムの再設計が含まれ、それらがハードウェアの並列処理ユニットを最大限に活用できるようにするかもしれません。例えば、SIMD（単一命令複数データ）アーキテクチャに適応するようにデータを再配置するか、特定のハードウェアアクセラレータを利用するためにアルゴリズムを調整します。ハードウェア特性の利用に加えて、命令最適化はコンパイラレベルの最適化、例えば自動ベクトル化、命令スケジューリング、およびレジスタ割り当てなどにも関与します。これらのコンパイラ技術は、最適化を自動的に識別し適用し、ハードウェアのポテンシャルをさらに解放することができます。

  • ストレージ最適化

ストレージ最適化においては、データアクセスパターンとメモリ階層構造を微調整することでシステム全体の性能を向上させようとします。メモリ遅延隠蔽技術は、メモリアクセスが完了するまでの待ち時間を最小化することを目的としています。これは通常、データをキャッシュにプリフェッチし、キャッシュヒット率を向上させるためにデータアクセスパターンを調整するか、または並列に他の計算タスクを実行することによって実現されます。

メモリ遅延隠蔽の目的は、プロセッサがデータ読み込みを待っている間も忙しく保ち、リソース利用率を向上させることです。ダブルバッファリングは2つのバッファを使用してデータフローをスムーズにし、遅延を隠蔽します。一方のバッファのデータが処理されている間に、もう一方のバッファを使用して新しいデータをロードできます。この方法は特にグラフィックスレンダリングやビデオ処理などの分野に適しており、そこでは連続的なデータフローと時間的に敏感な操作がシームレスに接続される必要があります。ストレージ最適化には、メモリリソースの適切な割り当て、メモリ使用量を削減するためのデータ構造の最適化、およびメモリプールを使用してメモリ割り当てと解放のオーバーヘッドを削減することが含まれます。これらの戦略は共同して、アプリケーションのメモリアクセス効率を大幅に向上させ、メモリボトルネックによる性能損失を減少させることができます。

DSLによる算子開発

手書き算子の開発は、開発者が下位レベルのハードウェアの詳細に深く関わることを要求し、データレイアウト、命令選択、インデックス計算など多くの側面の微調整を含むため、カーネルの作成の難易度が大幅に高まります。このプロセスを簡素化するために、TVMやTritonなどがこの分野の傑出した代表として発展してきました。

これらのDSLは、高度な抽象化を提供し、複数の一般的な最適化技術をカプセル化し、コンパイラの最適化段階でこれらの技術を自動的に識別し適用します。開発者はこれらのDSLを使用する際、高レベルの計算ロジックに集中し、DSLが提供するAPIインターフェースを利用することで、高性能な算子を実現できます。伝統的な手書きコードや極限までの最適化と比較して、この方法は極めて高い性能レベルに達することができないかもしれませんが、すでに90%以上の性能レベルを実現でき、開発効率は数十倍向上させることができます。このトレードオフは多くの場合合理的であり、開発者がより高い効率で性能の優れたアプリケーションを開発できるようにします。

TVMによる算子開発

TVMはOctoMLが開発した拡張可能でオープンなエンドツーエンドのAIコンパイラであり、ニューラルネットワークモデルの最適化とデプロイメントに使用され、企業が特定のハードウェア向けに開発し、ディープラーニングソフトウェアをデプロイするために費やすコストと時間を削減することを目的としています。

TVMはHalideの計算とスケジューリング思想を大幅に発展させ、実現可能な最適化をすべてスケジューリングAPIの形式で提示します。Tritonと比較して、TVMはCUDAに限定されず、より多くのバックエンドをサポートし、さらに多くのバックエンドを拡張しやすくなっています。TVM上の典型的なスケジューリングは以下の通りです：

def schedule_dense_packed(cfg, outs):
    """Packed dense schedule."""

    assert len(outs) == 1
    output = outs[0]
    const_ops = []
    ewise_inputs = []
    ewise_ops = []
    dense_res = []
    assert "int" in output.op.input_tensors[0].dtype

    def _traverse(op):
        if topi.tag.is_broadcast(op.tag):
            if not op.same_as(output.op):
                if not op.axis:
                    const_ops.append(op)
                else:
                    ewise_ops.append(op)
            for tensor in op.input_tensors:
                if isinstance(tensor.op, tvm.te.PlaceholderOp):
                    ewise_inputs.append((op, tensor))
                else:
                    _traverse(tensor.op)
        else:
            assert op.tag == "dense_pack"
            dense_res.append(op)

    _traverse(output.op)
    assert len(dense_res) == 1
    dense_stage = dense_res[0].output(0)
    s = te.create_schedule(output.op)

    ##### space definition begin #####
    b, c_o, _, _ = s[dense_stage].op.axis
    c_i, _ = s[dense_stage].op.reduce_axis
    cfg.define_split("tile_b", b, num_outputs=2)
    cfg.define_split("tile_ci", c_i, num_outputs=2)
    cfg.define_split("tile_co", c_o, num_outputs=2)
    cfg.define_knob("oc_nthread", [1, 2])
    ###### space definition end ######

    data, weight = dense_stage.op.input_tensors

    env = get_env()

    cdata = s.cache_read(data, env.inp_scope, [dense_stage])
    cweight = s.cache_read(weight, env.wgt_scope, [dense_stage])
    s[dense_stage].set_scope(env.acc_scope)

    # cache read input
    cache_read_ewise = []
    for consumer, tensor in ewise_inputs:
        cache_read_ewise.append(s.cache_read(tensor, env.acc_scope, [consumer]))

    # set ewise scope
    for op in ewise_ops:
        s[op].set_scope(env.acc_scope)
        s[op].pragma(s[op].op.axis[0], env.alu)

    for op in const_ops:
        s[op].compute_inline()

    # apply tiling for SRAM reuse
    x_b, x_c, _, _ = s[output].op.axis
    x_bo, x_bi = cfg["tile_b"].apply(s, output, x_b)
    x_co, x_ci = cfg["tile_co"].apply(s, output, x_c)
    s[output].reorder(x_bo, x_co, x_bi, x_ci)
    store_pt = x_co

    # set all compute scopes
    s[dense_stage].compute_at(s[output], store_pt)
    for op in ewise_ops:
        s[op].compute_at(s[output], store_pt)

    for tensor in cache_read_ewise:
        s[tensor].compute_at(s[output], store_pt)
        s[tensor].pragma(s[tensor].op.axis[0], env.dma_copy)

    # virtual threading along output channel axes
    if cfg["oc_nthread"].val > 1:
        _, v_t = s[output].split(x_co, factor=cfg["oc_nthread"].val)
        s[output].reorder(v_t, x_bo)
        s[output].bind(v_t, te.thread_axis("cthread"))

    x_bo, x_co, x_bi, _ = s[dense_stage].op.axis
    k_o, _ = s[dense_stage].op.reduce_axis
    s[dense_stage].reorder(x_bo, k_o, x_co)

    k_o, _ = cfg["tile_ci"].apply(s, dense_stage, k_o)
    s[cdata].compute_at(s[dense_stage], k_o)
    s[cweight].compute_at(s[dense_stage], k_o)

    # Use VTA instructions
    s[cdata].pragma(s[cdata].op.axis[0], env.dma_copy)
    s[cweight].pragma(s[cweight].op.axis[0], env.dma_copy)
    s[dense_stage].tensorize(x_bi, env.gemm)
    s[output].pragma(x_ci, env.dma_copy)

    return s

TVMが提供するスケジューリングAPIを使用すると、ループ、命令、ストレージレベルの最適化を簡単に実現でき、これらの最適化をパラメータ化し、自動チューナーを使用してより効率的な実装を検索できます。

TVMは現在、複数のアクセラレータベンダーによって採用されており、独自のハードウェアに適応したカスタム開発が行われています：

  • 希姆計算

    希姆計算はTVMベースのAIコンパイラを使用して希姆一代二代チップをエンドツーエンドでサポートし、カスタム算子ソリューションを実現し、そのモデル性能は手書きの極限に近いです。

  • 華為TBEテンソルアクセラレーションエンジン

    TBE（Tensor Boost Engine）は昇腾AIプロセッサ上のAI Coreで実行される算子を実行する責任を負い、TBEはTVMフレームワークに基づくカスタム算子開発機能を提供し、TBEが提供するAPIを通じて対応するニューラルネットワーク算子の開発を完了できます。

Tritonによる算子開発

TritonはOpenAIが開発した、ディープラーニングと高性能計算タスク向けに設計されたプログラミング言語とコンパイラです。GPU上で実行される複雑な操作の開発を簡素化し、最適化することを目的としています。Tritonの目標は、CUDAよりも高い生産性で高速なコードを記述するためのオープンソース環境を提供することです。

Tritonの核となる考え方は、ブロックベースのプログラミングパラダイムがニューラルネットワークの高性能計算コアの構築を効果的に促進できるということです。CUDAのプログラミングモデルは伝統的なSIMT（Single Instruction Multi Thread）GPU実行モデルであり、スレッドの細粒度でプログラミングしますが、Tritonはブロックの細粒度でプログラミングします。例えば、行列乗算の場合、CUDAとTritonには以下のような違いがあります。

このように、tritonはループ内でブロックごとに計算を行います。この方法の鍵となる利点は、ブロック構造の反復空間を生成し、既存のDSLと比較して、プログラマーがスパース操作を実装する際により多くの柔軟性を提供し、同時にコンパイラがデータの局所性と並列性のために積極的な最適化を行うことを可能にすることです。以下はTritonを使用した行列乗算の例です：

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, # M=N=K=1024
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, #BLOCK_M=BLOCK_N=BLOCK_K=32
):
    offs_m = tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
     c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, accumulator)

TritonのフロントエンドはPythonで実装されており、ユーザーの学習コストを大幅に低減します。そのバックエンドはMLIRに基づいて構築されています。Tritonの最適化思想は2つの部分から成り立っています：

  • レイアウト抽象：レイアウト抽象は計算リソースと入力、出力要素の座標マッピング関係を記述し、主にブロックエンコーディング、共有メモリエンコーディング、スライスエンコーディングなどのいくつかのカテゴリの定義が含まれます。これらのエンコーディング情報は属性としてTensorオブジェクトに付加され、このTensorが入力または出力として必要とするマッピング関係を記述します。Tensorが入力と出力の両方として使用される場合、マッピング関係が互換性がない場合、追加の変換オーバーヘッドを引き起こす可能性がある中間レイアウトを挿入して互換性の適合を完了します。
  • 最適化パス：主にNVIDIA GPU計算カーネル最適化の一般的なテクニックが含まれており、ベクトル化されたメモリアクセスを補助するためのcoalescing、計算メモリアクセスの差を緩和するためのpipeline/prefetch、共有メモリアクセスのbank-conflictを避けるためのswizzlingなどが含まれます。ユーザーはカーネルを開発する際、主にビジネスロジックに注力し、下位ハードウェア最適化の詳細はTritonコンパイラによって実現されます。非常に微細な最適化については、Tritonを使用しても実現できない場合があります。

アプリケーションシナリオでは、Tritonは複数の著名なフレームワークに統合されています：

  • jax-ml/jax-triton：JAXは数値計算を加速するためのPythonライブラリであり、Tritonを使用してJAXプログラムに埋め込むことができるカスタムGPUカーネルを記述できます。JAXではtriton_callを使用してTritonカーネルを簡単に呼び出すことができます。
  • PyTorch/inductor：InductorはTritonの統合においてより包括的かつ実用的です。InductorはTritonを3つの方法で使用しています。計算集約的でない算子に対しては、Inductor IRに基づいて比較的汎用的なCodegenのサポートを実現しています。GEMMに対しては、Jinja2に基づいてパターンマッチングの方法で半カスタマイズされたcodegenを実現しています。Convに対しては、事前に焼き付けられたTritonカーネルを呼び出し、カスタマイズ機能を提供していません。

Tritonの実装原理

Tritonが登場する前、算子エンジニアはDRAM、SRAM、計算ユニットを同時に処理する必要があり、多くの課題に直面していました：

  • メモリ管理：メモリ階層を合理的に利用し、頻繁にアクセスされるデータブロックをより速いストレージ領域にキャッシュし、アライメントとメモリアクセス要求のマージを行い、帯域幅の無駄を避けます。
  • スレッド管理：ハードウェア計算リソースを最大限に活用し、並列スレッドの数とスレッド束のサイズを計画します。
  • 命令使用：CUDAを使用して機能を実現するには複数の命令があり、異なる命令には異なる遅延とスループットがあります。

Tritonは算子開発時の効率を向上させ、開発者がハードウェアの詳細に縛られなくなるようにします。CUDAはThreadに直接向き合うのに対し、TritonはThread Blockに向き合うプログラミングを行い、開発者は1）カーネル起動のパラメータ、2）各データブロックのサイズ、3）データブロック間のインタラクションにのみ注力すればよいです。その下位の詳細はTritonによって実現されます。

TritonはMLIRに基づいて実装されており、そのアーキテクチャは以下の通りです：

Frontendは開発者がPythonで記述したカーネルを対応するTriton IR（Triton Dialect）に変換する役割を担います。@triton.jitを使用してカーネルをマークし、TritonはPython ASTを解析し、ユーザーが定義した計算プロセスをMLIR体系に持ち込み、その後さらに最適化を続けます。

Optimizerの大まかなワークフローは以下の通りです：

主に1）TritonIRの最適化、2）TritonIRからTritonGPU IRへの変換、3）TritonGPU IRの最適化の3つの部分に分かれています。中間を通るデータ構造はTritonGPU IRです。

TritonGPU DialectはTriton Dialectと比較して、主にGPUハードウェア関連のOpとTypeが追加されています。鍵となるOpはデータレイアウトの変換です。現在以下のいくつかのデータレイアウトがあります：

  • Blocked Layout：スレッド間でワークロードを均等に分配されていることを示し、各スレッドがメモリ上の連続したデータブロックを処理します。
  • Shared Layout：データが共有メモリのいくつかの特性を示します。
  • MMA Layout：Tensor Core中のMMA命令の結果のdata layoutを示します。
  • DotOperand Layout：TritonのDotOpの入力のlayoutを示します

いくつかの典型的なdata layoutの変換とその特徴を列挙します：

  • #shared -> #blocked：通常、データがshared memoryからregister fileにロードされることを示し、swizzleを考慮する必要があります
  • #blocked -> #shared：データがregister fileからshared memoryにストアされることを示し、前のステップと同じswizzle方式が必要です
  • #mma -> #blocked：DotOpの出力がより単純なlayoutに変換されることを示し、スレッド間のデータ転送を含むため、通常shared memoryを経由して中継します
  • #blocked -> #dot_operand：DotOpの入力に変換され、このステップもshared memoryを経由する必要がある場合があります

TritonIR上の最適化は主に計算自体に関するものであり、ハードウェアに依存しない最適化で、以下のPassが含まれます：1）Inliner Pass、カーネル呼び出しのサブ関数をインライン展開します；2）Combine Pass、特定のパターンの書き換え；3）Canonicalizer Pass、いくつかの簡略化パターンの書き換え；4）CSE Pass、MLIRのcse Pass、共通部分式を削除するために使用されます；5）LICM Pass、MLIRのLoopInvariantCodeMotion Pass、ループに依存しない変数をforloopの外に移動します。

TritonGPU IRの最適化は計算自体の最適化に加え、GPUハードウェア関連の最適化が追加されており、具体的なPassリストは以下の通りです：1）ConvertTritonToTritonGPU Pass、Triton IRをTritonGPU IRに変換し、主にTritonGPU特有のlayoutを追加します；2）Coalesce Pass、orderを再配置し、最大のcontiguityを持つ次元を最前面に並べます；3）Pipeline Pass、MMA命令に対応するglobal memoryからshared memoryへのN-Buffer最適化；4）Prefetch Pass、MMA命令に対応するshared memoryからregister fileへのN-Buffer最適化

参考文献

[1]: Superjomn's blog | OpenAI/Triton MLIR移行作業の概要

タグ: AIシステム 算子最適化 TVM Triton RoofLineモデル

5月15日 00:03 投稿

ホットタグ

©2026 異端開発室