Alias in XLA

Questo documento descrive l'API aliasing XLA, che consente di specificare l'aliasing tra i buffer di input e output quando si crea un programma XLA.

Definizione dell'aliasing al momento della compilazione

Ad esempio, considera un modulo HLO semplice che aggiunge semplicemente 1 al proprio input:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Questo modulo allocherà due buffer da 4 byte: uno per l'input %p e uno per l'output %out.

Tuttavia, spesso è consigliabile eseguire l'aggiornamento in loco (ad esempio, se nel frontend generando l'espressione la variabile di input non è più attiva. dopo il calcolo, come nell'incremento p++).

Per eseguire questo aggiornamento in modo efficiente, puoi specificare l'aliasing di input:

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Il formato specifica che l'intero output (contrassegnato da {}) ha un alias che rimanda al parametro di input 0.

Per specificare l'aliasing in modo programmatico, consulta XlaBuilder::SetUpAlias tramite Google Cloud CLI o tramite l'API Compute Engine.

Definizione dell'aliasing in fase di runtime

L'aliasing definito nel passaggio precedente viene specificato durante la compilazione. Durante l'esecuzione, puoi utilizzare LocalClient::RunAsync API per scegliere se donare il buffer.

I buffer di input del programma sono aggregati ExecutionInput che a loro volta contengono un albero di MaybeOwningDeviceMemory. Se la memoria è specificato come owning (la proprietà del buffer viene passata al runtime XLA), il buffer viene effettivamente donato e l'aggiornamento viene eseguito sul posto, come richiesto dall'API di aliasing in fase di compilazione.

Se, tuttavia, il buffer con alias in fase di compilazione non viene donato in runtime, si attiva la protezione dalla copia: è allocato un ulteriore buffer di output O, e i contenuti del buffer di input P da utilizzare come alias vengono copiati in O (quindi il programma può essere eseguito come se il buffer O fosse donato in fase di esecuzione).