カテゴリ: コンパイル時間: モザイクの入力/出力のブロックとタイリングの不整合
このエラーは、カーネルの入力または出力のブロック形状が、使用されている特定の TPU ハードウェアのデータ型のデフォルトのタイリングと一致しない場合に発生します。
エラー メッセージの例:
UNIMPLEMENTED: Mosaic failed to compile TPU kernel: Failed to set window params
for input 0: Operand of shape (..., 256, 8192) has tiling (16, 128), but its
block shape (..., 8, 8192) is not divisible by tiling evenly nor matches the
full shape.
XLA バックエンド: TPU
概要
TPU のテンソルコア（TC）には、2 次元ベクトル レジスタがあります。2 つのディメンションは、サブレーンとレーンと呼ばれます。TPU コンピューティング ユニット（MXU）はベクトル レジスタの粒度で動作し、XLA 配列は TPU メモリにタイルでレイアウトされます。
このタイリング（8x128）は、コンピューティング ユニットにデータを供給する際のデータ変換を最小限に抑えます。正確なタイリング ディメンション（サブレーン × レーン）は、ハードウェア世代とデータ型によって異なります。たとえば、ほとんどのタイプで一般的なタイリングは 8×128 です。
コンパイル時に、XLA は各カーネルの入力/出力のマイナー ディメンションと 2 番目のマイナー ディメンションに次の制約を適用します。
- 可分性: ブロック ディメンションは、基盤となるテンソルのタイル ディメンションの倍数である必要があります。
- 完全な形状の例外: ブロック ディメンションが割り切れない場合、基盤となるテンソルのそのディメンションのフルサイズと等しくなければなりません。
このエラーは、ブロックが両方の条件に違反した場合にトリガーされます。たとえば、シェイプ (8, 100) のブロックをシェイプ (8, 1024) の入力からシェイプ (8, 128) のタイリングでハードウェアに読み込むと、100 が 128 と 100 != 1024 で割り切れないため、失敗します。ただし、入力形状が (32, 100) の場合は許可されます。
デバッグ
このエラーを解決するには、カーネルのブロック形状が現在のハードウェア タイリングと一致していることを確認します。必要なタイリングの倍数になるように、ブロックサイズを調整するようカーネルコードを変更します。
- 例: エラーでタイリングが
(16, 128)と示されているが、ブロックの形状が(8, 128)の場合は、形状が(16, 128)と一致するようにブロック仕様を変更します。