triton-operator-code-review

Triton 演算子静的コードレビュー（Ascend NPU）

レビュー原則

• Ascend 固有の制約を優先：ハードウェアの違いに焦点を当て、Triton の一般知識は確認しません
• Mask はゼロ容認：Ascend では範囲外アクセスは一切許容されません

深刻度の分類

レベル	意味	典型的な問題
P0	必ずクラッシュするか、結果が誤る	Mask の漏れ、カーネル種別の不一致、Atomic ループのデッドロック
P1	精度／機能に問題が生じる可能性が高い	低精度のままのリダクション、max を減算しない Softmax
P2	パフォーマンス／保守性	冗長なメモリアクセス、BLOCK がアラインされていない

参照リソースの読み込み

フェーズ	読み込むもの	読み込まないもの
Phase 1：Host 側	ascend-triton-api-constraints.md	dtype-matrix
Phase 2：Device 側	ascend-api-dtype-matrix.md	test-patterns
項目ごとの照合	code-review-checklist.md	—
公式実装の参照	ascend-test-patterns.md	—
API シグネチャ／パラメーターの確認	triton-api-reference.md	—

読み込みの条件：tl/libdevice API のシグネチャ、パラメーター、または使用可能なデータ型を確認する必要がある場合は、triton-api-reference.md をすべて読みます。

レビューワークフロー

Phase 1：Host 側

確認項目	識別方法	レベル
コア数のハードコード	grid = (20,) などのリテラル	P0
コア種別の不一致	tl.dot を含む kernel で num_vectorcore を使用	P0
行列乗算が要素ごとの演算に退化	tl.dot がなく、Vector Core による要素ごとの乗算加算で matmul/GEMV を実装	P0
BLOCK_SIZE が tl.constexpr ではない	宣言を確認	P1
行列演算の BLOCK が 16 の倍数ではない	数値を確認	P2

コア種別：tl.dot を含む場合 → AI Core（num_aicore）；要素ごとの演算／リダクション／活性化の場合 → Vector Core（num_vectorcore）。注意：行列乗算系の演算子（GEMV を含む）は、必ず tl.dot + AI Core を使用します。Cube Core のスループットは Vector Core を大幅に上回ります。

Phase 2：Device 側

Mask の完全性（P0）：すべての tl.load／tl.store に mask= があるか、make_block_ptr を使用している必要があります。

データ型（P0-P1）：tl.dot の入力でサポートされるのは int8/fp16/fp32/bf16 のみです。dot_scaled は条件付きでサポートされます（bf16/fp16 の lhs/rhs、int8 の scale、fp32 の out）。permute／trans は int64 をサポートしません。

精度処理（P1）：FP16/BF16 のリダクション前に .to(tl.float32) を適用します。Softmax では必ず max を減算します。

リダクション系演算子の単一パス確認（P1）：GroupNorm/LayerNorm/RMSNorm などの演算子では、同じデータを複数回 load していないか確認します（先に統計量を計算し、その後に正規化のために再度 load するパターン）。D ≤ UB の場合は単一パスにする必要があります。1 回の load の後、tl.sum(x, 1) で UB 内ですべての計算を完了します。2 パスでは 5 倍遅くなることがあります。

制御フロー（P0）：Triton IR は MLIR の構造化制御フロー（scf.for／scf.while）を使用しており、early exit をサポートしません。

• ❌ for/while ループ内の return — コンパイルエラー：「Cannot have return statements inside while or for」（サブ関数内の return も含む）
• ❌ for ループ内の break — コンパイルエラー：「unsupported AST node type: Break」

Tensor インデックス（P0）：Triton tensor は Python 形式の [] 添字操作をサポートしません。

• ❌ tensor[i] = val — AssertionError
• ❌ val = tensor[i] — AssertionError
• ❌ tensor[i:j] スライス — コンパイルエラー

コードパターン：

• ❌ for ... : tl.atomic_cas/or/xor/and/xchg(...) — デッドロックする可能性があります（P0）
• ❌ 複数コアの kernel で tl.atomic_add の戻り値を使用（P0）
• ❌ kernel 内で import numpy（P0）
• ⚠️ Host のホットパスで tensor[i].item() を使用 — CPU-NPU 同期が発生します（P2）
• ⚠️ リダクション系演算子に複数の for ... : tl.load(x_ptr + ...) ループがある — 2 パスのアンチパターンです（P1）
• ⚠️ Post-dot 操作で tl.parallel(bind_sub_block=True) を使用していない — tl.dot 後に要素ごとの操作があるものの、並列分割がありません（P2）
• ⚠️ 整数型変換時のオーバーフロー — tl.cast で int8/int16 に変換する際、overflow_mode が指定されていません（P1）

Phase 3：パフォーマンス上の懸念（P2）

• 同じ ptr に対して tl.load を複数回実行 → GM アクセスが冗長です
• tl.arange(0, N) * stride（stride > 1）→ 非連続なメモリアクセスです
• pid をループなしで block に直接マッピング → 負荷が不均衡になります
• kernel 内で補助 tensor（cos/sin など）に対して broadcast stride でオフセットを計算（b * stride0 + n * stride1 + ...）→ Host 側で expand+contiguous に変更し、row * D + col のオフセットに統一することを推奨します
• リダクション系演算子で weight/bias を事前に展開せず、kernel 内で weight_ptr + ch_idx による gather を使用 → Host 側で連続レイアウトに展開することを推奨します

アンチパターン一覧（NEVER）

Host

• ❌ コア数のハードコード／行列乗算で num_vectorcore を使用／BLOCK_SIZE が tl.constexpr ではない
• ❌ 行列乗算（GEMV を含む）を Vector Core による要素ごとの乗算加算で実装 — 必ず tl.dot + AI Core を使用します

Device

• ❌ tl.load／tl.store に mask がない／tl.dot の入力が int32/int16/int64／dot_scaled
• ❌ atomic_or/xor/and/xchg/cas を for ループ内で使用／kernel 内でサードパーティーライブラリを使用
• ❌ FP16/BF16 のリダクションで FP32 に昇格しない／Softmax で max を減算しない
• ❌ 整数の切り詰め変換で overflow_mode を指定しない
• ❌ for/while ループ内の return／break — Triton の構造化制御フローでは early exit をサポートしません（サブ関数内の return も含む）
• ❌ tensor[i] インデックス操作（読み取り／代入／スライス）— 代わりに tl.where／tl.gather／tl.extract_slice を使用します

出力

code-review-report-template.md の形式に従ってレポートを出力します。