[TOPI][Relay][OP] Add a strided_set operation. #4303
Conversation
abergeron
force-pushed
the
strided_set
branch
from
06074cd
to
721086d
Compare
|
I did the necessary rebase. |
|
is it used in dl frameworks? |
|
cc @jroesch @junrushao1994 can you also take a look? |
python/tvm/relay/op/_transform.py
Outdated
| ls = len(strides) | ||
| if ls < n: | ||
| strides = list(strides) | ||
| strides[ls:] = [1] * (n - ls) |
There was a problem hiding this comment.
You can make it slightly shorter.
strides = list(strides) + [1] * (n - len(strides))
python/tvm/relay/op/_transform.py
Outdated
| if lb < n: | ||
| begin = list(begin) | ||
| for i in range(lb, n): | ||
| begin.append(0 if strides[i] >= 0 else inputs.shape[i]) |
There was a problem hiding this comment.
This is a bit different because the value to append might be different at each index. It depends on the stride for that index. I'm not sure how to shorten that.
python/tvm/relay/op/_transform.py
Outdated
| if le < n: | ||
| end = list(end) | ||
| for i in range(le, n): | ||
| lim = inputs[0].shape[i] + 1 | ||
| end.append(lim if strides[i] >= 0 else -lim) |
There was a problem hiding this comment.
Same as for begin.
|
@yzhliu I think this is something like |
| # Convert negative indexes | ||
| for i in range(n): | ||
| begin[i] = tvm.if_then_else(begin[i] < 0, | ||
| begin[i] + a.shape[i], |
There was a problem hiding this comment.
Even begin[i] + a.shape[i] could trigger OOB. I am not sure how to assert the bound...Could anyone help here?
There was a problem hiding this comment.
By OOB do you mean out of bounds of the indexed array? If yes, that is not a problem because the code will never try to fetch indices that are out the array shape.
This code just tries to handle numpy-style negative indexing (starts from the end of the array) just like strided_slice does.
There was a problem hiding this comment.
Yep, I don’t have much idea about this either
There was a problem hiding this comment.
Even if the logical result is very similar to strided_slice the implementation has a difference.
In this operation: a[begin:end:stride] = b
The core kernel loops over all valid indexes for a and check if that index is part of the values selected by the combination of begin, end and stride. If it is, it will compute the corresponding index in b and map the output to that value. Otherwise it will pick up the value from a at that index.
In all cases it doesn't matter if begin, end, or stride doesn't fall within the bounds of a because they are never used to directly or indirectly index into a.
| begin[i] + a.shape[i], | ||
| begin[i]) | ||
| end[i] = tvm.if_then_else(end[i] < 0, | ||
| end[i] + a.shape[i], |
There was a problem hiding this comment.
Same here. OOB is not completely detected.
topi/python/topi/transform.py
Outdated
| within_index(begin[i], end[i], strides[i], indices[i])) | ||
| index_tuple.append( | ||
| make_idx(begin[i], end[i], strides[i], a.shape[i], indices[i])) | ||
| #return tvm.all(*from_val) |
There was a problem hiding this comment.
Shall we remove those 4 lines?
|
@yzhliu It can be used for the gradient of strided_slice in Relay, but is also an operation we need for myia (https://github.com/mila-iqia/myia). @junrushao1994 Yes it is that operation. |
|
@abergeron Got it. Thanks. |
python/tvm/relay/op/transform.py
Outdated
| v : relay.Expr | ||
| The data to be set. | ||
|
|
||
| begin: list of int |
There was a problem hiding this comment.
Thank you for adding this op. Currently @yongwww is modifying strided_slice(#4312) to support begin, end and strides to be expression instead of just list of int. The reason is that in some DL frameworks, begin, end or strides can be a tensor. Also making it more dynamic can help us when building other ops, such as NMS. Considering this op is similar to strided_slice, should we keep it align with stride_slice, and allow begin, end and strides to be Expr?
There was a problem hiding this comment.
That could be done, yes. The underlying TOPI op does support TVM expressions so it shouldn't be too hard to do.
There was a problem hiding this comment.
This is now done.
| b_np = np.asarray(begin).astype('int32') | ||
| e_np = np.asarray(end).astype('int32') | ||
| out_npy = topi.testing.strided_set_python( | ||
| x_np, v_np, begin, end, strides) + 1 |
There was a problem hiding this comment.
just for curiosity, why is +1 needed here?
There was a problem hiding this comment.
It mirrors the +1 in topi expression above. As to why this does +1 at all, I don't know, but strided_slice does that so I did the same.
|
Can this move forward? |
|
cc @yzhliu |
|
@junrushao1994 please https://docs.tvm.ai/contribute/code_review.html#approve-and-request-changes-explicitly if there is not further comment, i propose to merge it in 24 hours |
This adds essentially the inverse of strided_slice where you set the values selected to a new set of values. The result is a new tensor that is a copy of the first input with some values replaced.