InfiniTensor/InfiniTensor

InfiniTensor is a high-performance inference engine tailored for GPUs and AI accelerators. Its design focuses on effective deployment and swift academic validation.

InfiniTensor – 軽量で拡張可能なAI計算グラフ用ランタイム

概要

  • InfiniTensorは、計算グラフ(例:ONNXモデル)を読み込み、変換し、実行するオープンソースのC++ライブラリです。ネイティブのCPUバックエンドが付属し、InfiniOps や InfiniRT と呼ばれるカスタムアクセラレータバックエンドに処理を委任できます。
  • コアライブラリはハードウェア非依存で、実際のデバイス固有のカーネルは、InfiniTensorと一緒にインストールする別のSDKに含まれています。

重要性

  • 現代のAIワークロードは、ONNXのようなグラフベースの中間表現(IR)で表現されることがよくあります。そのIRをインポートし、グラフレベルの最適化を適用し、CPUまたは専用アクセラレータで実行できるランタイムは、推論スタックの重要な構成要素です。
  • グラフエンジン(InfiniTensor)とデバイスカーネル(InfiniOps/InfiniRT)を分離することで、ハードウェアベンダーはコアランタイムを変更することなく独自の実装を組み込むことができます。

主な機能(READMEに記載されている内容)

  • ONNXインポーター – ONNXファイルを内部グラフ表現に解析します。
  • グラフ変換レイヤー – 汎用の書き換えをサポートしています。将来の作業では、より高度な導出ベースの最適化のために RefactorGraph と EinNet に言及しています。
  • CPUランタイム – 開発とフォールバック用の組み込みの純C++実行エンジン(pyinfinitensor Pythonバインディング)。
  • アクセラレータ委任 – InfiniOps/InfiniRTビルドが存在する場合、演算子はアクセラレータにディスパッチされます。サポートされていない演算子は、静かにフォールバックするのではなく、明示的なエラーを発生させます。
  • Python API – make install-python でインストールでき、ランタイムを pyinfinitensor パッケージとして公開します。
  • テストスイート – C++単体テスト、ONNXフロントエンドテスト、Python APIテストが提供されています。

はじめ方

  1. リポジトリをクローンし、サブモジュールを初期化します(git submodule update --init --recursive)。
  2. CPUのみのPythonパッケージをビルドしてインストールします。
    make install-python INFINI=OFF PYTHON=$(command -v python3)
    
  3. アクセラレータサポートの場合は、まず対応する InfiniOps と InfiniRT SDKをインストールし、INFINI=ON で再ビルドしてSDKプレフィックスを指定します。
  4. 提供されている make ターゲット(build、test-cpp、test-onnx、test-api)を使用して、インストールをコンパイルおよび検証します。

ロードマップのハイライト

  • RefactorGraph – 次世代グラフフレームワーク(別リポジトリとしてリンク)。
  • EinNet – 導出ベースのテンソルプログラム最適化(実験ブランチ NNET_e2e)。
  • PET – テンソルプログラムの自動修正のための部分等価変換。
  • InfiniOps/InfiniRTエコシステムによるアクセラレータサポートの継続的な拡張。

対象ユーザー

  • カスタムAI演算子を構築したり、新しいグラフ最適化パスを探求する研究者。
  • ONNXモデルに自社のカーネルを公開するための薄く再利用可能なフロントエンドを望むハードウェアベンダー。
  • CPUでの推論に軽量でPythonフレンドリーなランタイムを必要とし、独自アクセラレータをプラグインするオプションが必要な開発者。

詳細情報

  • インストールガイド、ユーザーガイド、分散例は /docs および /examples ディレクトリにあります。
  • 付随する研究論文 EinNet(OSDI 2023)および PET(OSDI 2021)は、プロジェクトが基盤とする最適化技術についてのより深い文脈を提供します。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト