Need `ttnn.index` or lowering of `aten.index.Tensor` performance too bad

[swimdi]

In beit, has this input varation of aten.index.Tensor

["Tensor<[732, 12]> self = ?", "List[Optional[Tensor]] indices = [<[38809]>]"]

It mean it will pick part of input 38809 times, and its output is [38809, 12]

The original run can pass

cd pytorch2.0_ttnn
pytest tests/models/beit/test_beit_image_classification.py
# pass

But if it apply the lowering method of #559 , it took very long time

git checkout swimdi/stage4-index
pytest tests/models/beit/test_beit_image_classification.py
# took very long time

And I think so the performance is unaceeptable, so cannot use the lowering method of #559 , TT should implement an operator named ttnn.index to completely fit this lowering

[swimdi]

Hi @ayerofieiev-tt :
I think maybe I can try to implement ttnn.index.
Here's example use of index and slice

• index(input, indices = [None, tensor([0, 1, 1]), tensor([2, 1, 2])]) result is
  [input[:][0][2], input[:][1][1], input[:][1][2]]
• ttnn.slice(input, begins=[0, 32, 0], ends=[1, 64, 16], step = [1, 2, 1]) result is
  input[0:1:1,32:64:2,0:16:1] # shape before: [1, 64, 16], after: [1, 16, 16]

I think there will have some similar tt-metal code between them
So maybe I can imitate slice code under ttnn/cpp/ttnn/operations/data_movement/slice to implement ttnn.index

[ayerofieiev-tt]

@swimdi , sure, lets try

[swimdi]

Hi @ayerofieiev-tt ,

I've been working on implementing ttnn.index, but I’m still in the process of studying it. Here are some notes of my study:

https://docs.google.com/document/d/1CpTNJhL7MeSY_hX5iFOEUNJNmBdQWt0d4U9G7ewPfmA/edit?tab=t.0

I’ve also tagged @jerrysky3 , @pschang-phy , and @jdh8 in case they have time to assist.

Just a notice, it might take me some time to complete the implementation. If this is urgent, it might be faster for the TT team to handle it.

[jdh8]

And I think so the performance is unaceeptable, so cannot use the lowering method of #559 , TT should implement an operator named ttnn.index to completely fit this lowering

I am with @swimdi. It is infeasible to compile an aten.index.Tensor instance to 38809 ops. I propose two ways to deal with this op.

For our year-end goals, we set up a threshold for aten.index.Tensor, e.g. N = 20. If lowering with #559 produces at most N ops, go for it. Otherwise, skip the conversion (letting it fall back).

For the ultimate goals, we still need ttnn.index in the kernel to fully support aten.index.Tensor.

[swimdi]

For our year-end goals, we set up a threshold for aten.index.Tensor, e.g. N = 20. If lowering with #559 produces at most N ops, go for it. Otherwise, skip the conversion (letting it fall back).

I try to merge #559 , but it encounter a problem
There's a usecase of index

indices = [None, tensor([0, 1, 1]), tensor([2, 1, 2])]
=> output is [input[:][0][2], input[:][1][1], input[:][1][2]]

So I need to add a getitem op like this

getitem(input, (slice(None), 0, 2)) # equal to input[:][0][2]

But I cannot add getitem node with slice(None) argument to graph module because it say slice(None) is unhasable

So currently #559 cannot handle indices with None, and then it become useless to increase the model coverage, so I'm not merge it

And also @ayerofieiev-tt is not suggest to merge mentioned at the meeting, so I think it can just implement on tt-metal