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

Overhaul XLA:CPU's calling convention.

This CL introduces a clean separation between calls to "thread local" and
"global" computations in XLA:CPU.

Global computations are:

 - kWhile body and condition computations
 - kConditional true and false computations
 - kCall callees

Parameters and results buffers for these calls are assigned a static
BufferAllocation::Slice by buffer assignment and so they don't require pointers
to result buffers and parameters to be explicitly passed in.  In fact, passing
in result and parameters buffers is actively misleading because in cases like:

  while_condition {
    val = (s32[], pred[]) infeed()
    ROOT result = get-tuple-element(val), index=0
  }

there is no instruction explicitly copying the result of the computation into
the result buffer.  Instead, it is up to the caller to pick up the correct
result buffer by asking buffer assignment (which would be buffer where infeed
wrote its second tuple component).

Thread local computations are all the other nested computations except fusion,
e.g. computations used by kMap and kReduce.

Parameters and result buffers for these calls are assigned a "thread local"
BufferAllocation::Slice which in XLA:CPU are mapped to allocas.  Since these are
not static addresses, we *do* need to pass in parameter and result buffers.  The
output is written to the result buffer by "allocating" the storage for the
root into the result buffer passed in by the caller.

There are two cleanup items that I kept off this CL to make reviews easier:

 - We should rename "temps" to something more generic, like "buffer_table".
   I'll do that in a followup CL.

 - We should use GatherComputationsByAllocationType from buffer_assignment.cc to
   CHECK that we use thread local calls for thread local callees and global
   calls for global callees.

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