hzfan commented on a change in pull request #5235:
URL: https://github.com/apache/incubator-tvm/pull/5235#discussion_r455514589
##########
File path: src/te/schedule/operation_inline.cc
##########
@@ -63,7 +63,9 @@ class OperationInliner final : public StmtExprMutator {
} else {
Map<Var, PrimExpr> vmap;
for (size_t i = 0; i < args_.size(); ++i) {
- vmap.Set(args_[i], op->indices[i]);
+ // indices into `operation_` must be in the range of its output
shape,
+ // so we can safely cast the indices without worrying about overflow
Review comment:
Sure. An example is
```
a_i32 = tvm.tir.const(2, “int32”)
b_i64 = tvm.tir.const(2 ** 32, “int64”)
a0 = te.placeholder((a_i32,), “a0”)
a = te.compute((a_i32,), lambda i: a0[i])
b = te.compute((b_i64,), lambda i: a[i % 2])
```
When we fuse `a` (denoted by `operation_` in the comment) into `b`, the
variable `i` (which is i32) in the computation of `a` is replaced by `i % 2`
(which is i64) in `b`. Since the original indexing variable in `a` `i` is i32,
we need to cast `i % 2` from i64 to i32.
By the comment, I mean `i % 2` must fit in i32, because the tensor it
indexes into has a shape of i32.
##########
File path: src/te/schedule/operation_inline.cc
##########
@@ -63,7 +63,9 @@ class OperationInliner final : public StmtExprMutator {
} else {
Map<Var, PrimExpr> vmap;
for (size_t i = 0; i < args_.size(); ++i) {
- vmap.Set(args_[i], op->indices[i]);
+ // indices into `operation_` must be in the range of its output
shape,
+ // so we can safely cast the indices without worrying about overflow
Review comment:
Sure. An example is
```
a_i32 = tvm.tir.const(2, “int32”)
b_i64 = tvm.tir.const(2 ** 32, “int64”)
a0 = te.placeholder((a_i32,), “a0”)
a = te.compute((a_i32,), lambda i: a0[i])
b = te.compute((b_i64,), lambda i: a[i % 2])
```
When we fuse `a` (denoted by `operation_` in the comment) into `b`, the
variable `i` (which is i32) in the computation of `a` is replaced by `i % 2`
(which is i64) in `b`. Since the original indexing variable `i` in `a` is i32,
we need to cast `i % 2` from i64 to i32.
By the comment, I mean `i % 2` must fit in i32, because the tensor it
indexes into has a shape of i32.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
