Cette page a été traduite par l'API Cloud Translation.
Switch to English

XLA appels personnalisés

Ce document décrit comment écrire et utiliser XLA « appels personnalisés ». appels personnalisés vous invoquons code écrit dans un langage de programmation comme C ++ ou CUDA d'un programme XLA.

Custom-appel CPU

Vous pouvez créer une instruction HLO qui représente un appel personnalisé via l'API client de XLA. Ce n'est pas exposé via tensorflow comme l'écriture.

Par exemple, le code suivant utilise un appel personnalisé pour calculer A[i] = B[i % 128] + C[i] sur la CPU. (Bien sûr, vous pouvez - et devrait - faire avec HLO régulièrement.)

 #include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/service/custom_call_target_registry.h"

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(0, xla::ShapeUtil::CreateShape(F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(1, xla::ShapeUtil::CreateShape(F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}));
}

void do_custom_call(void* out, const void** in) {
  float* out_buf = reinterpret_cast<float*>(out);
  const float* in0 = reinterpret_cast<const float*>(in[0]);
  const float* in1 = reinterpret_cast<const float*>(in[1]);
  for (int i = 0; i < 2048; ++i) {
    out_buf[i] = in0[i % 128] + in1[i];
  }
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "Host");
 

Notez que la fonction do_custom_call a besoin de connaître les dimensions des tampons qu'il opère sur. Dans cet exemple, nous hardcode la taille 128 et 2048. Si vous ne voulez pas faire cela, vous pouvez passer les dimensions en tant que paramètres à l'appel.

Custom-appel sur GPU

La coutume cadre d'appel GPU est un peu différent de celui de la CPU. Voici un exemple CUDA qui fait la même A[i] = B[i % 128] + C[i] calcul que le code CPU ci - dessus.

 void do_it() { /* same implementation as above */ }

__global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
  size_t idx = threadIdx.x * blockSize.x + gridIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, void** buffers,
                    const char* opaque, size_t opaque_len) {
  const float* in0 = reinterpret_cast<const float*>(buffers[0]);
  const float* in1 = reinterpret_cast<const float*>(buffers[1]);
  float* out = reinterpret_cast<float*>(buffers[2]);

  const int64 block_dim = 64;
  const int64 grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim,
                       /*dynamic_shared_mem_bytes=*/0, stream>>>(in0, in1, out);
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "CUDA");
 

Notez d' abord que la fonction d'appel personnalisé GPU est toujours une fonction exécutée sur le CPU. Notre do_custom_call fonction CPU est responsable du travail enqueueing sur le GPU. Ici, il lance un noyau CUDA, mais il pourrait aussi faire quelque chose d'autre, comme cublas d'appel.

des buffers est un tableau de pointeurs qui vit sur l'hôte, et chaque élément qu'il contient des points de dispositif ( par exemple GPU) de mémoire. Les paramètres viennent en premier, suivi de la valeur de sortie. Ceci est notamment différent de la convention d' appel du processeur, qui a deux params, ins et out . La principale raison pour laquelle nous diverger est de rendre possible de gérer les entrées tuple en forme / sorties de manière efficace; voir la section ci-dessous.

Comme dans l'exemple du processeur, nous avons Hardcoded l'entrée et la sortie tampon tailles dans notre appel personnalisé. Toutefois, contrairement à dans le cas du processeur, la mémoire tampon passant tailles comme opérandes à l'appel personnalisé ne fonctionne pas bien. En général, nous avons besoin de la mémoire tampon tailles disponibles pour nous sur la CPU; par exemple, lors du lancement d'un noyau, nous avons besoin de connaître les dimensions bloc / grille à utiliser. Mais si nous devions passer les tailles de tampon comme opérandes à notre appel personnalisé, leurs valeurs vivrions dans la mémoire du GPU. Nous aurions alors à faire un dispositif à hôte memcpy synchrone cher au début de notre opération juste pour lire les tailles.

Pour vous permettre de travailler autour de cela, nous fournissons l' opaque paramètre. Vous pouvez définir ce à une chaîne arbitraire d'octets lorsque vous créez l'appel personnalisé:

 std::string opaque = "...";
xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/ShapeUtil::CreateShape(F32, {2048}),
                opaque);
 

Depuis xla::Shape a une représentation tampon de protocole, vous pouvez stocker ceci à l' intérieur sérialisé proto d' opaque et désérialiser au sein de votre commande appel GPU. A noter cependant que , bien que xla::ShapeProto ne change pas souvent, cela ne change. Vérifiez le journal git pour voir comment il a changé dans le passé.

En passant tuples aux appels-personnalisés

Pensez à l'appel personnalisé suivant.

 using xla::ShapeUtil;
Shape p0_shape = ShapeUtil::MakeTuple({
    ShapeUtil::MakeShape(F32, {32}),
    ShapeUtil::MakeTuple({
        ShapeUtil::MakeShape(F32, {64}),
        ShapeUtil::MakeShape(F32, {128}),
    }),
    ShapeUtil::MakeShape(F32, {256}),
});
xla::XlaOp p0 = xla::Parameter(0, p0_shape, "p0");

Shape out_shape = ShapeUtil::MakeTuple({
  ShapeUtil::MakeShape(F32, {512}),
  ShapeUtil::MakeShape(F32, {1024}),
});
xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape);
 

Sur les deux CPU et le GPU, un tuple est représenté dans la mémoire comme un tableau de pointeurs. En C ++ - pseudocode, paramètre 0 ci-dessus est posé comme suit.

 // In-memory layout of parameter 0 from custom-call above.  True on both CPU
// and GPU.
float* subbuf0 = new float[32];
float* subbuf1 = new float[64];
float* subbuf2 = new float[128]
float* subbuf3 = new float[256];

void* subtuple = new void*[2];
(*subtuple)[0] = subbuf1;
(*subtuple)[1] = subbuf2;

void* p0 = new void*[3];
(*p0)[0] = subbuf0;
(*p0)[1] = subtuple;
(*p0)[2] = subbuf3;
 

Bien que la représentation en mémoire de tuples est le même processeur et GPU, ils sont traités différemment dans le CPU et le GPU les conventions d'appel d'appel personnalisé.

Tuple délivre en sortie en tant que tampons temporaires

entrées tuple à des appels personnalisés sont une commodité, mais ils ne sont pas strictement nécessaires. Si nous ne soutenons pas les entrées tuple aux appels personnalisés, vous pouvez toujours déballer les tuples utilisant get-tuple-élément avant de les transmettre à l'appel personnalisé.

D'autre part, les sorties tuple ne vous permettent de faire des choses que vous ne pouvait pas autrement.

La raison évidente d'avoir des sorties tuple est, voilà comment un appel personnalisé (ou tout autre op XLA) renvoie plusieurs tableaux indépendants.

Mais moins évidemment, une sortie de tuple est aussi une façon de donner à votre mémoire appel temp personnalisé. Oui, une sortie peut représenter un tampon temporaire. Considérons, une mémoire tampon de sortie a la propriété que l'op peut écrire, et il peut lire après qu'il a été écrit. C'est exactement ce que vous voulez à partir d'un tampon temporaire.

Dans l'exemple ci - dessus, supposons que nous voulions utiliser le F32[1024] comme un tampon temporaire. Ensuite, nous écrivions le HLO comme ci-dessus, et nous avions tout simplement jamais lu index de tuple 1 de la sortie de l'appel personnalisé.

Tuples dans les appels personnalisés CPU-

Dans le code du processeur, nous avons une fonction do_custom_call(const void** ins, void* out) . ins est un tableau avec un seul élément, qui pointe vers param0 . Les subbuffers de param0 sont accessibles par déréférencement ce pointeur, et les subbuffers de output_tuple sont accessibles par déréférencement out .

Tuples dans les appels personnalisés-GPU

Dans le code GPU, nous avons une fonction do_custom_call(..., void** buffers, ...) . Dans ce cas , des buffers est un réseau hôte de six pointeurs de périphériques, un pour chaque mémoire tampon de la feuille dans l'entrée / sortie. Pour générer la liste plate, nous itérer sur les paramètres et la sortie, et pour chacun nous faire une précommande traversal de sa forme. concrètement:

 // Layout of `buffers` parameter to GPU custom call function for custom-call
// above.
buffers[0] == subbuf0
buffers[1] == subbuf1
buffers[2] == subbuf2
buffers[3] == subbuf3
buffers[4] == output_subbuf0
buffers[5] == output_subbuf1