タイル表示レイアウト

図 1

図 1 は、配列 F32[3,5] が 2x2 タイリングでメモリ内でどのように配置されるかを示しています。このレイアウトのシェイプは F32[3,5]{1,0:T(2,2)} として記述されます。1,0 は寸法の物理的な順序（レイアウトの minor_to_major フィールド）に対応し、コロンの後の (2,2) は物理寸法の 2×2 タイルによるタイル分を示します。

直感的に、上記の例のように、タイルがシェイプを覆うようにレイアウトされ、各タイル内で要素がタイルなしでレイアウトされます。この例の右側は、メモリ内のレイアウトを示しています。元の配列境界が等しくない場合であっても、完全な 2x2 タイルを表示するために追加される白いパディング要素が含まれています。

パディングの追加要素に特定の値を含める必要はありません。

形状とタイルを指定してタイリングするための線形インデックスの式

タイリングなしの場合、配列境界 d=(d_n, d_n-1, ... , d₁)（d1 は最小の次元）を持つ配列内の要素 e=(e_n, e_n-1, ... , e₁) は、位置にメジャーからマイナー 順で配置されます。

linear_index(e, d)
= linear_index((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d₁ + e_n-1 ...

このドキュメントでは、表記をわかりやすくするために、タイルの次元数が配列と同じであると仮定しています。XLA によるタイリングの実装では、これはサイズの小さいタイルに一般化されています。そのため、最初の最大サイズを変更せずに最下位のディメンションのみにタイルを適用し、指定されたタイルに、タイル状の物理サイズの接尾辞を指定します。

サイズ（t_n、t_n-1、...、t₁）のタイルを使用する場合、インデックス（e_n、e_n-1、...、e₁）を持つ配列内の要素は、最終的なレイアウトの次の位置にマッピングされます。

_n

レイアウトは、(⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋) の 2 つの部分を持つと考えることができます。これは、サイズの配列内のタイル インデックスに対応します（Эd_n/t_n ₡, ... , 💯?d₁ / t_e に mod1, t mod 1/t₁ に、_{ceil 関数は、大きい配列の境界にタイルが重なると、図 1 のようにパディングが挿入されるため、Ъdi/ti の形で現れます。タイルとタイル内の要素はどちらもタイリングなしで再帰的に配置されます。}

たとえば、図 1 の例では、要素 (2,3) には、結合された座標ベクトル (1,1,0,1) のタイル インデックス (1,1) とタイル内インデックス (0,1) があります。タイル インデックスには境界（2,3）があり、タイル自体は（2,3,2,2）の結合ベクトルが（2,2）です。論理シェイプでインデックス（2,3）を持つ要素のタイルを使用した線形インデックスは次のようになります。

linear_index_with_tile((2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= linear_index((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index((0,1))
.

パッド - 形 - 転置としてのタイリング

タイルベースのレイアウトは次のように動作します。
寸法（d_n、d_n-1、...、d1）（d1 が最もマイナーな次元）の配列を考えてみましょう。サイズ（t_n、t_n-1、...、t₁）（t₁ が最も短い寸法）のタイルでレイアウトされた場合、そのタイルはパッドの形状の変更と転置によって次のように記述できます。

1. 配列のパディングが次のとおりに行われるようになりました。(⬇d_n/t_nↁ ∙ t_n, ... , Эd₁/t₁Ъ∙t₁)。
2. 各次元 i は (Ъd_i/ti!, t_i) に分割されます。つまり、配列の形状が
   に再形成されます (↘d_n/t_nЪ, t_n, ... , Эd₁/t₁↘, t₁)。
   この形状変更だけでは、物理的なレイアウトの変更はないため、ビットキャストです。タイリングについて明示的に考えていない場合、この形状変更により、パディングされたシェイプと同じ数の要素を持つ任意のシェイプを表現できます。この例は、この方法でタイルを表現する方法です。
3. 転置は、相対的な順序を保ちながら t_n、...、t₁ を最小のディメンションに移動することで行われます。つまり、最も大から大へのディメンションの順序は
   となります（Ъd_n/t_n↘, ... , Ъd₁/t₁、t_n、...、t₁）。

最後のシェイプには接頭辞
（Эd_n/t_n↘, ... , Ъd₁/t₁💯?）という接頭辞があり、これは各次元のタイルの数を表します。配列内の要素（e_n、...、e₁）は、最終的な形状
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... mod t₁, ... mod t₁) でこの要素にマッピングされます。要素の線形インデックスが上記の式に従うことは簡単です。

繰り返しタイル張り

XLA のタイリングは、繰り返し適用することでさらに柔軟なものになります。

図 2

図 2 は、サイズ 4x8 の配列が 2 つのレベルのタイル（最初の 2x4 の次に 2x1 のタイル）でタイルされる方法を示しています。この繰り返しのタイリングを (2,4)(2,1) と表現します。それぞれの色は 2x4 のタイルを示し、赤の枠線ボックスはそれぞれ 2x1 のタイルを示します。番号は、その要素のメモリ内の線形インデックスをタイル形式で示します。この形式は、TPU の BF16 に使用される形式と一致しますが、初期タイルの方が大きく、つまり（8,128）（2,1）タイリングになります。2x1 でタイリングする目的は、TPU のアーキテクチャに合わせて 2 つの 16 ビット値を一緒に収集して 1 つの 32 ビット値を形成することです。

2 番目以降のタイルは、この例の (8,128)(2,1) のようにタイル内のデータを再配置するだけのマイナーなタイル内ディメンションの両方を参照できますが、前のタイルの主要なクロスタイル ディメンションを参照することもできます。

タイルを使用してディメンションを組み合わせる

XLA のタイリングでは、ディメンションを組み合わせることもできます。たとえば、F32[2,7,8,11,10]{4,3,2,1,0} のディメンションを F32[112,110]{1,0} に結合してから、(2,3) と並べることができます。使用されるタイルは (*,∗,2,∗,3) です。ここでは、タイルのアスタリスクは、そのディメンションを取得し、次のマイナーなディメンションと組み合わせることを意味します。隣接する複数のディメンションを 1 つのディメンションにまとめることができます。包括的なディメンションは、タイルのそのディメンションのタイル値 -1 として表されます。それ以外の場合は、ディメンション サイズとしてタイルでは無効です。

より正確に言うと、シェイプのディメンション i がタイルのアスタリスクによって削除されると、以前のタイリングの定義が適用される前に、そのディメンションはタイルされるシェイプとタイルベクトルの両方から削除され、シェイプのディメンション i-1 の配列境界は d_i-1 から d_id_i-1 に増加します。このステップが、タイルベクトルのアスタリスクごとに繰り返されます。