Commit f7e4143a authored by Bixia Zheng's avatar Bixia Zheng Committed by TensorFlower Gardener
Browse files

[XLA:GPU] Remove the implicit reference of the last kernel in the thunk sequence

when ParallelLoopEmitter is used.

Currently, when IrEmitterUnnested handles an HLO for which ParallelLoopEmitter
is used to emit the kernel, we first create the thunk for the kernel, push the
thunk to the end of the thunk sequence and then call the handler for the same
HLO in IrEmitter. The handler in IrEmitter emits code to compute the base
pointers for parameters before constructing the elemental ir emitter, and then
passes the elemental ir emitter to IrEmitterUnnested::EmitTargetElementLoop,
which implicitly references the last kernel in the thunk and eventually invokes
ParallelLoopEmitter with the elemental ir emitter.

To remove such an implicit reference of the last kernel in the thunk sequence,
we postpone the creation of the thunk to routine
IrEmitterUnnested::EmitTargetElementLoop, and delay the emission of code that
computes the base pointers for parameters by wrapping the code emission inside
a generator and pass the generator to IrEmitterUnnested::EmitTargetElementLoop
through the FusedIrEmitter.
PiperOrigin-RevId: 219249592
parent 10f25986
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment