Unverified Commit e72c9ebe authored by Todd Wang's avatar Todd Wang Committed by GitHub
Browse files

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
parent da1b48dd
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment