[RELAY][Fix] i64 indices #5235
[RELAY][Fix] i64 indices #5235
Conversation
|
@hzfan can you followup? |
|
@tqchen ok, I will follow up. |
hzfan
force-pushed
the
relay-large-tensor_pr
branch
from
318fea3
to
c79dc79
Compare
hzfan
force-pushed
the
relay-large-tensor_pr
branch
3 times, most recently
from
0b16f80
to
e941d98
Compare
hzfan
force-pushed
the
relay-large-tensor_pr
branch
from
e941d98
to
112bdec
Compare
| return relay.Function([x], y) | ||
|
|
||
| orig = before() | ||
| fuse0(tvm.IRModule.from_expr(orig)) |
| return relay.Function([x], y) | ||
|
|
||
| orig = before() | ||
| fuse0(tvm.IRModule.from_expr(orig)) |
src/te/schedule/operation_inline.cc
Outdated
| @@ -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 | |||
There was a problem hiding this comment.
I dont understand this comment, can you try to clarify it?
There was a problem hiding this 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 (a) has a shape of i32.
There was a problem hiding this comment.
i see, it makes sense that we need to cast to the original indexing type of the operator being fused into. Why though can we "safely cast the indices without worrying about overflow"? Isn't it still possible to overflow given a sufficiently large index?
There was a problem hiding this comment.
You are right. The comment is a bit misleading. I think I can change it to 'cast indices to the type of the original indexing variable'
|
@jwfromm please take another look and approve explicitly if everything looks good to you. |
|
LGTM |
* fix * resolve comments
* fix * resolve comments
* fix * resolve comments
* fix * resolve comments
Fix i64 indices issues. Tests are taken from #5219.