/note/tech

I Made Zig Compute 33 Million Satellite Positions in 3 Seconds. No GPU Required.

要約:

■ 1. プロジェクト概要

  • astroz は最速の汎用SGP4実装でネイティブZigで1秒あたり1100万〜1300万回の伝播計算を実現
  • Python経由では約700万回毎秒でpip install astrozのみで利用可能
  • GPUを使用せずCPUのSIMD命令のみで高速化を達成
  • 汎用実装の定義:
    • heyoka.pyは複数衛星の同時バッチ処理で高速だが一般的なODE積分器でLLVMのJITコンパイルとC++依存関係が必要
    • astrozは時間バッチ伝播で2倍高速で単一衛星の複数時刻計算に特化
    • GPU加速SGP4実装は大規模バッチワークロードで高速だがCUDA/OpenCLセットアップが必要で汎用的ではない

■ 2. SGP4最適化の動機

  • SGP4は1980年代からTLEデータから衛星位置を予測する標準アルゴリズム
  • 既存実装は元のリファレンスコードの移植版で機能するが高密度時間解像度では遅い
  • 実用的な課題:
    • 1ヶ月のエフェメリスデータを1秒間隔で生成すると衛星1個あたり260万回の伝播計算が必要
    • 地上局ネットワークでの通過予測は数週間にわたりサブ秒精度が必要
    • 接近スクリーニングの軌道解析は近接アプローチを捕捉するため細かい時間ステップが必要
  • 典型的な良い実装で200万〜300万回毎秒だと衛星1個あたり数秒かかり反復解析やインタラクティブツール構築では遅い

■ 3. 初期最適化とスカラー実装

  • SIMD導入前のスカラー実装で既にRust sgp4クレート(最速のオープンソース実装)と同等以上の性能
  • 主要な設計選択:
    • 分岐のないホットパス:
      • SGP4アルゴリズムには多数の条件分岐が存在
      • 深宇宙対近地球や異なる摂動モデルやケプラーソルバーの収束チェックなど
      • 可能な限り分岐なし表現として記述
      • 当初はパフォーマンス目的ではなくコードの理解しやすさのため
      • 偶然にも現代CPUが予測可能な命令ストリームを好むため高速化
    • コンパイル時事前計算:
      • Zigのcomptimeで任意のコードをコンパイル時に実行可能
      • SGP4のセットアップ作業の多く(重力定数や多項式係数や派生パラメータ)を一度計算してバイナリに焼き込む
      • ランタイム初期化や繰り返し計算が不要
      • constでマークした変数は自動的にcomptimeとして扱われる
  • スカラー実装で520万回毎秒を達成しRustの500万回毎秒を若干上回る

■ 4. SIMD実装の発見と基礎

  • SGP4は状態に依存せず各衛星と各時点が独立しているためSIMDに適している
  • ZigがSIMDを第一級市民として扱う設計:
    • シンプルな型宣言のみで開始可能
    • const Vec4 = @Vector(4, f64)で4つの64ビット浮動小数点数のベクトルを定義
    • イントリンシクスやプラットフォーム検出や条件付きコンパイルが不要
    • LLVMバックエンドがコード実行環境に適した命令セットを自動的にターゲット
  • 組み込み演算:
    • スカラーを全レーンにブロードキャストする@splat
    • LLVMイントリンシクスを通じた自動ベクトル化された超越関数
    • @sinと@cosはLinux x86_64でのlibmvecなどプラットフォーム最適実装を使用
  • レーン単位の思考への転換:
    • スカラーコードではif文で一方のパスを実行
    • SIMDでは全4レーンが同時に実行されるため両方の結果を計算してレーン毎に選択
    • 両方のパスを計算することは無駄に見えるが分岐予測ミスより高速な場合が多い
    • SGP4では大半の衛星が同じパスを取るため無駄な計算は少ない
  • 収束ループの処理:
    • SGP4のケプラーソルバーは各結果が収束するまで反復
    • スカラーでは許容誤差以下になるまでループ
    • SIMDでは異なるレーンが異なる速度で収束
    • 解決策はレーン毎に収束をマスクで追跡し@reduceで全員完了を確認
  • パターン理解後はスカラー実装を一行ずつ変換し元のコードをテストスイートで比較できるよう保持

■ 5. 3つの伝播モード

  • 時間バッチ propagateV4:
    • 最も一般的なワークロードに対応
    • 単一衛星の複数時点を同時に4時点処理
    • エフェメリスデータ生成や通過予測や軌道解析や接近スクリーニングに使用
    • 衛星の軌道要素がレジスタに留まり4つの出力を計算するため最もキャッシュフレンドリー
    • 最大の初期高速化を実現
  • 衛星バッチ propagateSatellitesV4:
    • 多数の衛星を1つの特定時点で処理
    • 衝突スクリーニングスナップショットやカタログ全体の可視性チェックに使用
    • 4つの異なる衛星を同一時点で同時処理
    • ElementsV4という異なるデータレイアウトが必要
    • 各フィールドがVec4で4つの異なる衛星の値を保持
  • コンステレーションモード propagateConstellationV4:
    • 多数の衛星を多数の時点で伝播
    • ライブデモでは13000個の衛星を1440時点で処理
    • キャッシュを意識したタイリング戦略:
      • 各衛星の全時点を計算する素朴なアプローチはキャッシュをスラッシング
      • 64時点を全衛星で処理してから次の64時点を処理
      • 時間値がL1キャッシュにホットな状態で衛星バッチをスイープ
      • タイルサイズ64はL1サイズを作業データサイズで割りSIMDフレンドリーな数に丸めたもの
    • 大規模カタログでは素朴なループより15〜20%高速

■ 6. atan2問題と解決策

  • SGP4のケプラーソルバーはatan2を必要とするがLLVMはベクトル化されたビルトインを提供していない
  • スカラー関数を呼び出すとSIMD実装が壊れる
  • 多項式近似による解決:
    • SGP4の精度要件(モデルに固有の制限がある)では完全な精度は不要
    • 引数を[0,1]の範囲に保ち多項式精度を向上
    • ホーナー法で計算
    • @selectを使用した分岐なしの象限補正
  • 多項式近似は約1e-7ラジアンの精度でLEO距離で約10mmの位置誤差に相当
  • これはSGP4の固有精度限界内で数日間の伝播では数キロメートルの不確実性がある
  • この数学は複雑でAIの支援を受けて実装

■ 7. Struct of Arraysデータ構造

  • 複数衛星処理のためStruct of Arraysレイアウトを使用
  • ElementsV4構造体:
    • 各フィールドがVec4で4つの衛星の値を保持
    • 重力モデルや離心率や傾斜角や昇交点赤経などの軌道要素
    • 事前スプラット済み定数をinit時に一度計算
  • 事前スプラット最適化:
    • ホットパスでの繰り返し@splat呼び出しを排除
    • 毎秒数百万回実行されるコードでは小さな最適化も重要

■ 8. ベンチマーク結果

  • ネイティブ実装比較(Zig vs Rust):
    • 1日(分間隔)でastroz 0.27ms対Rust 0.31msで1.16倍高速
    • 1週間(分間隔)でastroz 1.99ms対Rust 2.04msで1.03倍高速
    • 2週間(分間隔)でastroz 3.87ms対Rust 4.03msで1.04倍高速
    • 2週間(秒間隔)でastroz 222ms対Rust 234msで1.05倍高速
    • 1ヶ月(分間隔)でastroz 8.37ms対Rust 8.94msで1.07倍高速
    • 両実装ともスカラー処理で約500万回毎秒を達成
    • Zig実装がRustを若干上回るのはホットパス最適化とcomptimeの積極的な事前計算による
  • ネイティブSIMDスループット:
    • astroz Zig SIMD実装で1200万回毎秒
    • astroz Zigスカラー実装で520万回毎秒
    • Rust sgp4で510万回毎秒
    • @Vector(4, f64)を使用した複数衛星または時点の並列処理でスループットがスカラー実装の2倍以上に向上
  • Pythonバインディング性能:
    • 2週間(秒間隔)でastroz 160ms対python-sgp4 464msで2.9倍高速
    • 1ヶ月(分間隔)でastroz 5.9ms対python-sgp4 16.1msで2.7倍高速
  • Pythonバインディングスループット:
    • astroz Python 700万回毎秒
    • satkit(Rust)370万回毎秒
    • python-sgp4 220万回毎秒
  • heyoka.pyとの比較:
    • 8衛星×1440時点でheyoka.py 1620万回毎秒対astroz 750万回毎秒
    • 1衛星×1440時点でheyoka.py 380万回毎秒対astroz 850万回毎秒
    • 100衛星×100時点でheyoka.py 1550万回毎秒対astroz 840万回毎秒
    • heyoka.pyは複数衛星バッチで勝利しastrozは時間バッチワークロードで勝利
    • ユースケースによる選択:
      • 数百の衛星を同時に同一時点で伝播する場合はheyoka.pyが有利
      • 個別衛星のエフェメリス生成や通過予測や軌道解析の場合はastrozが2倍高速
      • 簡単なデプロイメントが必要な場合はastrozがpip installのみでNumPyのみ依存しheyoka.pyはLLVMとC++スタックが必要

■ 9. 実用的なデモ

  • インタラクティブCesiumビジュアライゼーション:
    • CelesTrakの全アクティブカタログ約13000個の衛星を処理
    • 1440時点(1日を分解像度)で約1900万回の伝播計算
    • Pythonバインディング経由で約2.7秒で完了(約700万回毎秒)
    • TEME→ECEF座標変換に約0.6秒追加で合計約3.3秒
    • ネイティブZigでは1100万〜1300万回毎秒に到達
  • 高速化により計算のバッチ処理やスケジューリングを考える必要がなくなり必要な時に必要な計算を実行可能

■ 10. 今後の展望と利用方法

  • 今後の開発:
    • SDP4実装で深宇宙天体に対応(現在は軌道周期225分未満の近地球衛星のみ)
    • マルチスレッド化でコア全体にスケール
    • SIMD作業は単一スレッドで実行されており別の倍率が待機中
  • 利用方法:
    • PyPIで利用可能でpip install astroz
    • Zigプロジェクトへの追加はzig fetch --save
    • GitHubでオープンソースとして公開
    • examplesを参照して統合可能
    • ライブデモで動作確認可能