Dizine Ekleme Analizi

HLO dizine ekleme analizi, bir tensörün öğelerinin "dizine ekleme haritaları" aracılığıyla birbiriyle nasıl ilişkili olduğunu açıklayan bir veri akışı analizidir. Örneğin, bir HLO talimatı çıkışının dizinleri, HLO talimatı işlenenlerinin dizinleriyle nasıl eşlenir?

Örnek

tensor<20xf32> - tensor<10x20x30xf32> arası yayın için

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

Çıkıştan girişe dizin oluşturma eşlemesi i in [0, 10], j in [0, 20] ve k in [0, 30] için (i, j, k) -> (j)'dır.

Motivasyon

XLA, birleştirme, işlenen kullanımı ve döşeme şemaları hakkında akıl yürütmek için çeşitli özel çözümler kullanır (daha fazla ayrıntı aşağıda verilmiştir). Dizin oluşturma analizinin amacı, bu tür kullanım alanları için yeniden kullanılabilir bir bileşen sağlamaktır. Dizin oluşturma analizi, MLIR'nin Affine Map altyapısı üzerine kurulmuştur ve HLO semantiğini ekler.

Birleşme

Girişlerin hangi öğelerinin/dilimlerinin çıkışın bir öğesini hesaplamak için okunduğunu bildiğimizde, bellek birleştirme hakkında akıl yürütmek önemsiz olmayan durumlarda mümkün hale gelir.

İşlenen Kullanımı

XLA'daki işlenen kullanımı, talimatın her girişinin ne kadar kullanıldığını gösterir. Bu, çıkışının tamamen kullanıldığı varsayılarak hesaplanır. Şu anda kullanım, genel bir durum için de hesaplanmamaktadır. Dizin oluşturma analizi, kullanımı hassas bir şekilde hesaplamamıza olanak tanır.

Döşeme

Döşeme/dilim, ofsetler, boyutlar ve adımlarla parametrelendirilmiş bir tensörün hiper dikdörtgen alt kümesidir. Döşeme yayılımı, işlemin döşeme parametrelerini kullanarak işlemin üreticisi/tüketicisinin döşeme parametrelerini hesaplamanın bir yoludur. Softmax ve nokta için bunu yapan bir kütüphane zaten var. Döşeme yayılımı, dizin oluşturma haritaları aracılığıyla ifade edilirse daha genel ve sağlam hale getirilebilir.

Dizine ekleme haritası

Dizin oluşturma haritası,

  • Bir tensörün A her öğesini B tensöründeki öğe aralıklarıyla eşleyen, sembolik olarak ifade edilmiş bir işlev;
  • İşlevin alanı da dahil olmak üzere geçerli işlev bağımsız değişkenleriyle ilgili kısıtlamalar.

İşlev bağımsız değişkenleri, doğalarını daha iyi ifade etmek için 3 kategoriye ayrılır:

  • Tensörün A veya eşleme yaptığımız bir GPU ızgarasının boyut değişkenleri; değerler statik olarak bilinir. Dizin öğeleri boyut değişkenleri olarak da adlandırılır.

  • range değişkenleri. Bunlar bire çok eşleme tanımlar ve B içinde A'nin tek bir değerini hesaplamak için kullanılan bir dizi öğe belirtir. Değerler statik olarak bilinir. Bir matris çarpımının daraltma boyutu, bir aralık değişkeni örneğidir.

  • Yalnızca yürütme sırasında bilinen çalışma zamanı değişkenleri. Örneğin, gather işleminin dizinler bağımsız değişkeni.

İşlevin sonucu, hedef B tensörünün bir dizinidir.

Kısacası, x işlemi için A tensöründen B tensörüne yönelik bir dizin oluşturma işlevi

map_ab(index in A, range variables, runtime variables) -> index in B.

Eşleme bağımsız değişkeni türlerini daha iyi ayırmak için bunları şu şekilde yazarız:

map_ab(index in A)[range variables]{runtime variables} -> (index in B)

Örneğin, azaltma işlemi için dizine ekleme haritalarına bakalım f32[4, 8] out = reduce(f32[2, 4, 8, 16] in, 0), dimensions={0,3}:

  • in öğelerini out ile eşlemek için işlevimiz (d0, d1, d2, d3) -> (d1, d2) olarak ifade edilebilir. Değişkenlerin kısıtlamaları d0 in [0, 1], d1 in [0, 3], d2 in [0, 7], d3 in [0, 15], in şekliyle tanımlanır.

  • out öğelerini in ile eşlemek için: out yalnızca iki boyuta sahiptir ve boyut azaltma, boyutları azaltmayı kapsayan iki aralık değişkeni sunar. Bu nedenle, eşleme işlevi (d0, d1)[s0, s1] -> (s0, d0, d1, s1)'dır. Burada (d0, d1), out'nin dizinidir. s0, s1, işlemin semantiği tarafından tanımlanan ve in tensörünün 0 ile 3 arasındaki aralığını kapsayan aralıklardır. Kısıtlamalar d0 in [0, 3], d1 in [0, 7], s0 in [0,1], s1 in [0, 15].

Çoğu senaryoda çıktı öğelerinden eşleme yapmak istediğimizi belirtmek önemlidir. Hesaplama için

C = op1(A, B)
E = op2(C, D)

Bu nedenle, "B'nin dizine eklenmesi" ifadesiyle "E öğelerinin B öğeleriyle eşlenmesi" kastedilir. Bu, girişten çıkışa doğru çalışan diğer veri akışı analizi türlerine kıyasla sezgisel olmayabilir.

Değişkenlerdeki kısıtlamalar, optimizasyon fırsatlarını etkinleştirir ve kod oluşturmaya yardımcı olur. Belgelerde ve uygulama kısıtlamalarında, eşleme işlevinin tüm geçerli kombinasyonlarını veya bağımsız değişken değerlerini tanımladıkları için alan olarak da adlandırılırlar. Kısıtlamalar, birçok işlem için yalnızca tensörlerin boyutlarını tanımlar ancak bazı işlemler için daha karmaşık olabilir. Aşağıdaki örneklere bakın.

İşlevlerin ve bağımsız değişken kısıtlamalarının sembolik olarak ifade edilmesi ve işlevlerle kısıtlamaların birleştirilebilmesi sayesinde, rastgele büyük bir hesaplama (birleştirme) için kompakt bir dizin oluşturma eşlemesi hesaplayabiliriz.

Sembolik işlevin ve kısıtlamaların ifade gücü, uygulama karmaşıklığı ile daha hassas bir gösterimden elde ettiğimiz optimizasyon kazançları arasındaki dengeyi ifade eder. Bazı HLO işlemleri için erişim kalıplarını yalnızca yaklaşık olarak yakalarız.

Uygulama

Yeniden hesaplamayı en aza indirmek istediğimiz için sembolik hesaplamalar için bir kitaplığa ihtiyacımız var. XLA zaten MLIR'ye bağlı olduğundan, başka bir sembolik aritmetik kitaplığı yazmak yerine mlir::AffineMap'i kullanırız.

Tipik bir AffineMap şöyle görünür:

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap iki tür parametre içerir: Boyutlar ve semboller. Boyutlar, boyut değişkenleri d'ye, semboller ise aralık değişkenleri r ve çalışma zamanı değişkenleri rt'ye karşılık gelir. AffineMap, parametrelerin kısıtlamalarıyla ilgili herhangi bir meta veri içermediğinden bunları ayrı olarak sağlamamız gerekir.

struct Interval {
 int64_t lower;
 int64_t upper;
};

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  mlir::AffineMap affine_map_;

  // DimVars represent dimensions of a tensor or of a GPU grid.
  std::vector<Variable> dim_vars_;

  // RangeVars represent ranges of values, e.g. to compute a single element of
  // the reduction's result we need a range of values from the input tensor.
  std::vector<Variable> range_vars_;

  // RTVars represent runtime values, e.g. a dynamic offset in
  // HLO dynamic-update-slice op.
  std::vector<Variable> rt_vars_;
  llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};

dim_vars_ Genellikle transpoze, azaltma, öğe bazında, nokta gibi işlemler için çıkış tensörünün şekliyle çakışan, indeksleme haritasının boyut değişkenleri d için kapsayıcı kutu kısıtlamalarını kodlayın ancak HloConcatenateInstruction gibi bazı istisnalar vardır.

range_vars_, aralık değişkenlerinin s aldığı tüm değerler. Aralık değişkenleri, eşleme yaptığımız tensörün tek bir öğesini hesaplamak için birden fazla değer gerektiğinde (ör. indirgemelerin çıkış->giriş dizin oluşturma eşlemesi veya yayınların giriş->çıkış eşlemesi için) gereklidir.

rt_vars_, olası değerleri çalışma zamanında kodlayın. Örneğin, 1 boyutlu HloDynamicSliceInstruction için ofset dinamiktir. İlgili RTVar, 0 ile tensor_size - slice_size - 1 arasında uygun değerlere sahip olacaktır.

constraints_ biçimindeki değerler arasındaki ilişkileri yakalayın <expression> in <range> (ör. d0 + s0 in [0, 20]). Variable.bounds ile birlikte, dizine ekleme işlevinin "alanını" tanımlar.

Yukarıdakilerin hepsinin ne anlama geldiğini anlamak için örneklerle inceleyelim.

Birleştirilmemiş işlemler için Haritalar'ı dizine ekleme

Elementwise

Öğe bazında işlemler için dizin oluşturma eşlemesi bir kimliktir.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  output = f32[10, 20] add(p0, p1)

Çıkış-giriş eşlemesi output -> p0:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

Giriş-çıkış eşlemesi p0 -> output:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

Yayın

Yayın, çıkışı girişe eşlediğimizde bazı boyutların kaldırılacağı, girişi çıkışa eşlediğimizde ise ekleneceği anlamına gelir.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

Çıkış-giriş haritası:

(d0, d1, d2) -> (d1),
domain:
d0 in [0, 9],
d1 in [0, 19],
d2 in [0, 29]

Girişten çıkışa eşleme:

(d0)[s0, s1] -> (s0, d0, s1),
domain:
d0 in [0, 19],
s0 in [0, 9],
s1 in [0, 29]

Giriş-çıkış eşlemesi için sağ tarafta artık aralık değişkenleri s olduğunu unutmayın. Bunlar, değer aralıklarını temsil eden sembollerdir. Örneğin, bu özel durumda, d0 dizinli girişin her öğesi, çıkışın 10x1x30 dilimine eşlenir.

Iota

Iota'nın giriş tensörü işleneni olmadığından giriş dizini bağımsız değişkeni yoktur.

iota = f32[2,4] iota(), dimensions={1}

Çıkış-giriş haritası:

(d0, d1) -> ()
domain:
d0 in [0, 1]
d1 in [0, 3]

Girişten çıkışa eşleme:

()[s0, s1] -> (s0, s1)
domain:
s0 in [0, 1]
s1 in [0, 3]

DynamicSlice

DynamicSlice'ın yalnızca çalışma zamanında bilinen ofsetleri vardır.

src = s32[2, 2, 258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = s32[1, 2, 32] dynamic-slice(src, of1, of2, of3), dynamic_slice_sizes={1, 2, 32}

ds ile src arasındaki giriş-çıkış eşlemesi:

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

Giriş-çıkış eşlemesi için sağ tarafta rt olduğunu unutmayın. Bunlar, çalışma zamanı değerlerini temsil eden sembollerdir. Örneğin, bu özel durumda, girişin dizinini hesaplamak için d0, d1, d2 dizinli çıkışın her öğesi için dilim ofsetlerine (of1, of2 ve of3) erişiriz. Çalışma zamanı değişkenlerinin aralıkları, dilimin tamamının sınırlar içinde kaldığı varsayılarak elde edilir.

of1, of2 ve of3 için çıkış-giriş haritası:

(d0, d1, d2) -> (),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31]

DynamicUpdateSlice

src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
    s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)

src için giriş haritasına çıkış önemsizdir. Alan, güncellenmemiş dizinlerle sınırlandırılarak daha hassas hale getirilebilir ancak şu anda dizine ekleme haritaları eşitsizlik kısıtlamalarını desteklemiyor.

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 19],
d1 in [0, 29]

upd için giriş/çıkış haritası:

(d0, d1){rt0, rt1} -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

Artık çalışma zamanı değerlerini temsil eden rt0 ve rt1 olduğunu unutmayın. Bu özel durumda, çıkışın d0, d1 dizinli her öğesi için girişin dizinini hesaplamak üzere dilim ofsetlerine of1 ve of2 erişiriz. Çalışma zamanı değişkenlerinin aralıkları, dilimin tamamının sınırlar içinde kaldığı varsayılarak türetilir.

of1 ve of2 için çıkış-giriş eşlemesi:

(d0, d1) -> (),
domain:
d0 in [0, 19],
d1 in [0, 29]

Toplama

Yalnızca basitleştirilmiş toplama desteklenir. gather_simplifier.h dosyasına bakın.

operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
  offset_dims={1,2,3},
  collapsed_slice_dims={},
  start_index_map={0,1},
  index_vector_dim=1,
  slice_sizes={7,8,4}

operand için giriş/çıkış haritası:

(d0, d1, d2, d3){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

Artık çalışma zamanı değerlerini temsil eden rt sembollerimiz olduğunu unutmayın.

indices için giriş/çıkış haritası:

(d0, d1, d2, d3)[s0] -> (d0, s0),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
s0 in [0, 1]

Aralık değişkeni s0, çıkışın bir öğesini hesaplamak için indices tensörünün tüm satırına (d0, *) ihtiyacımız olduğunu gösterir.

Transpoze

Transpoze için indeksleme 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}

Çıkış-giriş haritası:

(d0, d1, d2, d3) -> (d0, d3, d1, d2),
domain:
d0 in [0, 2],
d1 in [0, 5],
d2 in [0, 127],
d3 in [0, 12287],

Girişten çıkışa eşleme:

(d0, d1, d2, d3) -> (d0, d2, d3, d1),
domain:
d0 in [0, 2],
d1 in [0, 12287],
d2 in [0, 5],
d3 in [0, 127]

Tersine çevir

Geri alınan değişiklikler 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}

Çıkış-giriş haritası:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

Girişten çıkışa eşleme:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

(Variadic)Reduce

Çok değişkenli azaltma işleminde birden fazla giriş ve birden fazla başlangıç değeri vardır. Çıkıştan girişe eşleme, azaltılmış boyutları ekler.

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
out = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=max

Giriş haritalarına çıkış:

  • out[0] -> p0:
(d0)[s0] -> (s0, d0),
domain:
d0 in [0, 9],
s0 in [0, 255]
  • out[0] -> p0_init:
(d0) -> (),
domain:
d0 in [0, 9]

Girişten çıkışa eşlemeler:

  • p0 -> out[0]:
(d0, d1) -> (d1),
domain:
d0 in [0, 255],
d1 in [0, 9]
  • p0_init -> out[0]:
()[s0] -> (s0),
domain:
s0 in [0, 9]

Dilim

Dilim sonuçları için çıkıştan girişe dizine ekleme, çıkışın her öğesi için geçerli olan adımlı bir dizine ekleme haritası oluşturur. Girişten çıkışa eşleme, girişin öğelerinde adımlı bir aralıkla sınırlı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]}

Çıkış-giriş haritası:

(d0, d1, d2) -> (d0 + 5, d1 * 7 + 3, d2 * 2),
domain:
d0 in [0, 4],
d1 in [0, 2],
d2 in [0, 24]

Girişten çıkışa eşleme:

(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2),
domain:
d0 in [5, 9],
d1 in [3, 17],
d2 in [0, 48],
(d1 - 3) mod 7 in [0, 0],
d2 mod 2 in [0, 0]

Yeniden şekillendirme

Yeniden şekillendirme işlemleri farklı türlerde yapılabilir.

Şekli daraltma

Bu, N boyutlu diziden 1 boyutlu diziye "doğrusallaştırıcı" bir yeniden şekillendirmedir.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

Çıkış-giriş haritası:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

Girişten çıkışa eşleme:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

Şekli genişletme

Bu, ters "şekli daraltma" işlemidir. 1 boyutlu girişi N boyutlu çıkış olarak yeniden şekillendirir.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

Çıkış-giriş haritası:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

Girişten çıkışa eşleme:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

Genel yeniden şekillendirme

Bunlar, tek bir genişletme veya daraltma şekli olarak gösterilemeyen yeniden şekillendirme işlemleridir. Bu şekiller yalnızca 2 veya daha fazla genişletme ya da daraltma şeklinin birleşimi olarak gösterilebilir.

1. örnek: Doğrusallaştırma-doğrusallaştırmayı kaldırma.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

Bu yeniden şekillendirme, tensor<4x8xf32> şeklinin tensor<32xf32> şekline daraltılması ve ardından tensor<2x4x4xf32> şekline genişletilmesi olarak gösterilebilir.

Çıkış-giriş haritası:

(d0, d1, d2) -> (d0 * 2 + d1 floordiv 2, d2 + (d1 mod 2) * 4),
domain:
d0 in [0, 1],
d1 in [0, 3],
d2 in [0, 3]

Girişten çıkışa eşleme:

(d0, d1) -> (d0 floordiv 2, d1 floordiv 4 + (d0 mod 2) * 2, d1 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7]
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 birleşimi olarak gösterilebilir. Birincisi, en dıştaki boyutları tensor<4x8x12xf32> ile tensor<32x12xf32> arasında daraltır. İkincisi ise en içteki boyutu tensor<32x12xf32> ile tensor<32x3x4xf32> arasında genişletir.

Çıkış-giriş haritası:

(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2),
domain:
d0 in [0, 31],
d1 in [0, 2]
d2 in [0, 3]

Girişten çıkışa eşleme:

(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7],
d2 in [0, 11]

Bitcast

Bir bitcast işlemi, devrikleştirme-yeniden şekillendirme-devrikleştirme dizisi olarak gösterilebilir. Bu nedenle, dizine ekleme haritaları yalnızca bu dizinin dizine ekleme haritalarının bir bileşimidir.

CONCATENATE

Birleştirme için çıkış-giriş eşlemesi tüm girişler için tanımlanır ancak çakışmayan alanlarla (yani girişlerden yalnızca biri aynı anda kullanılır).

p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT output = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}

Girişlere karşılık gelen çıkışlar:

  • output -> p0:
(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]
  • output -> p1:
(d0, d1, d2) -> (d0, d1 - 5, d2),
domain:
d0 in [0, 1],
d1 in [5, 15],
d2 in [0, 6]
  • output -> p2:
(d0, d1, d2) -> (d0, d1 - 16, d2),
domain:
d0 in [0, 1],
d1 in [16, 32],
d2 in [0, 6]

Çıkış haritalarının girişleri:

  • p0 -> output:
(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]
  • p1 -> output:
(d0, d1, d2) -> (d0, d1 + 5, d2),
domain:
d0 in [0, 1],
d1 in [0, 10],
d2 in [0, 6]
  • p2 -> output:
(d0, d1, d2) -> (d0, d1 + 16, d2),
domain:
d0 in [0, 1],
d1 in [0, 16],
d2 in [0, 6]

Dot

Nokta için haritaların indekslenmesi, azaltma için haritaların indekslenmesine çok benzer.

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
output = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

Girişlere karşılık gelen çıkışlar:

  • output -> p0:
(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]
  • output -> p1:
(d0, d1, d2)[s0] -> (d0, s0, d2),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

Çıkış haritalarının girişleri:

  • p0 -> output:
(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 255],
s0 in [0, 63]
  • p1 -> output:
(d0, d1, d2)[s0] -> (d0, s0, d1),
domain:
d0 in [0, 3],
d1 in [0, 255],
d2 in [0, 63],
s0 in [0, 127]

Pad

PadOp'un dizine eklenmesi, SliceOp'un dizine eklenmesinin tersidir.

p0 = f32[4, 4] parameter(0)
p1 = f32[] parameter(1)
pad = f32[12, 16] pad(p0, p1), padding=1_4_1x4_8_0

Dolgu yapılandırması 1_4_1x4_8_0, lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1 değerini gösterir.

Giriş haritalarına çıkış:

  • output -> p0:
(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4),
domain:
d0 in [1, 7],
d1 in [4, 7],
(d0 - 1) mod 2 in [0, 0]
  • output -> p1:
(d0, d1) -> (),
domain:
d0 in [0, 11],
d1 in [0, 15]

ReduceWindow

XLA'daki ReduceWindow da dolgu işlemi gerçekleştirir. Bu nedenle, dizine ekleme haritaları, dolgu yapmayan ReduceWindow dizine ekleme ve PadOp'un dizine eklemesinin birleşimi olarak hesaplanabilir.

c_inf = f32[] constant(-inf)
p0 = f32[1024, 514] parameter(0)
outpu = f32[1024, 3] reduce-window(p0, c_inf),
  window={size=1x512 pad=0_0x0_0}, to_apply=max

Giriş haritalarına çıkış:

  • output -> p0:
(d0, d1)[s0] -> (d0, d1 + s0),
domain:
d0 in [0, 1023],
d1 in [0, 2],
s0 in [0, 511]
  • output -> c_inf:
(d0, d1) -> (),
domain:
d0 in [0, 1023],
d1 in [0, 2]

Füzyon için Haritaları Dizine Ekleme

Birleştirme işlemi için dizine ekleme haritası, kümedeki her işlem için dizine ekleme haritalarının bir bileşimidir. Bazı girişler farklı erişim kalıplarıyla birkaç kez okunabilir.

Tek giriş, birden fazla dizine ekleme haritası

p0 + transpose(p0) için bir örnek verilmiştir.

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ış-giriş indeksleme eşlemleri (d0, d1) -> (d0, d1) ve (d0, d1) -> (d1, d0) olacaktır. Bu, çıkışın bir öğesini hesaplamak için giriş parametresini iki kez okumamız gerekebileceği anlamına gelir.

Tek giriş, yinelenenleri kaldırma işlemi uygulanmış dizine ekleme haritası

img

Dizine ekleme haritaları, hemen anlaşılmasa da aslında aynı olabilir.

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 output = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}

Bu durumda p0 için çıkış-giriş indeksleme eşlemesi yalnızca (d0, d1, d2) -> (d2, d0, d1) olur.

Softmax

img

Softmax için parameter 0 çıkış-giriş dizinleri eşlemesi:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 1],
d1 in [0, 64],
d2 in [0, 124],
s0 in [0, 124]

ve

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 64],
d2 in [0, 124]

Burada s0, girişin en içteki boyutunu ifade eder.

Daha fazla örnek için indexing_analysis_test.cc sayfasına bakın.

Endeksleme Haritası Basitleştiricisi

mlir::AffineMap yukarı akışı için varsayılan basitleştirici, boyut/sembol aralıkları hakkında herhangi bir varsayımda bulunamaz. Bu nedenle, mod ve div içeren ifadeleri verimli bir şekilde basitleştiremez.

Daha da basitleştirmek için afin haritalardaki alt ifadelerin alt ve üst sınırlarıyla ilgili bilgilerden yararlanabiliriz.

Sadeleştirici, aşağıdaki ifadeleri yeniden yazabilir.

  1. [0, 6] x [0, 14] adlı kuruluş biriminde d için (d0, d1) -> (d0 + d1 floordiv 16, d1 mod 16), (d0, d1) -> (d0, d1) olur
  2. di in [0, 9] için (d0, d1, d2) -> ((100d0 + 10d1 + d2) floorDiv 100, ((100d0 + 10d1 + d2) mod 100) floordiv 10, d2 mod 10), (d0, d1, d2) -> (d0, d1, d2) olur.
  3. d_i in [0, 9] için (d0, d1, d2) -> ((16d0 + 4d1 + d2) floordiv 8, (16d0 + 4d1 + d2) mod 8), (d0, d1, d2) -> (2d0 + (4d1 + d2) floordiv 8,(4d1 + d2) mod 8) olur.
  4. [0, 9] x [0, 10] ülkesindeki d için (d0, d1) -> (-(-11d0 - d1 + 109) floordiv 11 + 9), (d0, d1) -> (d0) olur.

Dizin oluşturma haritası basitleştiricisi, HLO'daki bazı zincirlenmiş 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)

Dizin oluşturma haritalarının oluşturulup basitleştirilmesinin ardından

(d0, d1, d2) -> (d0, d1, d2).

Dizin oluşturma haritasının basitleştirilmesi, kısıtlamaları da basitleştirir.

  1. lower_bound <= affine_expr (floordiv, +, -, *) constant <= upper_bound türündeki kısıtlamalar updated_lower_bound <= affine_expr <= updated_upped_bound olarak yeniden yazılır.
  2. Her zaman karşılanan kısıtlamalar (ör. d0 + s0 in [0, 20] for d0 in [0, 5] ve s0 in [1, 3]) ortadan kaldırılır.
  3. Kısıtlamalardaki afin ifadeler, yukarıdaki dizine ekleme afin haritası olarak optimize edilir.

Daha fazla örnek için indexing_map_test.cc sayfasına bakın.