Add support for _unsafe_index.

[ysiraichi]

This PR adds support for _unsafe_index by aliasing index.

Problem: _unsafe_index was hitting the CPU fallback execution flow, which moved the tensors to CPU (including self)

• The CPU fallback was calling the dispatcher again: at::index
• That, was getting dispatched to the XLA implementation (error, since self lives on CPU)

This problem affected a few Torchbench benchmarks running with openxla backend: Super_SloMo, doctr_det_predictor, pytorch_unet, ...

[JackCaoG]

do we know when will this method being used(how does it get dispatched to this) and create a test for it?

[wonjoolee95]

+1 to adding the tests. For more info, the cpp unit tests can verify that this new op you lowered is properly invoked via metric assertions as such https://github.com/pytorch/xla/blob/master/test/cpp/test_aten_xla_tensor_1.cpp#L35.

[ysiraichi]

I get it as a node in the FX graph generated from AOTAutograd (using openxla backend). I will investigate it further and add tests.

[lezcano]

@peterbell10 added this function. These functions are an artefact to signal the compiler to avoid performing any boundary checks on the indices that they use. These are used in decompositions where we know that the indices will be in the right range, like the backward of functions, or a decomposition where we know that the inputs are bounded to the right range.

[lezcano]

Note that you also have its cousin _unsafe_index_put

[wonjoolee95]

Thanks for adding the cpp tests! For some reason, I can see that the CI fails at build stage but I can't seem to see any logs. I've just re-triggered the CI, let's see if this run is good.

[JackCaoG]

Thanks!

[peterbell10]

That, was getting dispatched to the XLA implementation (error, since self lives on CPU)

This sounds like a bug, either the fallback creates a CPU tensor with XLA dispatch keys, or the falback mechanism doesn't handle lists of tensors and some of the arguments were not moved to CPU.