Questo documento descrive l'API di aliasing XLA, che consente di specificare l'aliasing tra i buffer di input e 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 alloca due buffer a 4 byte: uno per l'input %p e uno per l'output %out.
Tuttavia, spesso è preferibile 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'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 al
parametro di input 0.
Per specificare l'alias 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 MaybeOwningDeviceMemory. Se la memoria viene specificata come proprietario (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.
Tuttavia, se, tuttavia, il buffer con alias al momento della compilazione non viene donato in fase di runtime, si attiva copy-protection: viene allocato un buffer di output aggiuntivo O e i contenuti del buffer di input P destinato ad avere alias vengono copiati in O (quindi il programma può essere eseguito come se il buffer O fosse stato donato in fase di runtime).