StableHLO'da dinamiklik

Dinamizmin mevcut durumu Dinamizm RFC'sinde daha resmi bir şekilde açıklanmaktadır. Bu sayfada, RFC'ye genel bir bakış sunulmakta ve dinamik programlarla etkileşime geçmek için önemli API'ler ve araçlar ele alınmaktadır.

Dinamizm Terminolojisi ve Desteğe Genel Bakış

Öncelikle, bu dokümanda yer alacak birkaç terimi ve bunların StableHLO'da desteklenmesiyle ilgili kısa bir girişi ele alalım:

Dinamik boyutlar

Dinamik boyutlar, boyutu bilinmeyen tüm boyutları ifade eder. StableHLO'da dinamik boyutları ? (ör. tensor<16x?xf32>) kullanarak temsil ederiz.

Sınırlı dinamizm

Sınırlı dinamizm, değeri bilinen bir üst sınıra sahip dinamik bir boyutu ifade eder. Bu genellikle, yürütme sırasında tenzoru doldurmak için yararlıdır. StableHLO'da tensör kodlaması olarak #stablehlo.bounds kullanarak sınırlı dinamizmi temsil ediyoruz. Yani dinamik boyutu 16'da sınırlı olan sıra 2 tensörü, dinamik boyutu 16'da sınırlanmış, diğeri ise sınır bulunmayan tensor<?x?xf32, #stablehlo.bounds<16, ?>> olarak temsil edilebilir.

StableHLO, sınırlı dinamizmi temsil edebilir ancak TensorFlow'dan gelen sınırlı çerçeve desteği ve PyTorch/XLA'da bazı destekler vardır.

Sınırsız dinamizm

Adından da anlaşılacağı gibi sınırsız dinamizm, boyutu için bilinen bir sınırı olmayan dinamik bir boyutu ifade eder. Bu dinamiklik türü, genellikle dinamik grup boyutuna veya dizi uzunluğuna sahip modelleri dışa aktarmak için kullanılan JAX, PyTorch/XLA ve TF destekli StableHLO'da çok yaygındır.

StableHLO'da, bu dinamizm biçimi için kodlanan sınırları (ör. tensor<?x?xf32>) kaldırırız.

Şekil polimorfizmi

Şekil polimorfizmi, JAX'tan aldığımız bir terimdir.

Polimorfizmi şekillendirmenin iki önemli sonucu vardır:

  1. Programdaki tüm dinamizm, giriş bağımsız değişkenlerine dayanır.
  2. Tüm dinamizm yalnızca tensör şekilleriyle ilgilidir, yani verilere bağlı değildir.

Bu iki kural sayesinde, bir programın statik şekilleri bilindiğinde dinamik bir programı derleme için tamamen statik bir programa dönüştürebiliriz ("Dinamik programları hassaslaştırmak için derleyici geçişleri" bölümüne bakın).

Genellikle şekil polimorfizmi sınırsız dinamizm kullanır. Bilinen bağımsız değişken şekilleri tamamen statik bir programa yol açabilirse değerlerin nasıl sınırlandırılacağını tahmin etmeniz gerekmez.

Verilere dayalı dinamizm

Veriye dayalı dinamiklik, bir tensör içindeki verilerle ilgili dinamik boyut boyutlarını ifade eder. Standart örnek, bir tensör değerinde 0 olan tüm öğelerin dizinlerini döndüren bir nonzeros işlevidir. Veriler değerlendirilmeden şekil bilinemez ancak genellikle sınırlı dinamizm kullanılarak derlenebilir. Bu durumda, potansiyel çıkış tenör boyutu için fazladan bellek harcanır.

Veriye dayalı birçok dinamik işlem, sınırlı dinamizm kullanılarak modellenebilir. Burada, bir tensör boyutunun üst sınırı belirtilir ve donanım, genellikle tensör dolgusu aracılığıyla bunu uygular. Şu anda PyTorch/XLA ve TensorFlow'da verilere dayalı dinamizm için bazı destekler mevcuttur ancak JAX şu anda verilere dayalı dinamizme yol açan işlemleri izlememektedir.

Dinamik boyutlara sahip programları dışa aktarma

Dinamik toplu boyutları veya sıra uzunlukları olan programların nasıl dışa aktarılacağı hakkında bilgi edinmek için StableHLO eğitimlerimizi inceleyin:

Dinamik programları hassaslaştırmak için derleyici geçişleri

Dinamizm geçiş ardışık düzenini kaldırma

Şekilleri hassaslaştırmak için birkaç yararlı geçiş vardır. Bunların tümü, geçiş ardışık düzeninde createStablehloRemoveDynamismPipeline gruplandırılmıştır:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Dinamikliği artırmak için bireysel kartlar

Şekil hassaslaştırma için yararlı olan geçişler tek tek şunlardır:

Güncel bilgiler ve örnekler için bağlantılı dokümanlara bakın.

Örnek: Dinamik dinamik nasıl faydalıdır ve nasıl kullanabilirim?

Dinamikliğin birçok kullanım alanı vardır. Bu bölümde, genellikle Şekil Çokgenliği için yaygın kullanım alanına odaklanacağız. Burada, genellikle dinamik grup boyutunu veya dizi uzunluğunu temsil etmek için kullanılan, dışa aktarılan esnek bir model temsili oluşturuyoruz.

Statik add_one modeli

Bunu göstermek için aşağıdaki basit add_one modelini kullanacağız:

def add_one(x):
  return x + 1

tensor<4xf32> kullanılarak izlendiğinde aşağıdaki StableHLO programı elde edilir:

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

Bu model yalnızca tensor<4xf32> şekline sahip giriş bağımsız değişkenleri için çalışır. Grup boyutumuzu veya dizi uzunluğumuzu değiştirseydik kaynak kodu yeniden izleyip StableHLO'ya geri düşürmemiz gerekir. Bu durumda kaynak koda hâlâ erişimimizin olacağını garanti edemeyiz.

Dinamik add_one modeli

Burada şekil polimorfik dinamikleri devreye girer. Bunun yerine JAX ve PyTorch/XLA, add_one modelini dinamik olarak geçerli IR ile yayınlayabilir. Bu durumda sabit, dinamik giriş şekliyle eşleşecek şekilde aşağıdaki gibi yayınlanır:

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

Bu model temsili çok daha esnektir ve toplu boyut veya sıra uzunluğu gibi değerlerin ertelenen şekilde belirtilmesine olanak tanır. Bu model, dinamik şekil desteğine sahip platformlarda (AI Edge gibi) dağıtılabilir veya bu dokümanda belirtilen dinamizm geçişleri kullanılarak hassaslaştırılabilir.

Dinamik modeli hassaslaştırma

Örneğin, aşağıdaki geçiş sıralaması bu programı tamamen hassaslaştırabilir:

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

Program aşamalı olarak şu şekilde dönüştürülür:

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}