1.12.0-rc2 cherry-pick request: Various XLA scatter improvements. (#23235)
* [XLA] Update Tf2Xla bridge to use Scatter HLO. PiperOrigin-RevId: 215687800 * [XLA:GPU] Add an implementation of scatter for GPU This simple has a kernel that runs on every element of the updates tensor, figure out the right indices to perform the update, and applies it with an atomic operation. Currently we emit a CAS for plain (i.e. non-add) updates, which is inefficient. Also TuplePointsToAnalysis doesn't know that it should alias the operand and output buffers of a scatter, which would avoid a copy. PiperOrigin-RevId: 216412467 * [XLA] Allow scatter to share the operand buffer with the output This avoids a copy. PiperOrigin-RevId: 216437329 * [XLA:GPU] Elide the SequentialThunk when emitting scatter with no copy We have a 1-element thunk sequence if we're not copying. That's still two thunks and hlo profiling gets confused if it sees two thunks for the same instruction and one of them claims to be the whole instruction. PiperOrigin-RevId: 216448063 * [XLA:GPU] Allow input fusion into scatter We fuse everything into the scatter now, and emit two kernels. The first kernel fills the output buffer with the computation fused into the scatter operand. The second kernel is a regular scatter, which also contains the fused operations from the updates and scatter_indices inputs. PiperOrigin-RevId: 216624225 * [XLA:GPU] Adding a test case for Scatter where GPU implementation fails. PiperOrigin-RevId: 216798034 * [XLA:GPU] Fix scatter oob check computation This was comparing the index after adding it to the window, and then comparing against the window dimension. This means that the bounds check was only correct for the first element of a window. Instead compare the scatter index, which is the same for all elements of a window. PiperOrigin-RevId: 216921512 * [XLA:GPU] Elide tuple roots of the entry computation The tuple buffer is never read, so stop emitting code to fill it. A typical root tuple consists of a H2D memcpy and a host callback, both of which are somewhat slow. This helps tiny models and inference benchmarks, where the host/device syncs can be a significant part of the runtime of the entire computation. PiperOrigin-RevId: 216968475
Loading
Please sign in to comment