[XLA:GPU] Enhance the tiled 0-2-1 transpose algorithm to handle fusion.
Add class TiledParameterInfo to provide information for FusedIrEmitter to read the content of a tiled parameter from the tile buffer instead of the original input memory. Reimplement the tiled 0-2-1 transpose algorithm for copy instructions only in a more general way so that it can handle both fusion instructions and copy instructions. The original tiled 0-2-1 transpose implementation incorrectly used (tile_size+1) rows for a tile buffer to reduce share memory bank conflicts while it should be (tile_size+1) column instead. This is a performance issue and is fixed in the new implementation. The original tiled 0-2-1 transpose implementation did not generate LLVM alias meta data for the loads and stores of the tensors. This was due to a bug where function IrArray::CastToShape miss copying meta data to the new IrArray object. This is also a performance issue and is fixed in this change. Modified KernelSupportLibrary to support emitting an if-stmt with a given branch name prefix. Add test cases to test the new implementation. PiperOrigin-RevId: 203310403
Loading
Please sign in to comment