Add support for aten::index operator

[ftxj]

There are many Tensor indexing API in PyTorch. We already support the index_select operator, but we doesn't support index operator, which is supported by Inductor + Triton and widely used in GNN. So we need to add this operator.

Although we can share most of logic with index_select, there are some difficulties. The index signature is aten::index.Tensor(Tensor self, Tensor?[] indices) which use List data type. But nvfuser doesn't support it. So the TorchScript test will raise an error until we support List[T].

• Support List[T] type in nvfuser.

[csarofeen]

What's the status of this PR? CC @jjsjann123

[jjsjann123]

Do we actually support index with backend?!?! I remember we went over the idea on GNN meeting and it's the actual codegen part where we are lacking. (frontend list of tensor as inputs is also needed of course.)

asking @naoyam for codegen support.

[naoyam]

Generally no, we don't support aten::index. However, I think what @ftxj is trying to do is to add a limited partial support of aten::index only when it can be converted to index_select, which is supported with limitations. It seems that's common in GNNs, so I guess it probably makes sense.

[jjsjann123]

That sounds good. I think the example included is also indicating a single tensor indexing.

Looks like we are using python API here. Are we doing this via dynamo/primtorch or TorchScript? We can plumb it through python API, but frankly speaking I'm less comfortable having a "limited support" thing there, which makes it hard to get a robust runtime.

So for dynamo/primtorch, we can probably hack it via a decomposition to convert compatible index to index_select. Need to be careful with not going into an infinite recursion with the dispatch.
For TorchScript, I think we can hack it in parser.cpp. We can reject index op when len(indices) > 1 and parser index as index_select.

[naoyam]

We would likely need to do this anyway as supporting aten::index completely would be quite challenging. I'm also working on limited support of slice, which is also part of aten::index.

[jjsjann123]

Even with the suggested changes, we still need to somehow wire Tensor[] in our integration.

We can do the proper thing and plumb it through, which might take some time.
Alternatively, we can fake it by converting list of tensors to a single tensor and work with that... How urgent do we need this working? @ftxj

[jjsjann123]

awesome. Looks like we are already parsing it as index_select.

[ftxj]

But I didn't make restrictions on the situation that cannot be converted...

[jjsjann123]

index here is a list of tensors. so we need to access the first and only item instead of the whole thing.

[jjsjann123]

should return a gatherOp type here.

[jjsjann123]

we should also return false when second input length is > 1.

[jjsjann123]

This signature doesn't look right. It should be a list of Tensors instead.

[ftxj]

This signature isn't right. I don't know how to declare the list of Tensor type...

[jjsjann123]

This really isn't the expected index signature?

[ftxj]

Even with the suggested changes, we still need to somehow wire Tensor[] in our integration.

We can do the proper thing and plumb it through, which might take some time. Alternatively, we can fake it by converting list of tensors to a single tensor and work with that... How urgent do we need this working? @ftxj

Support aten::index is not urgent at present, I can replace the aten::index using aten::index_select in the python level.