Dynamisme dans StableHLO

L'état actuel du dynamisme est décrit plus en détail dans la RFC sur le dynamisme. Cette page fournit un aperçu général de la RFC et présente les API et les outils importants pour interagir avec les programmes dynamiques.

Terminologie et présentation de l'assistance concernant le dynamisme

Tout d'abord, pour couvrir quelques termes qui apparaîtront dans ce document, ainsi qu'une brève introduction à leur prise en charge dans StableHLO :

Dimensions dynamiques

Les dimensions dynamiques désignent toutes les dimensions dont la taille est inconnue. Dans StableHLO, nous représentons les dimensions dynamiques à l'aide de ?, c'est-à-dire tensor<16x?xf32>.

Dynamisme limité

Le dynamisme limité fait référence à une dimension dynamique dont la valeur a une limite supérieure connue. En général, cela est utile pour ajouter un remplissage au Tensor lors de l'exécution. Dans StableHLO, nous représentons le dynamisme limité à l'aide de #stablehlo.bounds en tant que tensor encoding, c'est-à-dire qu'un Tensor de rang 2 avec une dimension dynamique limitée à 16 et l'autre sans limite peut être représenté par tensor<?x?xf32, #stablehlo.bounds<16, ?>>.

StableHLO est capable de représenter le dynamisme limité, mais la prise en charge du framework est limitée. Elle provient de TensorFlow et est partiellement disponible dans PyTorch/XLA.

Dynamisme illimité

Comme son nom l'indique, le dynamisme illimité fait référence à une dimension dynamique dont la taille n'est pas limitée. Ce type de dynamisme est très courant dans StableHLO, avec la prise en charge de JAX, PyTorch/XLA et TF, souvent utilisé pour exporter des modèles avec une taille de lot ou une longueur de séquence dynamiques.

Dans StableHLO, nous élidons simplement l'encodage des limites pour cette forme de dynamisme, c'est-à-dire tensor<?x?xf32>.

Polymorphisme de forme

Le polymorphisme de forme est un terme que nous avons hérité de JAX.

Le polymorphisme de forme a deux implications clés :

Tout le dynamisme du programme provient de ses arguments d'entrée.
Tout dynamisme ne concerne que les formes des Tensors, c'est-à-dire qu'il n'est pas dépendant des données.

Grâce à ces deux règles, une fois les formes statiques d'un programme connues, nous pouvons prendre un programme dynamique et l'affiner complètement en un programme statique pour la compilation (voir "Passes du compilateur pour affiner les programmes dynamiques").

En général, le polymorphisme de forme utilise un dynamisme illimité. Si les formes d'arguments connues peuvent conduire à un programme entièrement statique, il n'est pas nécessaire de deviner comment limiter les valeurs.

Dynamisme dépendant des données

Le dynamisme dépendant des données fait référence aux tailles de dimensions dynamiques qui se rapportent aux données à l'intérieur d'un Tensor. L'exemple canonique est une fonction nonzeros qui renvoie les index de tous les éléments 0 dans une valeur Tensor. La forme ne peut pas être connue sans évaluer les données, mais elle peut souvent être compilée à l'aide du dynamisme limité, en dépensant de la mémoire supplémentaire sur la taille potentielle du Tensor de sortie.

De nombreuses opérations dynamiques dépendant des données peuvent être modélisées à l'aide du dynamisme limité, où une limite supérieure de la taille d'un Tensor est spécifiée, et le matériel implémentera généralement cela via le remplissage du Tensor. Aujourd'hui, PyTorch/XLA et TensorFlow prennent en charge le dynamisme dépendant des données, mais JAX ne trace pas actuellement les opérations qui conduisent à un dynamisme dépendant des données.

Exporter des programmes avec des dimensions dynamiques

Consultez nos tutoriels StableHLO pour savoir comment exporter des programmes avec des tailles de lot ou des longueurs de séquence dynamiques :

Passages du compilateur pour affiner les programmes dynamiques

Supprimer le pipeline de transmission du dynamisme

Il existe quelques passes utiles pour affiner les formes. Elles sont toutes regroupées dans un pipeline de passes createStablehloRemoveDynamismPipeline :

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Passes individuelles pour affiner le dynamisme

Les passes qui sont généralement utiles pour affiner la forme sont les suivantes :

stablehlo-refine-arguments pour remplacer les arguments d'entrée par des types de Tensor concrets.
stablehlo-refine-shapes pour propager les nouvelles informations sur la forme de l'argument d'entrée dans l'ensemble du programme.
stablehlo-canonicalize-dynamism pour remplacer les opérations dynamiques par leurs variantes statiques.
stablehlo-check-shape-assertions pour vérifier et supprimer les appels personnalisés d'assertions de forme.

Consultez la documentation associée pour obtenir des informations et des exemples à jour.

Exemple : Comment le dynamisme est-il utile et comment puis-je l'utiliser ?

Le dynamisme a de nombreuses utilisations. Ici, nous nous concentrerons principalement sur le cas d'utilisation courant du polymorphisme de forme, qui consiste à créer une représentation de modèle exportée flexible, généralement utilisée pour représenter la taille de lot ou la longueur de séquence dynamiques.

Modèle statique add_one

Pour ce faire, nous allons utiliser le modèle add_one simple suivant :

def add_one(x):
  return x + 1

Lorsque nous effectuons le traçage à l'aide d'un tensor<4xf32>, nous obtenons le programme StableHLO suivant :

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

Ce modèle ne fonctionnera que pour les arguments d'entrée ayant une forme tensor<4xf32>. Si nous devions modifier la taille de notre lot ou la longueur de notre séquence, nous devrions retracer le code source et le réabaisser vers StableHLO. De plus, rien ne garantit que nous aurons toujours accès au code source.

Modèle dynamique add_one

C'est là que le dynamisme polymorphe des formes entre en jeu. À la place, JAX et PyTorch/XLA peuvent émettre le modèle add_one avec un IR dynamique valide qui diffusera la constante pour correspondre à la forme d'entrée dynamique comme suit :

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

Cette représentation de modèle est beaucoup plus flexible et permet de spécifier des valeurs telles que la taille du lot ou la longueur de la séquence de manière différée. Ce modèle peut être déployé sur des plates-formes compatibles avec les formes dynamiques (comme AI Edge) ou affiné à l'aide des passes de dynamisme mentionnées dans cette documentation.

Affiner le modèle dynamique

Par exemple, l'ordre de transmission suivant peut affiner complètement ce programme :

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

Voici comment le programme est transformé de manière incrémentielle :

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

Dynamisme dans StableHLO Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.