StableHLO'da dinamiklik

Dinamizmin mevcut durumu, Dinamizm RFC'de daha resmi bir şekilde açıklanmaktadır. Bu sayfada, RFC'ye ilişkin üst düzey bir genel bakış sunulacak ve dinamik programlarla etkileşim kurmak için önemli API'ler ve araçlar ele alınacaktır.

Dinamizm Terminolojisi ve Destek Ekibine Genel Bakış

Öncelikle, bu dokümanda yer alacak birkaç terimi ve StableHLO'daki destekleriyle ilgili kısa bir giriş yapalım:

Dinamik boyutlar

Dinamik boyutlar, boyut boyutu bilinmeyen tüm boyutları ifade eder. StableHLO'da dinamik boyutları ? ile gösteririz. Örneğin, tensor<16x?xf32>.

Sınırlı dinamizm

Sınırlı dinamizm, değeri bilinen bir üst sınıra sahip olan dinamik boyutu ifade eder. Genellikle bu, yürütme sırasında tensörü doldurmak için yararlıdır. StableHLO'da sınırlı dinamizmi tensör kodlaması olarak #stablehlo.bounds ile gösteririz. Örneğin, bir boyutu 16 ile sınırlanmış ve diğer boyutu sınırlanmamış dinamik bir boyuta sahip 2 sıralı tensör tensor<?x?xf32, #stablehlo.bounds<16, ?>> olarak gösterilebilir.

StableHLO, sınırlı dinamizmi temsil edebilir ancak TensorFlow'dan kaynaklanan ve PyTorch/XLA'da sınırlı destek sunan bir çerçeve desteği vardır.

Sınırsız dinamizm

Adından da anlaşılacağı gibi sınırsız dinamizm, boyutuyla ilgili bilinen bir sınır olmayan dinamik bir boyutu ifade eder. Bu tür dinamizm, JAX, PyTorch/XLA ve TF desteğiyle StableHLO'da çok yaygındır. Genellikle dinamik toplu iş boyutu veya dizi uzunluğuna sahip modelleri dışa aktarmak için kullanılır.

StableHLO'da bu dinamizm biçimi için sınır kodlaması atlanır. Örneğin, tensor<?x?xf32>.

Şekil polimorfizmi

Şekil polimorfizmi, JAX'ten devraldığımız bir terimdir.

Şekil polimorfizmiyle ilgili iki önemli sonuç vardır:

  1. Programdaki tüm dinamizm, giriş bağımsız değişkenlerinden kaynaklanır.
  2. Tüm dinamizm yalnızca tensör şekilleri ile ilgilidir, yani veriye bağlı değildir.

Bu iki kural sayesinde, bir programın statik şekilleri bilindiğinde dinamik bir programı alıp derleme için statik bir programa dönüştürebiliriz (bkz. "Dinamik programları iyileştirmeye yönelik derleyici geçişleri").

Genel olarak ş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ğı tahmin edilmesine gerek yoktur.

Verilere bağlı dinamizm

Verilere bağlı dinamizm, 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. Şekil, veriler değerlendirilmeden bilinemez ancak genellikle sınırlı dinamizm kullanılarak derlenebilir. Bu durumda, olası çıkış tensörü boyutu için ek bellek harcanır.

Verilere bağlı birçok dinamik işlem, sınırlı dinamizm kullanılarak modellenebilir. Bu durumda, tensör boyutu için bir üst sınır belirtilir ve donanım genellikle bunu tensör dolgusu aracılığıyla uygular. PyTorch/XLA ve TensorFlow'da veriye bağlı dinamizm için şu anda bir miktar destek bulunmaktadır ancak JAX, veriye bağlı dinamizme yol açan işlemleri şu anda izlemez.

Dinamik boyutlara sahip programları dışa aktarma

Dinamik toplu iş boyutlarına veya sıra uzunluklarına sahip programları dışa aktarma hakkında bilgi edinmek için StableHLO eğitimlerimize göz atın:

Dinamik programları iyileştirmeye yönelik derleyici geçişleri

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

Şekilleri iyileştirmek için kullanabileceğiniz birkaç faydalı geçiş vardır. Bunların tümü, geçiş hattında createStablehloRemoveDynamismPipeline bir araya getirilmiştir:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Dinamizmi hassaslaştırmak için tek tek kartlar

Şekil iyileştirme için yararlı olan geçişler şunlardır:

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

Örnek: Dinamizm ne işe yarar ve nasıl kullanabilirim?

Dinamizmin birçok kullanım alanı vardır. Burada, Shape Polymorphism'in yaygın kullanım alanı olan, genellikle dinamik toplu iş boyutu veya sıra uzunluğunu temsil etmek için kullanılan esnek bir dışa aktarılmış model temsili oluşturmaya odaklanacağız.

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. Toplu iş boyutumuz veya sıra uzunluğumuz değişirse kaynak kodu yeniden izlememiz ve StableHLO'ya yeniden düşürmemiz gerekir. Kaynak koduna erişimimizin devam ettiğine dair bir garanti yoktur.

Dinamik add_one modeli

Bu noktada şekil polimorfik dinamizmi devreye girer. Bunun yerine JAX ve PyTorch/XLA, sabiti dinamik giriş şekliyle eşleşecek şekilde yayınlayacak olan dinamik olarak geçerli IR ile add_one modelini yayınlayabilir:

// 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 gösterimi çok daha esnektir ve toplu iş boyutu veya sıra uzunluğu gibi değerlerin ertelenmiş olarak belirtilmesine olanak tanır. Bu model, dinamik şekil desteği olan platformlarda (ör. AI Edge) dağıtılabilir veya bu dokümanda belirtilen dinamizm geçişleri kullanılarak iyileştirilebilir.

Dinamik modeli iyileştirme

Örneğin, aşağıdaki geçiş sıralaması bu programı tamamen iyileştirebilir:

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

Programın dönüşümü aşamalı olarak şu şekilde gerçekleşir:

// 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>
}