このページは Cloud Translation API によって翻訳されました。
Switch to English

タイルレイアウト

図1

図1は、配列F32 [3,5]がメモリに2x2タイリングでどのように配置されるかを示しています。このレイアウトの形状はF32 [3,5] {1,0:(2,2)}として記述されます。ここで、1,0は次元の物理的な順序(レイアウトのminor_to_majorフィールド)に関連し、(2,2)はコロンは、2x2タイルによる物理寸法のタイリングを示します。

直感的に、タイルは形状をカバーするようにレイアウトされ、次に各タイル内に要素がタイリングなしでレイアウトされます。上記の例では、例の右側に追加された白いパディング要素を含むメモリ内のレイアウトが示されています元の配列の境界が均等でない場合でも、完全な2x2タイルを作成するため。

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

与えられた形状とタイルのタイリングのための線形インデックス式

タイリングせずに、要素e =(E N、E N-1、...、E 1)アレイの配列の境界は、=(D、N、D N-1、...、D 1)(d1があるdの最もマイナーな次元)は、次の位置にメジャーからマイナーの順序で配置されます。

linear_index(e、d)
= linear_index((e n 、e n-1 、...、e 1 )、(d n 、d n-1 、...、d 1 ))
= e n d n-1 ... d 1 + e n-1 d n-2 ... d 1 + ... + e 1

このドキュメントでは表記を簡単にするために、タイルは配列と同じ次元数を持っていると想定しています。 XLAのタイリングの実装では、これは、初期の最長次元を変更せずに最少の次元にのみタイルを適用することで、より少ない次元のタイリングに一般化されます。そのため、指定されたタイリングは、物理的な次元のサフィックスに言及します。タイル張りされている形状。

サイズ(t n 、t n-1 、...、t 1 )のタイリングを使用すると、インデックス(e n 、e n-1 、...、e 1 )を持つ配列の要素がこれにマッピングされます最終レイアウトでの位置:

linear_index_with_tile(e、d、t)
= linear_index((⌊e/t⌋、e mod t)、(⌈d/t⌉、t))(算術は要素ごと、(a、b)は連結)
= linear_index((⌊eN / T N⌋、...、⌊e1 / T 1⌋、E N MOD T N、...、E 1 MOD T 1)、(⌈dN / T N⌉、 ...、⌈d1 / T 1⌉、T nは 、T N-1、···、T 1))
= linear_index((⌊eN / T N⌋、...、⌊e1 / T 1⌋)、(⌈dN / T N⌉、...、⌈d1 / T 1⌉))∙T N t n-1 ... t 1 + linear_index((e n mod t n 、...、e 1 mod t 1 )、(t n 、t n-1 、...、t 1 ))

サイズのタイルのアレイにタイルインデックスに対応する(⌊eN / T N⌋、...、⌊e1 / T 1⌋)、(⌈dN:レイアウトは二つの部分を有すると考えることができます/ T N⌉、...、⌈d1 / T 1⌉)、および(E N MOD T N、...、E 1 MOD T 1)、その内のタイルインデックスに対応します。 CEIL関数は / T⌈dに表示⌉タイルがより大きな配列の境界をオーバーラン場合ので、図1のタイル内のタイルとの両方の要素をタイリングせずに再帰的にレイアウトされているように、パディングが挿入されます。

図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)、(2,2))
=(1∙3 + 1)∙2∙2 +(0∙2 + 1)
= 17。

pad-reshape-transposeとしてタイリング

タイリングベースのレイアウトは次のように動作します。
次元の配列(d n 、d n-1 、...、d1)を考えます(d1が最も小さい次元です)。サイズのタイリング(t n 、t n-1 、...、t 1 )(t 1が最も小さい寸法)でレイアウトされている場合、そのタイリングは次のようにpad-reshape-transposeの観点から説明できます。仕方。

  1. アレイは、(⌈dN / T N⌉∙T N、...、⌈d1 / T 1⌉∙T 1)に埋め込まれています。
  2. 各次元iは(⌈dI / T iは⌉、Iを T)、すなわちアレイをに再形成さに分割され
    (⌈dN / T N⌉、T N、...、⌈d1 / T 1⌉、T 1)。
    このリシェイプ自体には物理的なレイアウトの変更がないため、このリシェイプはビットキャストです。タイリングを明確に考えていない場合、この形状変更は、パディングされた形状と同じ数の要素を持つ任意の形状を表現できます。ここでの例は、この方法でタイルを表現する方法です。
  3. 転置は、相対的な順序を維持しながらt n 、...、t 1を最短の次元に移動することで発生します。これにより、最長の次元から最短の次元への次元の順序が次のようになります。
    (⌈dN / T N⌉、...、⌈d1 / T 1⌉、T N、···、T 1)。

最終形状には接頭辞が付きます
各次元におけるタイルの数を記載している(⌈dN / T N⌉、...、⌈d1 / T 1⌉)。配列の要素(e n 、...、e 1 )は、最終的な形状でこの要素にマップされます。
(⌊eN / T N⌋、...、⌊e0 / T 0⌋、E N MOD T N、...、E 1 MOD T 1)。要素の線形インデックスが予想どおり上記の式に従うことが簡単にわかります。

繰り返しタイリング

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

図2

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

2番目以降のタイルは、(8,128)(2,1)を使用したこの例のように、タイル内のデータを再配置するタイル内の両方の小さな次元を参照できますが、主要なクロスタイルも参照できることに注意してください前のタイリングからの寸法。

タイルを使用して寸法を組み合わせる

XLAのタイリングは、ディメンションの組み合わせもサポートしています。たとえば、F2 [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 i d i-1に増加します。この手順は、タイルベクトルの各アスタリスクに対して繰り返されます。