Commit 4009f82f authored by Sanjoy Das's avatar Sanjoy Das Committed by TensorFlower Gardener
Browse files

Implement constant buffer allocation for XLA:GPU

This CL teaches XLA:GPU to use "normal" buffer assignment for constant
instructions.  Constant instructions are mapped to a BufferAllocation, like all
other instructions, except the storage for this buffer is allocated statically
as a global in the generated PTX.

This CL does not change how we access the constants -- in
IrEmitterUnnested::BuildKernelThunk (used for top level computations) and in
HloToIrBindings::EmitBasePointersForHlos (used for nested computations) we bind
the kConstant instructions to the llvm::GlobalVariable backing them.  So users
of constant instructions still access the globals corresponding to the constants
directly.

However, we no longer emit the constant literals inline.  Instead we emit a
constant with a zero initializer and then memcpy in the contents of the literal
when we load the CUBIN/PTX.  This works around compile time issues in LLVM and
ptxas caused by large constants.

We also populate `BufferAllocations` with the device pointers for the constant
globals.  This is at least needed for TupleThunk today because TupleThunk wants
the addresses for the sub-buffers on the host.  I'm not sure if there are other
places in XLA:GPU that rely on there being an entry in BufferAllocations for
every BufferAllocation.

PiperOrigin-RevId: 206243319
parent e336ee65
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment