[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
Loading
Please sign in to comment