[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
Loading
Please sign in to comment