Dinamismo in StableHLO

Lo stato attuale del dinamismo è descritto in modo più formale nella RFC sul dinamismo. Questa pagina fornisce una panoramica di alto livello della RFC e descrive API e strumenti importanti per interagire con i programmi dinamici.

Terminologia del dinamismo e panoramica dell'assistenza

Innanzitutto, per trattare alcuni termini che appariranno in questo documento, nonché una breve introduzione al loro supporto in StableHLO:

Dimensioni dinamiche

Le dimensioni dinamiche si riferiscono a qualsiasi dimensione di cui non è nota la dimensione. In StableHLO rappresentiamo le dimensioni dinamiche utilizzando ?, ovvero tensor<16x?xf32>.

Dinamismo limitato

Il dinamismo limitato si riferisce a una dimensione dinamica il cui valore ha un limite superiore noto. In genere, questa opzione è utile per il padding del tensore durante l'esecuzione. In StableHLO rappresentiamo il dinamismo limitato utilizzando #stablehlo.bounds come codifica del tensore, ovvero un tensore di rango 2 con una dimensione dinamica limitata a 16 e l'altra senza limite può essere rappresentato come tensor<?x?xf32, #stablehlo.bounds<16, ?>>.

StableHLO è in grado di rappresentare il dinamismo limitato, ma il supporto del framework è limitato, ha origine in TensorFlow e con un certo supporto in PyTorch/XLA.

Dinamismo illimitato

Come suggerisce il nome, il dinamismo illimitato si riferisce a una dimensione dinamica senza limiti noti per le dimensioni. Questo tipo di dinamismo è molto comune in StableHLO, con supporto JAX, PyTorch/XLA e TF, spesso utilizzato per esportare modelli con dimensioni batch o lunghezza della sequenza dinamiche.

In StableHLO omettiamo semplicemente la codifica dei limiti per questa forma di dinamismo, ovvero tensor<?x?xf32>.

Polimorfismo delle forme

Il polimorfismo della forma è un termine che abbiamo ereditato da JAX.

Il polimorfismo della forma ha due implicazioni chiave:

  1. Tutto il dinamismo del programma risale ai suoi argomenti di input.
  2. Tutto il dinamismo riguarda solo le forme dei tensori, ovvero non dipende dai dati.

Con queste due regole, una volta note le forme statiche di un programma, siamo in grado di prendere un programma dinamico e perfezionarlo completamente in un programma statico per la compilazione (vedi "Passaggi del compilatore per perfezionare i programmi dinamici").

In genere, il polimorfismo delle forme utilizza un dinamismo illimitato. Se le forme degli argomenti noti possono portare a un programma completamente statico, non è necessario indovinare come limitare i valori.

Dinamismo basato sui dati

Il dinamismo dipendente dai dati si riferisce alle dimensioni dinamiche delle dimensioni che riguardano i dati all'interno di un tensore. L'esempio canonico è una funzione nonzeros che restituisce gli indici di tutti gli elementi che sono 0 in un valore tensore. La forma non può essere nota senza valutare i dati, ma spesso può essere compilata utilizzando il dinamismo limitato, spendendo memoria aggiuntiva per le dimensioni potenziali del tensore di output.

Molte operazioni dinamiche dipendenti dai dati possono essere modellate utilizzando il dinamismo limitato, in cui viene specificato un limite superiore per la dimensione di un tensore e l'hardware in genere implementa questa operazione tramite il padding del tensore. Oggi esiste un certo supporto per il dinamismo dipendente dai dati in PyTorch/XLA e TensorFlow, ma JAX al momento non traccia le operazioni che portano al dinamismo dipendente dai dati.

Esportazione di programmi con dimensioni dinamiche

Consulta i nostri tutorial su StableHLO per informazioni su come esportare programmi con dimensioni batch o lunghezze di sequenza dinamiche:

Passaggi del compilatore per perfezionare i programmi dinamici

Rimuovere la pipeline di trasferimento dinamico

Esistono alcuni pass utili per perfezionare le forme, tutti raggruppati in una pipeline di pass createStablehloRemoveDynamismPipeline:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Passaggi individuali per perfezionare il dinamismo

Singolarmente, i passaggi che tendono a essere utili per il perfezionamento della forma sono:

Per informazioni ed esempi aggiornati, consulta la documentazione collegata.

Esempio: in che modo il dinamismo è utile e come posso utilizzarlo?

Il dinamismo ha molti utilizzi. Qui ci concentreremo principalmente sul caso d'uso comune per il polimorfismo delle forme: la creazione di una rappresentazione flessibile del modello esportato, generalmente utilizzata per rappresentare le dimensioni del batch dinamiche o la lunghezza della sequenza.

Modello statico add_one

Per dimostrarlo, utilizzeremo il seguente modello add_one semplice:

def add_one(x):
  return x + 1

Se viene tracciato utilizzando un tensor<4xf32>, otterremo il seguente programma StableHLO:

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

Questo modello funzionerà solo per gli argomenti di input che hanno una forma tensor<4xf32>. Se dovessimo modificare le dimensioni del batch o la lunghezza della sequenza, dovremmo tracciare nuovamente il codice sorgente e abbassarlo di nuovo a StableHLO e non è garantito che abbiamo ancora accesso al codice sorgente.

Modello dinamico add_one

È qui che entra in gioco il dinamismo polimorfico delle forme. Invece, JAX e PyTorch/XLA possono emettere il modello add_one con IR valido dinamicamente che trasmetterà la costante in modo che corrisponda alla forma dell'input dinamico come segue:

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

Questa rappresentazione del modello è molto più flessibile e consente la specifica differita di valori come la dimensione del batch o la lunghezza della sequenza. Questo modello può essere implementato su piattaforme con supporto di forme dinamiche (come AI Edge) oppure può essere perfezionato utilizzando i passaggi di dinamismo menzionati in questa documentazione.

Perfezionamento del modello dinamico

Ad esempio, il seguente ordine di pass può perfezionare completamente questo programma:

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

In modo incrementale, ecco come viene trasformato il programma:

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