Commit 9a7a63af authored by Justin Lebar's avatar Justin Lebar Committed by TensorFlower Gardener
Browse files

[XLA:GPU] Assume that tuple sub-buffers are available at runtime.

Previously we assumed this was not the case, and allowed front-ends to
pass in a pointer to tuple without also passing in pointers to
sub-buffers.

This mostly worked: Whenever we wanted a tuple sub-buffer, we'd just
chase the tuple's pointers in our emitted kernel.

But this doesn't work if we ever need a pointer to that sub-buffer on
the host.  Which we do if e.g. the sub-buffer is an input to a cudnn
call.

There are various ways to make this work, but by far the simplest and
most efficient is simply to specify away this problem, and say that the
front-end *must* give us all the pointers we want.  This is what the
earlier change, "Assert that all buffers and sub-buffers passed to XLA
have an explicit pointer" did.

This change adds a testcase and lets us skip some pointer chasing when
we have a tuple whose sub-buffers are known statically.

PiperOrigin-RevId: 190949743
parent fb8c7b1d
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment