masahi commented on a change in pull request #8971:
URL: https://github.com/apache/tvm/pull/8971#discussion_r707850237
##########
File path: src/te/tensor.cc
##########
@@ -39,15 +39,26 @@ IterVar reduce_axis(Range dom, std::string name) { return
IterVar(dom, Var(name)
Var var(std::string name_hint, DataType t) { return Var(name_hint, t); }
// Tensor
-PrimExpr Tensor::operator()(Array<Var> indices) const {
+PrimExpr Tensor::operator()(Array<Var> indices, bool support_negative_indices)
const {
Array<PrimExpr> arr(indices.begin(), indices.end());
- return operator()(arr);
+ return operator()(arr, support_negative_indices);
}
-PrimExpr Tensor::operator()(Array<PrimExpr> indices) const {
- if (ndim() != 0) {
- ICHECK_EQ(ndim(), indices.size()) << "Tensor dimension mismatch in read "
- << "ndim = " << ndim() << ",
indices.size=" << indices.size();
+PrimExpr Tensor::operator()(Array<PrimExpr> indices, bool
support_negative_indices) const {
+ Array<PrimExpr> shape = (*this)->shape;
+
+ if (shape.size() != 0) {
+ ICHECK_EQ(shape.size(), indices.size())
+ << "Tensor dimension mismatch in read "
+ << "ndim = " << ndim() << ", indices.size=" << indices.size();
+ }
+
+ 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),
+ indices[i] + shape[i], indices[i]);
+ indices.Set(i, new_index);
Review comment:
Negative indices handling is also done in
https://github.com/apache/tvm/blob/d9fe67259ead828f174e0be719906e3f2d3c2137/python/tvm/relay/op/transform.py#L926-L927
https://github.com/apache/tvm/blob/cbe3dcad5e3f9358af8d2d79e2880f92718a5d0b/include/tvm/topi/detail/strided_slice.h#L45-L48
https://github.com/apache/tvm/blob/cbe3dcad5e3f9358af8d2d79e2880f92718a5d0b/include/tvm/topi/detail/strided_slice.h#L105
I believe there are other cases like this spread across the code base. Maybe
we should revisit all index-taking op and centralize negative indices handling.
Generally I think people prefer not making a change down the stack.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
