HLO オペレーション統計ツール

HLO Op Stats を使用すると、プロファイリング セッション中に実行された High Level Optimizer（HLO）オペレーションのパフォーマンス統計情報を確認できます。このツールは、プロファイリング期間中の HLO グラフ内で最も時間のかかるオペレーションをハイライト表示することで、パフォーマンスの最適化の対象となる可能性のあるターゲットを特定するのに役立ちます。

サポートされているプラットフォーム

TPU: サポート対象

GPU: サポート対象

GPU の場合、HLO オペレーションは、実際に実行されるカーネルと N:M の関係にあります。カーネル レベルの統計情報については、GPU カーネル統計ツールをご覧ください。

HLO オペレーション統計ツールのコンポーネント

HLO Op Stats ツールには、次の主要コンポーネントがあります。

• HLO オペレーション統計テーブル: これは主なコンポーネントで、プロファイリング セッション中に実行されたすべての HLO オペレーションの詳細な内訳を表形式で示します。個別の HLO オペレーションごとに 1 行があり、そのオペレーションに関するさまざまな詳細をキャプチャする列があります。
  • プルダウン セレクタを使用すると、各 HLO オペレーションで可視化する列を決定できます。
  • 2 つ目のプルダウン セレクタを使用して、HLO Op カテゴリで行をフィルタすることもできます。
  • 検索ボックスを使用すると、特定のプログラム ID、HLO オペレーション、HLO オペレーション テキスト、フレームワーク オペレーション名でフィルタできます。フィルタは、指定された文字列が各列内の任意の場所に表示されるように選択します。
• 概要グラフ セクション: このセクションには、次の方法でオペレーションごとの詳細な表を要約したグラフが表示されます。
  • さまざまなカテゴリの HLO オペレーションで消費された時間の割合を示す円グラフ。
  • さまざまな個別の HLO オペレーションによって消費された時間の割合を示す円グラフ（上位 N 個のオペレーションに切り捨てられ、読みやすくするために残りは「その他」に分類されます）。
  • 再マテリアライゼーションの一部であるオペレーションに費やされた合計時間の割合を示す円グラフ。XProf は、プロファイルに関連付けられたコンパイラ メタデータからこの情報を取得します。
  • 再マテリアル化時間（存在する場合）が費やされる HLO オペレーション カテゴリを示す円グラフ。
  • 外部コンパイルに費やす時間: 外部コンパイルは TensorFlow の機能で、XLA 計算内の特定のオペレーションをアクセラレータ デバイスではなくホスト CPU で透過的に実行できるようにします（デバイスにない I/O アクセスを必要とする tf.summary や tf.print など）。
  • すべての HLO オペレーションの GFLOPS/s をプロットしたグラフ（合計自己時間で並べ替え）。
  • TPU のみ、レプリカ グループごとの HLO あたりの時間: プルダウンから、プロファイリング セッション中に実行されたさまざまな集約オペレーションを選択できます。その集約オペレーションの異なるインスタンスが、異なるレプリカ グループ間で実行されている可能性があります（AllGather）が表示されます。円グラフには、これらのインスタンス間の時間の分布が表示されます。

HLO オペレーションの統計情報テーブルの詳細

列の見出しをクリックすると、HLO オペレーションの統計情報の表を並べ替えることができます。デフォルトの順序は、オペレーションの合計セルフタイム（表の「ランク」）に基づいています。

この表には、各 HLO オペレーションに関する次の情報が含まれています。

• オペレーション名: XLA コンパイラによって HLO オペレーションに割り当てられた一意の名前。
• プログラム ID: このオペレーションが関連付けられている HLO モジュールの ID
• HLO Op カテゴリ: これらは主に XLA コンパイラによって定義されます。XProf は、特定のオペレーション（畳み込みの融合など）を特定して分類するためにヒューリスティクスも使用します。
• HLO Op テキスト: XLA コンパイラによって提供され、入力/パラメータの型や形状などの詳細が含まれます。
• フレームワーク オペレーション名: フレームワーク レベルのオペレーション（例: JAX）が、この HLO オペレーションの生成につながった。
• 発生回数: プロファイリング期間中にこの特定の HLO オペレーションが実行された回数の合計。
• 合計時間（μs）: このオペレーションのすべての発生で費やされた合計時間。このオペレーションに子オペレーションがある場合（融合内など）、この時間には子オペレーションで費やされた時間も含まれます。
• 平均時間（μs）: この HLO オペレーションの実行あたりの平均時間（子オペレーションで費やされた時間を含む）。
• 合計自己時間（μs）: この HLO オペレーションの本体内でのみ費やされた時間の累積値（子オペレーションで費やされた時間は除く）。
• 平均自己時間（μs）: この HLO オペレーションの実行あたりの平均時間（子オペレーションで費やされた時間は除く）。
• 合計自己時間（%）: すべてのオペレーションでデバイスの合計時間に対するオペレーションの自己時間の割合。
• 累積合計自己時間（%）: 「ランク」順で上位に表示されるすべてのオペレーションの合計自己時間（%）の累積値。
• DMA の停止（%）: ダイレクト メモリアクセス（DMA）オペレーションが原因でオペレーションが停止した合計時間の割合。
• FLOP、HBM、内部 TPU メモリの帯域幅使用量（使用量/秒）（利用可能な場合）（例: CMEM は TPU v4 にのみ存在します）。これらはすべて、コンパイラの静的費用分析（分子）とプロファイリングされた実行時間（分母）を使用して計算されます。
• モデルの GFLOPS/秒: GFLOP の費用は XLA コンパイラによって計算され、時間はプロファイルによって測定されます。
• 正規化された GFLOPS/秒: 数値の精度と、さまざまな精度でデバイスがサポートするピーク FLOPS に基づいて、オペレーションごとにコンパイラによって計算された FLOPS を正規化します。たとえば、TPU v6e は、bf16 のピーク FLOP の 2 倍の int8 をサポートしています。XProf はデフォルトで bf16 に正規化します。
• メモリ帯域幅: 任意のメモリ（オンチップ TPU VMEM や HBM など）からオペレーションで消費される 1 秒あたりのバイト数。XLA コンパイラは「バイト」数を提供し、プロファイルで時間の長さが測定されます。
• HBM BW: HBM から消費されるバイト数/秒。
• 動作の強度: ルーフライン分析で定義されているように、バイトあたりの FLOPS の比率として計算されます。
• コンピューティング/メモリバウンド: Roofline モデルに基づいて、この列は、オペレーションのパフォーマンスが主にコンピューティング能力またはメモリ帯域幅によって制限されているかどうかを示します。
• 再マテリアル化: オペレーションが再マテリアル化の一部であったかどうかを示します。
• 外部コンパイル: Tensorflow の外部コンパイル機能を使用してホスト CPU でオペレーションが実行されたかどうかを示します。
• Autotuned: XLA によってオペレーションが自動チューニングされたかどうかを示します。詳細については、自動チューニングと永続化された自動チューニングをご覧ください。