Questo documento descrive l'API di aliasing XLA, che consente di specificare l'aliasing tra il buffer di input e quello di output durante la creazione di un programma XLA.
Definizione degli alias in fase di compilazione
Ad esempio, considera un modulo HLO banale che aggiunge semplicemente 1 al suo input:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Questo modulo allocherà due buffer a 4 byte: uno per l'input %p e uno per l'output %out.
Tuttavia, spesso è opportuno eseguire l'aggiornamento in loco (ad esempio, se nel frontend che genera 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'alias 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 l'API XlaBuilder::SetUpAlias.
Definizione degli alias in fase di runtime
L'aliasing definito nel passaggio precedente viene specificato durante la compilazione.
Durante l'esecuzione, puoi utilizzare l'API LocalClient::RunAsync per scegliere se donare il buffer.
I buffer di input nel programma sono aggregati in ExecutionInput, che a loro volta contengono una struttura ad albero di MaybeOwningDeviceMemory. Se la memoria è specificata come proprietario (la proprietà del buffer viene passata al runtime XLA), il buffer viene effettivamente donato e l'aggiornamento viene eseguito, come richiesto dall'API di aliasing in fase di compilazione.
Tuttavia, se il buffer con alias in fase di compilazione non viene donato in fase di runtime, si attiva la protezione dalla copia: viene allocato un buffer di output O in più e i contenuti del buffer di input P destinato ad avere l'alias vengono copiati in O (in questo modo il programma può essere eseguito come se il buffer O fosse stato donato in fase di runtime).