RedFuser — AIアクセラレータ向け自動演算子フュージョン

AI/ML
RedFuser — AIアクセラレータ向け自動演算子フュージョン

RedFuser — AIアクセラレータ向け自動演算子フュージョン

RedFuserは、深層学習コンパイラ上で「連鎖的なリダクション（cascaded reductions）」を自動的に検出・フュージョンし、AIアクセラレータ向けに高効率なカーネルを生成するフレームワークです。Apache TVM上に構築され、演算子の結合性・スケジューリング・メモリ最適化を考慮した最適化パスと検索モデルを備え、ハードウェア固有の命令やメモリ階層を活かしたコード生成を目指します（ASPLOS 2026の成果物）。

https://github.com/alibaba/redfuser

概要

RedFuserは、深層学習における連続したリダクション（例：一連の和/平均などの集約操作）が生成する低効率なメモリ・計算パターンをターゲットにした自動演算子フュージョンのフレームワークです。Apache TVM上に実装され、演算子グラフから「フュージョン可能なリダクション連鎖」を抽出し、合法性チェック・コストモデル・スケジューリング探索を通じて、単一または少数の高性能カーネルへ変換します。これによりメモリアクセス削減やデータ再利用向上、AIアクセラレータ特有の算術ユニットやベクトル命令の活用を図り、推論/学習ワークロードの性能を改善します。

GitHub

リポジトリの統計情報

  • スター数: 10
  • フォーク数: 0
  • ウォッチャー数: 10
  • コミット数: 30
  • ファイル数: 33
  • メインの言語: Python

主な特徴

  • Cascaded reductions（連鎖リダクション）の自動検出と安全なフュージョンを実現
  • Apache TVMと連携したスケジューラ／変換パスの実装
  • ハードウェア指向のコストモデルと探索でアクセラレータ特性を活かす最適化
  • フュージョン後のコード生成でメモリ帯域を削減しデータ再利用を向上

技術的なポイント

RedFuserの技術核は「連鎖的リダクションに対する安全かつ効果的なフュージョン戦略」です。多段リダクションは中間テンソルを介した大量のメモリ読書きを生むため、単純な演算子単位の最適化では限界がある点に着目しています。実装はTVM上で、演算子グラフ解析フェーズにより候補となるリダクション系列を抽出し、データ依存と数学的性質（結合性・可換性など）を利用してフュージョンの合法性を判定します。

次に、ハードウェア特性を組み込んだコストモデルにより、フュージョンが有利かを定量評価します。コストモデルはメモリ転送量、算術密度、アクセラレータの特有命令（ベクトル/マトリクス積）やスレッド並列性を考慮します。最適化されたスケジューリング探索では、タイルサイズ、ループ順序、並列化戦略、レジスタ/共有メモリの割当てなどを探索し、候補ごとに生成されるカーネルの推定性能を比較します。探索はヒューリスティックとコスト推定の組み合わせで実行され、過度な探索膨張を抑えつつ実用的な解を得る設計となっています。

生成フェーズでは、フュージョン後の計算を単一カーネルや少数の効率的カーネルへ展開するためのテンプレート化されたコード生成パスを持ちます。ここで、ループ分割やアンローリング、SIMD化、メモリバンクの競合回避、ならびに複数リダクションを同時に処理するための累積ローカル変数（例：レジスタに保持する部分和）を利用し、メモリアクセスを極力削減します。さらに、アクセラレータのスレッドモデルやメモリ階層に応じたマッピングを行い、現代的なAI専用ハードでの性能を狙います。

論文成果（ASPLOS 2026）に紐づくアーティファクトとして、ベンチマーク・実験スクリプト・検証用のテストケースが含まれており、研究結果の再現性に配慮した構成になっています。

プロジェクトの構成

主要なファイルとディレクトリ：

  • .asf.yaml: file
  • .clang-format: file
  • .gitattributes: file
  • .github: dir
  • .gitignore: file
  • LICENSE: file（Apache 2.0）
  • README.md: file（プロジェクトの概要と使用方法）
  • artifact: dir（論文付属の実験スクリプトやベンチマーク）
  • redfuser: dir（フレームワーク本体、変換パス・スケジューラ実装）
    • fusion.py: リダクション検出とフュージョン判定のロジック
    • schedule_search.py: スケジュール探索とコスト評価
    • codegen/: カーネルテンプレートとコード生成モジュール
  • tests: dir（単体テストと再現実験ケース）
  • examples: dir（簡単な使用例とデモ）
  • setup.py / pyproject.toml: パッケージングと依存定義
  • docs/: ドキュメント（使用方法・論文へのリンク）

リポジトリはTVM依存の箇所があるため、実行には対応するTVMバージョンと、場合によっては特定のバックエンド（XPU/カスタムアクセラレータ）に関するセットアップが必要です。artifactディレクトリには実験再現に便利なDockerfileやスクリプトが含まれる場合が多く、論文の評価環境を再現する手助けになります。

…他 28 ファイル

まとめ

連鎖リダクションに特化した実用的なフュージョン基盤で、TVM上でのハードウェア寄せ最適化に貢献する良質なアーティファクト。

リポジトリ情報：

READMEの抜粋：

RedFuser

License ASPLOS 2026

This repository contains the artifact for our ASPLOS 2026 paper: “RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators”.

Overview

RedFuser is a novel framework for optimizing cascaded reductions in deep learning compilers. Built on top of [Apache TVM](htt…