Gather to Slice Fusion

[Lafi7e]

This PR is to optimize the running for below code from Huggingface's XLNet model.

x = torch.index_select(x, 3, torch.arange(klen, device=x.device, dtype=torch.long))

The code will be exported to Range->Gather, which can be fused to a Slice Op. Slice kernel is much faster than Gather, especially for backward run. The main reason is for Gather, the data in indices can be duplicated so that it needs sum during backward, but Slice node cannot have such case.

Use Huggingface's XLNet model for profiling.

• Before the fuse
  forward, ~753us
  
  backward, ~46101us
  

• After the fuse
  forward, ~627us
  
  backward, ~677us
  

[baijumeswani]

Can you add a note here explaining that this fusion is primarily helpful for gradient computation for context.

[Lafi7e]

Forward is also faster, but not that big. The model code also has below comment to show this for PyTorch. That's why I also put this fusion to the transformer utils file for inference.

        # Note: the tensor-slice form was faster in my testing than torch.index_select
        #       However, tracing doesn't like the nature of the slice, and if klen changes
        #       during the run then it'll fail, whereas index_select will be fine.
        x = torch.index_select(x, 3, torch.arange(klen, device=x.device, dtype=torch.long))

[baijumeswani]

Just a small comment to add some context. You can get it in another PR if you like.

Looks good.