Bu belgede, HLO işlemleri için dizine ekleme haritalarını sembolik olarak hesaplamanızı sağlayan HLO dizine ekleme analizi açıklanmaktadır. Dizine ekleme haritası, bir tensörün dizinlerini başka bir tensörün dizinleriyle eşleyen bir işlevdir. Örneğin, bir HLO talimatı çıktısının dizinlerini HLO talimatı girişlerinin dizinleriyle (veya tam tersi) eşler.
Örnek
tensor<20xf32> - tensor<10x20x30xf32> yayını için
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
çıkıştan girişe kadar olan dizin haritası, \((i, j, k) \mapsto (j)\) $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$ içindir.
Motivasyon
XLA GPU; birleştirme, işlenen kullanım ve döşeme şemaları hakkında nedenler belirlemek için özel olarak hazırlanmış çeşitli çözümler kullanır (aşağıda daha ayrıntılı bilgi verilmiştir). Dizine ekleme analizinin amacı, bu tür kullanım alanları için yeniden kullanılabilir bir bileşen sağlamaktır. Dizine ekleme analizi, MLIR'nin Affine Map altyapısı üzerine kuruludur ve HLO semantiği ekler.
Birleşme
Çıktının bir öğesini hesaplamak için giriş öğelerinin hangi öğelerinin/dilimlerinin okunduğunu bildiğimizde, bellek birleştirmeyle ilgili akıl yürütme, sıradan olmayan durumlarda uygun hale gelir.
Yönetilen Kullanım
XLA'da işlem yapılan kullanım, çıktısının tamamen kullanıldığı varsayıldığında her talimat girişinin ne kadar kullanıldığını gösterir. Şu anda kullanım genel bir durum için de hesaplanmamaktadır. Dizine ekleme analizi, kullanımın tam olarak hesaplanmasını sağlar.
Döşeme
Parça/dilim, bir tensörün ofset, boyut ve basamaklara göre parametreleştirilmiş hiper dikdörtgen alt kümesidir. Parça yayılımı, işlemin kendisinin döşeme parametrelerini kullanarak işlemin üreticisinin/tüketicisinin parça parametrelerini hesaplamanın bir yoludur. Softmax ve nokta için bunu yapan bir kitaplık zaten var. Karo yayılımı, dizine ekleme haritaları aracılığıyla ifade edilirse daha genel ve güçlü hale getirilebilir.
İşlev ve Alan Adı
Dizine ekleme haritası, \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\) bir tensörün çoklu dizinini \(\boldsymbol{d}\) tensörün \(A\) öğeleri/aralıkları \(B\)ile eşleyen bir işlevdir. Parametre, \(\boldsymbol{s}\) tensörde bulunan \(B\)ancak tensörde bulunmayan boyutların dizin aralıklarına başvuruda bulunur \(A\).
Örneğin, tensor<2x4x8x16xf32> değerinden tensor<4x8xf32> değerine bir düşüş varsa 2D çıkıştan 4D girişe olan dizine ekleme haritası\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\)olur. Burada \(d_i\) , çıkış tensörünün dizinlerine karşılık gelen boyut parametreleridir. Parametreler \(s_j\)
birden çok değeri kodlar.Yani çıktının bir \((d_0, d_1)\) öğesini hesaplamak için \((s_0, d_0, d_1, s_1)\) giriş öğelerine ihtiyacımız vardır. Bu öğeler \(s_0 \in [0, 2)\) ve\(s_1 \in [0, 16)\)olur.
Bu eşleme, HLO talimatlarının özelliklerinden ya da birleşik talimatların eşlemeleri kullanılarak, bir füzyon için dizine eklemek amacıyla oluşturulabilir. Eşlemede, eşlemenin hangi tensör öğeleri için mevcut olduğunu belirten bir alan adı da bulunur.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
Yeniden hesaplamayı en aza indirmek istediğimizden sembolik hesaplamalar için bir kitaplığa ihtiyacımız var. XLA zaten MLIR'ye bağlı olduğundan, sembolik aritmetik kitaplığı yazmak yerine mlir::AffineMap'i kullanıyoruz.
Tipik bir AffineMap görünümü
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
AffineMap, kolayca kullanılabilen iki tür parametreye sahiptir: sırasıyla \(\boldsymbol d\) ve \(\boldsymbol s\) kullanabileceğimiz boyutlar ve semboller.
AffineMap, boyut aralıklarıyla ilgili meta veri içermediğinden bu verileri bizim sağlamamız gerekir.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges, dizine ekleme haritasının boyut parametreleri \(\boldsymbol{d}\) için kapsayıcı kutu kısıtlamalarını kodlar. Bu kısıtlama genellikle ters çevir, azalt, öğe düzeyinde, nokta gibi işlemler için çıkış tensörünün şekliyle örtüşür. Ancak HloConcatenateInstruction gibi bazı istisnalar vardır.
symbol_ranges, parametrelerin alabileceği olası değerleri \(\boldsymbol {s}\) kodlar.
Yukarıdakilerin gerçekte ne anlama geldiğini anlamak için tek tek inceleyelim.
Etkin Olmayan İşlemler için Haritalar Dizine Ekleme
Öğe yönü
Elementwise işlemleri için, dizine ekleme haritası bir kimliktir.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
Giriş haritalarına çıkış:
- çıkış -> giriş_0: \((d_0, d_1) \mapsto (d_0, d_1)\) için $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(output)$
- çıkış -> giriş_1: \((d_0, d_1) \mapsto (d_0, d_1)\) için $\boldsymbol{d} \in {\rm Dom} (çıkış)$
Çıkış haritalarına giriş
- giriş_i -> çıkış: \((d_0, d_1) \mapsto (d_0, d_1)\) $\boldsymbol{d} \in {\rm Dom}(output)$
Anons yap
Yayınlama, çıkışı girişle eşlediğimizde bazı boyutların kaldırılacağı ve girişi çıkışla eşleştirdiğimizde ekleneceği anlamına gelir.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
Giriş haritası için çıkış:
- çıkış -> giriş: \((d_0, d_1, d_2) \mapsto (d_1)\) $\boldsymbol{d} için \in {\rm Dom}(output)$
Çıkış haritası girişi
- giriş -> çıkış: \((d_0) \mapsto (s_0, d_1, s_1)\) $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9] \time [0, 29]$.
Şimdi ise giriş-çıkış eşlemesi \(\boldsymbol s\) için sağ taraftayız. Bunlar, değer aralıklarını temsil eden simgelerdir. Örneğin, bu özel örnekte dizin içeren \(d_0\) her giriş öğesi, çıkışın 10x1x30 boyutundaki bir dilimle eşlenir.
Constant ve Iota
Bu URL'lerde giriş parametreleri bulunmaz. Bu nedenle, dizine eklemenin hesaplanacağı bir durum yoktur.
Ters çevir
Ters çevirme için dizine ekleme haritası, giriş/çıkış boyutlarının permütasyonudur.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
Giriş haritası için çıkış:
- çıkış -> giriş:\(\boldsymbol{d} \in {\rm Dom}(output)\)için \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\)
Çıkış haritası girişi:
- giriş -> çıkış:\(\boldsymbol{d} \in {\rm Dom}(input)\)için \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\)
Ters Smaç
Ters yönde değişiklik için dizine ekleme haritası, geri alınan boyutları $upper_bound(d_i) - d_i$ olarak değiştirir:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
Giriş haritası için çıkış:
- çıkış -> giriş: $(d_0, d_1, d_2; d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(çıkış)$
Çıkış haritası girişi:
- giriş -> çıkış: $(d_0, d_1, d_2; d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(giriş)$
(Çeşitli)Azaltma
Değişken azaltmada birkaç giriş ve birkaç başlatma bulunur. Çıkıştan girdiye kadar olan harita, indirgenmiş boyutları ekler. Bir şekilde yayının tersi gibi davranır.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
Giriş haritalarına çıkış:
- çıkış -> giriş_j: \((d_0) \mapsto (s_0, d_0)\) için $\boldsymbol{d} \in {\rm Dom}(çıkış)\( and \)\boldsymbol{s} \0, 9]$
- çıktı -> init_j: \((d_0) \mapsto ()\) için $\boldsymbol{d} \in {\rm Dom}(output)$
Çıkış haritalarına giriş:
- giriş_i -> çıkış_j: \((d_0, d_1) \mapsto (d_1)\) için $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i -> exit_j: \(() \mapsto (s_0)\) için \(\boldsymbol{s} \in [0, 9]\)
\(i, j = 0, \ldots, INPUT\\_COUNT\)için.
Dilim
Dilim sonuçları için çıkıştan girişe dizine ekleme, çıkışın her öğesi için geçerli olan aşamalı bir dizine ekleme haritası oluşturur. Girişten çıkışa eşleme, girişteki öğelerin çizgili aralığıyla kısıtlıdır.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
Giriş haritası için çıkış:
- çıkış -> giriş:\(\boldsymbol{d} \in {\rm Dom}(output)\)için \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\)
Çıkış haritası girişi:
- giriş -> çıkış: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) için $[1, 7, 2]$aşamalı olmalıdır.
TBD: Girişten çıkışa dizine ekleme
Şekillendir
Şekillendirmeler farklı şekillerde sunulur.
Şekli daralt
Bu, N-D'den 1D'ye "doğrusallaştırma" türünde bir yeniden şekillendirmedir.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
Giriş haritası için çıkış:
- çıkış -> giriş: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) $\boldsymbol{d} için \in {\rm Dom}(output)$
Çıkış haritası girişi:
- giriş -> çıkış: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) $\boldsymbol{d} \in {\rm Dom}(input)$ için.
Şekli genişlet
Bu, ters bir "daraltma şekli" işlemidir, 1D girdiyi N-D çıkışına yeniden şekillendirir.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
Giriş haritası için çıkış:
- çıkış -> giriş: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) $\boldsymbol{d} için \in {\rm Dom}(output)$
Çıkış haritası girişi:
- giriş -> çıkış: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) $\boldsymbol{d} için {\rm Dom}(input)$.
Genel şekil değiştirme
Bunlar, tek bir genişletme veya daraltma şekli olarak temsil edilemeyecek şekilde yeniden şekillendirme işlemleridir. Yalnızca 2 veya daha fazla genişletme veya daraltma şeklinin bir bileşimi olarak temsil edilebilirler.
1. Örnek: Lineerleştirme-delineerleştirme.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
Bu yeniden şekil, tensor<4x8xf32> - tensor<32xf32> daralma şeklinin ve ardından şeklin tensor<2x4x4xf32> boyutuna genişlemesinin bir bileşimi olarak temsil edilebilir.
Giriş haritası için çıkış:
- çıkış -> giriş: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
\(\boldsymbol{d} \in {\rm Dom}(output)\)için
Çıkış haritası girişi:
- giriş -> çıkış: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
\(\boldsymbol{d} \in {\rm Dom}(input)\)için.
2. Örnek: Genişletilmiş ve daraltılmış alt şekiller
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
Bu yeniden şekillendirme, iki yeniden şekillendirmenin bileşimi olarak temsil edilebilir. İlki, en dıştaki boyutları tensor<4x8x12xf32> tensor<32x12xf32> olacak şekilde daraltır, ikincisi ise en içteki tensor<32x12xf32> boyutunu tensor<32x3x4xf32> olarak genişletir.
Giriş haritası için çıkış:
- çıkış -> giriş: \(\boldsymbol{d} \in {\rm Dom}(output)\)için \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\)
Çıkış haritası girişi:
- giriş -> çıkış: \(\boldsymbol{d} \in {\rm Dom}(input)\)için \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\).
Bitcast
Bitcast işlemi, transpose-reshape-transpose sırası olarak temsil edilebilir. Bu nedenle, dizine ekleme haritaları yalnızca bu adım sırası için dizine ekleme haritalarının bir bileşimidir.
Birleştir
Conctat için çıkıştan girişe eşleme, tüm girişler için tanımlanır ancak örtüşmeyen alan adlarıyla, yani aynı anda girişlerden yalnızca biri kullanılır.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
Giriş haritası için çıkış:
- çıkış -> giriş 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)süreyle
- çıkış -> giriş 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) $\boldsymbol{d} \in [0, 2] \times [50, 79]$
Çıkış haritası girişleri:
- giriş 1 -> çıkış: \((d_0, d_1) \mapsto (d_0, d_1)\) için $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- giriş 2 -> çıkış: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) için $\boldsymbol{d} \in {\rm Dom}(input_2)$.
Nokta (çıktıdan girişe uygulandı
Nokta için haritaların dizine eklenmesi, azaltma için kullanılanlara çok benzer.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
Giriş eşlemelerinin çıkışı:
- çıkış -> giriş_1:\(\boldsymbol{d} \in {\rm Dom}(output)\) ve \(\boldsymbol{s} \in [0, 255]\)için \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\)
- çıkış -> giriş_2:\(\boldsymbol{d} \in {\rm Dom}(output)\) ve \(\boldsymbol{s} \in [0, 255]\)için \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\)
Çıkış haritaları için girişler:
- giriş_1 -> çıkış: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) \(\boldsymbol{d} \in {\rm Dom}(input_1)\) ve \(\boldsymbol{s} \in [0, 63]\)için
- giriş_2 -> çıkış: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) \(\boldsymbol{d} \in {\rm Dom}(input_2)\) ve \(\boldsymbol{s} \in [0, 127]\)için
Pencereyi küçültme (TBD)
Ped (TBD)
Fusion için Haritalar Dizine Ekleme
Füzyon işlemi için dizine ekleme haritası, kümedeki her işlem için haritaların bir bileşimidir. Bazı girişler farklı erişim kalıplarıyla birkaç kez okunabilir.
Bir giriş, birden fazla dizine ekleme haritası
\(p_0 + p_0^T\)örneği
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
p0 için çıkıştan girişe dizine ekleme haritaları $(d_0, d_1) \mapsto (d_0, d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$ olacaktır. Bu, çıkışın bir öğesini hesaplamak için giriş parametresini iki kez okumamız gerekebileceği anlamına gelir.
Tek giriş, tekilleştirilmiş dizin haritası
Hemen belirgin olmasa bile, dizine ekleme haritalarının aslında aynı olduğu durumlar vardır.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
Bu durumda p0 için çıkıştan girişe dizin haritası yalnızca $(d_0, d_1, d_2)
\mapsto (d_2, d_0, d_1)$ şeklindedir.
Softmaks
Softmax için parameter 0 için çıkıştan girişe dizine ekleme eşlenir:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
\(\boldsymbol{d} \in {\rm Dom}(output)\) ve \(\boldsymbol{s} \in [0, 124]\) girinin en iç boyutunu ifade eder.
Dizine Ekleme Haritası Basitleştiricisi
mlir::AffineMap yukarı akış için varsayılan basitleştirici, boyut/simge aralıkları hakkında herhangi bir varsayımda bulunamaz. Bu nedenle, mod ve div ile ifadeleri verimli bir şekilde basitleştiremez.
Afin haritalardaki alt ifadelerin alt ve üst sınırları hakkındaki bilgiden bunları daha da basitleştirmek için kullanabiliriz.
Basitleştirici, aşağıdaki ifadeleri yeniden yazabilir.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1; d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \d_0 \d_0 \d_1 d_1 d_2 \d_1 d_1 d_1 d_1\( for \)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16 d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \0, 9]\( becomes \)(d_0, d_1, d_2 (d_2) d_2 (d_2) d_1 (d_2) d_2
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
Harita oluşturmayla ilgili basit bir uygulama, HLO'daki bazı zincirleme yeniden şekillendirmelerin birbirini iptal ettiğini anlamamızı sağlar.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
Dizine ekleme haritalarının birleştirilmesinden ve basitleştirildikten sonra,
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).