[TE] Support negative indices #9023

AndrewZhaoLuo · 2021-09-15T20:28:19Z

Negative indices are a pretty common feature in most modern programming languages and libraries. This proposed PR would add optional support for negative indices out of the box for tensors in TE.

A lot of operators from other frontends support negative indices. See ONNX, numpy, and PyTorch where their tensors support negative indices out of the box. Unfortunately TE and the rest of TVM does not.

One benefit of this change is now it is trivial to change operators to support negative indices. This makes adapting frontends to our operators simpler. For example an operator using this simply has to turn the support_negative_indices flag to true for every relevant indexing operation. An example of doing this can be found in https://github.com/AndrewZhaoLuo/tvm/blob/02f1870d7f2e274cfd7e04678691da019ff201f0/include/tvm/topi/transform.h#L1265. Right now how we handle negative indices is by manually converting them to positive indices which can be cumbersome and results in lots of code duplication.

The downside is that it technically adds a little more computation per indexing operation. While these operations are probably memory and not compute bound I am not 100% confident that setting support_negative_indices=true would not result in performance regressions. Furthermore, it will make some of printouts very ugly since it will but the comparison operation in each index statement.

@jroesch also has great points that this might make loop level analysis hard and result in some optimizations harder.

Pros:

Clear way to support negative indices

Cons:

Printouts will be ugly
Might be slower (but can turn off easily)

Unknowns:

Dynamic shapes?
Effect on compile time optimizations

masahi · 2021-09-15T20:56:21Z

For the context, this PR was split from another PR following the discussion #8971 (comment), since this could be potentially a controversial change (change the semantics of te::Tensor indexing)

cc @tqchen @junrushao1994

mbrookhart · 2021-09-15T22:19:40Z

I like this idea. I've been slowly throwing relay.where around the ops and the importers to solve this problem when I hit it, but that adds some complication to ops, fusion, and the risk of making things slower than they need to be. This would enable things to be much simpler at the Relay level in a number of places, I'm very curious to hear Tianqi and Junru's thoughts.

tqchen · 2021-09-15T23:26:32Z

src/te/tensor.cc

+
+  if (support_negative_indices) {
+    for (size_t i = 0; i < shape.size(); i++) {
+      PrimExpr new_index = if_then_else(indices[i] < make_const(indices[i]->dtype, 0),


you want to use select here, because it does not involves a memory operation(thus can evaluate both sides)

Done.

I understand the difference between if_then_else and Select

Don't fully understand why we want to be able to evaluate both sides however. Is this to prevent branching?

tqchen · 2021-09-15T23:29:48Z

I think this can be useful for some of the indexing operations. We could try to make it more explicit though. e.g. introduce a new API. tensor.LookupWithNegativeIndices(indices) and explicitly call into them in these cases

yzh119 · 2021-09-15T23:44:03Z

Is the negative indices used inside TE? If so I wonder can we make it a transformation pass rather than runtime behavior?

tqchen · 2021-09-15T23:56:33Z

@yzh119 I believe the intended usage was for the case where the value of indices are not known at compile time, otherwise compiler will be able to prove and simplify the conditionals

junrushao · 2021-09-16T01:01:07Z

I think I sort of understand the usecase here: some negative indices are not known to be negative until in runtime. This forces us to defer the conversion from compile-time to runtime.

On the other hand, I am not 100% sure if it is the best fix by adding a new argument in the public interface, given that in most cases indices are just positive and well in-range.

I was thinking, if the issue comes from an importer, is it possible to add an operator like normalize_indices, mark it as injective which makes it fusible, so that there isn't architectural change in TE? What do you guys think?

AndrewZhaoLuo · 2021-09-16T17:55:30Z

I think I sort of understand the usecase here: some negative indices are not known to be negative until in runtime. This forces us to defer the conversion from compile-time to runtime.

On the other hand, I am not 100% sure if it is the best fix by adding a new argument in the public interface, given that in most cases indices are just positive and well in-range.

I was thinking, if the issue comes from an importer, is it possible to add an operator like normalize_indices, mark it as injective which makes it fusible, so that there isn't architectural change in TE? What do you guys think?

We do have a normalize_indices on the relay level. I think a problem though is the format might be slightly different so you would need a few different implementations. For example take tf's gather vs gather_nd.

AndrewZhaoLuo · 2021-09-20T17:49:14Z

I'll get ahead with this and start responding to comments if there are no other objections.

AndrewZhaoLuo · 2021-09-20T22:26:18Z

This is ready for another look though I might need some guidance for unit tests

include/tvm/te/tensor.h

junrushao

It looks good to me given we want to fully support numpy-style negative indexing. On the API side, although I still think we can put this as a helper function in topi instead of a member function of te::Tensor, the difference is minimal because we are not making architectural change here. So nice work!

BTW, shall we add a unittest for this?

AndrewZhaoLuo · 2021-09-22T18:36:32Z

It looks good to me given we want to fully support numpy-style negative indexing. On the API side, although I still think we can put this as a helper function in topi instead of a member function of te::Tensor, the difference is minimal because we are not making architectural change here. So nice work!

BTW, shall we add a unittest for this?

I'll look into the changes you suggest and add a unit test.

junrushao · 2021-09-22T18:45:14Z

Ah no worries! I’m not suggesting changes :-) just unittests

AndrewZhaoLuo · 2021-10-01T06:12:05Z

@junrushao1994 looked at the unittests for cpp and they seem a little bare. I added a single simple test which logs the following results:

A is of shape (x, y)

A(0, 0):
[23:09:53] /home/andrewzhaoluo/Desktop/dev/tvm/tests/cpp/tensor_test.cc:59: A[0, 0]

A.IndexWithNegativeIndices(-1, -1):
[23:09:53] /home/andrewzhaoluo/Desktop/dev/tvm/tests/cpp/tensor_test.cc:60: A[select((bool)1, (-1 + x), -1), select((bool)1, (-1 + y), -1)]

A.IndexWithNegativeIndices(0, -1):
[23:09:53] /home/andrewzhaoluo/Desktop/dev/tvm/tests/cpp/tensor_test.cc:61: A[select((bool)0, x, 0), select((bool)1, (-1 + y), -1)]

Some guidance on what is sufficient would be helpful.

masahi · 2022-01-09T23:20:23Z

What's the status of this PR?

AndrewZhaoLuo · 2022-01-10T18:04:17Z

Oh, forgot about this as I moved to working on other things. We still cool with merging @junrushao1994 ?

junrushao · 2022-01-10T18:05:20Z

@AndrewZhaoLuo yeah let's get it merged!

AndrewZhaoLuo · 2022-01-12T18:23:41Z

@junrushao1994 this is now ready to merge

junrushao

LGTM

junrushao · 2022-01-12T20:57:08Z

Thanks @AndrewZhaoLuo! It's merged :-)

* initial change * more explicit api * switch to select * add support for negative indices * reduce things further * lint * to CamelCase * unit test Co-authored-by: Andrew Zhao Luo <andrewzhaoluo@system76-pc.localdomain>

AndrewZhaoLuo requested review from junrushao, kparzysz-quic, masahi, tqchen, vinx13 and ZihengJiang as code owners September 15, 2021 20:28

AndrewZhaoLuo mentioned this pull request Sep 15, 2021

[Onnx] Fix NLL Loss tests #8971

Merged

tqchen reviewed Sep 15, 2021

View reviewed changes

AndrewZhaoLuo force-pushed the aluo/support-negative-indices branch from 903a7bc to 84e5ab9 Compare September 21, 2021 16:54

tqchen reviewed Sep 21, 2021

View reviewed changes

include/tvm/te/tensor.h Outdated Show resolved Hide resolved

junrushao reviewed Sep 22, 2021

View reviewed changes

AndrewZhaoLuo requested review from areusch, comaniac, jroesch, merrymercy and yzhliu as code owners October 1, 2021 06:02

Andrew Zhao Luo and others added 8 commits January 10, 2022 10:04

initial change

fcc1c83

more explicit api

3d77b0c

switch to select

788d1f0

add support for negative indices

2acbcf5

reduce things further

3bef867

lint

5f75a02

to CamelCase

72ba914

unit test

9b7d01a

AndrewZhaoLuo force-pushed the aluo/support-negative-indices branch from e19fdba to 9b7d01a Compare January 10, 2022 18:04

AndrewZhaoLuo requested a review from icemelon as a code owner January 10, 2022 18:04

junrushao approved these changes Jan 12, 2022

View reviewed changes

junrushao merged commit 80f9b9f into apache:main Jan 12, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TE] Support negative indices #9023

[TE] Support negative indices #9023

AndrewZhaoLuo commented Sep 15, 2021 •

edited

Loading

masahi commented Sep 15, 2021 •

edited

Loading

mbrookhart commented Sep 15, 2021

tqchen Sep 15, 2021

AndrewZhaoLuo Sep 20, 2021

tqchen commented Sep 15, 2021 •

edited

Loading

yzh119 commented Sep 15, 2021

tqchen commented Sep 15, 2021

junrushao commented Sep 16, 2021 •

edited

Loading

AndrewZhaoLuo commented Sep 16, 2021

AndrewZhaoLuo commented Sep 20, 2021

AndrewZhaoLuo commented Sep 20, 2021

junrushao left a comment

AndrewZhaoLuo commented Sep 22, 2021

junrushao commented Sep 22, 2021

AndrewZhaoLuo commented Oct 1, 2021

masahi commented Jan 9, 2022

AndrewZhaoLuo commented Jan 10, 2022

junrushao commented Jan 10, 2022

AndrewZhaoLuo commented Jan 12, 2022

junrushao left a comment

junrushao commented Jan 12, 2022

[TE] Support negative indices #9023

[TE] Support negative indices #9023

Conversation

AndrewZhaoLuo commented Sep 15, 2021 • edited Loading

masahi commented Sep 15, 2021 • edited Loading

mbrookhart commented Sep 15, 2021

tqchen Sep 15, 2021

Choose a reason for hiding this comment

AndrewZhaoLuo Sep 20, 2021

Choose a reason for hiding this comment

tqchen commented Sep 15, 2021 • edited Loading

yzh119 commented Sep 15, 2021

tqchen commented Sep 15, 2021

junrushao commented Sep 16, 2021 • edited Loading

AndrewZhaoLuo commented Sep 16, 2021

AndrewZhaoLuo commented Sep 20, 2021

AndrewZhaoLuo commented Sep 20, 2021

junrushao left a comment

Choose a reason for hiding this comment

AndrewZhaoLuo commented Sep 22, 2021

junrushao commented Sep 22, 2021

AndrewZhaoLuo commented Oct 1, 2021

masahi commented Jan 9, 2022

AndrewZhaoLuo commented Jan 10, 2022

junrushao commented Jan 10, 2022

AndrewZhaoLuo commented Jan 12, 2022

junrushao left a comment

Choose a reason for hiding this comment

junrushao commented Jan 12, 2022

AndrewZhaoLuo commented Sep 15, 2021 •

edited

Loading

masahi commented Sep 15, 2021 •

edited

Loading

tqchen commented Sep 15, 2021 •

edited

Loading

junrushao commented Sep 16, 2021 •

edited

Loading