[XLA:GPU] Remove a map data structure used to support the code generation of
GetTupleElement inside a fused computation. The map initially records the fusion parameters and their base addresses. Then during the forward traversal of the fused computation that produces elemental ir generators, we generate code to calculate the base address for the GetTupleElement instructions that produce tuple results, and add the information to the map. As a result, this map records all the tuple producing instructions whose results can be used in GetTupleElement instructions. During code generation, if a GetTupleElement instruction produces a non-tuple result, we look up the map for the base pointer of the tuple to generate code that produces the instruction result. On the other hand, if a GetTupleELement instruction produces a tuple result, we do not generate any code as the code for such instructions is already generated when we produce the map data structure. Such a map data structure is not necessary. The base addresses for a parameter can be derived from the IrArray of the parameter, which can be looked up through the parameter number. We change the elemental ir emitter for GetTupleElement instructions to return the base address for the tuple producing GetTupleElement instructions to avoid looking up the map for such instructions. In doing that, we also avoid generating code for tuple producing GetTupleElement instructions duing the phase that produces elemental ir generators to clean up the code generation for GetTupleElement instructions. PiperOrigin-RevId: 218922219
Loading
Please sign in to comment