カテゴリ: コンパイル時間: Mosaic の未検証のメモリアクセス アライメント
このエラーは、コンパイラがメモリ アクセス オペレーション（vector.load、vector.store、tpu.load、tpu.store など）を分析し、特定のディメンションに使用される動的インデックスが、必要なタイリング サイズの倍数であることを静的に証明できない場合に発生します。
エラー メッセージの例:
INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128
at location: ...
The MLIR operation involved:
%14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space<vmem>>, index, index) -> vector<1x32xf32>
XLA バックエンド: TPU
概要
カーネルがベクトルを読み込みまたは保存するとき、メモリ アドレス（ベースポインタと動的インデックスから計算）は、ハードウェア上のベクトルのタイリング サイズと一致する必要があります。たとえば、ディメンションが 128 個の要素でタイル化されている場合、それにアクセスするために使用される動的インデックスは 0、128、256 などである必要があります。多くのオペレーション（ベクトル ロードやストアなど）では、静的インデックスにこのような要件はありません。
コンパイラは、静的解析を使用してこの要件を適用します。インデックス変数の履歴を、その変数を生成した算術演算（乗算、加算など）まで遡ってトレースします。コンパイラが、結果の値が常にタイリング サイズで割り切れることを（コンパイル時に）保証できない場合、このエラーが発生します。
コンパイラは、「証明済みの不整合」と「不明なアライメント」を同じように扱います。したがって、数学的にアライメントが保証されていないインデックス（i * 128 + 32）の場合、コンパイラは同じエラーを発生させます。
そのため、このエラーは次の場合に発生する可能性があります。
- メモリにアクセスするには、ランタイム変数（動的インデックス）を使用します。
- インデックス計算ロジックが複雑すぎて、コンパイラで分析できない。
- インデックスは数学的に有効ですが、コードに明示的な証明がありません。
- 静的分析により「証明済みの不整合」が特定されます。
デバッグ
このエラーを解決するには、次のいずれかの方法をお試しください。
1. アライメントを明示的にアサートする
インデックスが有効であることがわかっているが、コンパイラがそれを証明できない場合は、tpu.assume_multiple オペレーションを使用します。これは、値が特定の係数で割り切れることをコンパイラに約束するものです。
2. Aligned Loads と Rotate を使用する
位置合わせのずれが意図的なシナリオでは、位置合わせされていない小さなベクトル セグメントを読み込む代わりに、
- 完全に配置された大きなタイルを読み込み、動的な量だけ値を回転させて、目的のデータを所定の位置に移動します（動的な開始インデックスを持つベクトル スライスはサポートされていないため）。または
- データがインデックス 0 から始まり、アクセス間のストライドがハードウェア アライメントと一致するように、テンソルを再形成またはパディングします。
- 例: オフセット 1 から始まるサイズ 32 のチャンクを反復処理する場合、オフセットは 1、33、65... になります。（未調整）。
- 修正: 最初のチャンクが 0 で、ディメンションが 128 にパディングされた新しいテンソルにデータを再パックします。オフセットは 0、128、256... となり、アライメント要件を満たします。
これらのメソッドはメモリを多く消費しますが、カーネル ロジックを簡素化し、手動によるアライメント アサーションの必要性をなくすことができます。