yongfeng-nv commented on a change in pull request #5382:
URL: https://github.com/apache/incubator-tvm/pull/5382#discussion_r421228975
##########
File path: src/te/operation/op_util.cc
##########
@@ -164,9 +164,21 @@ MakeLoopNest(const Stage& stage,
value_map[iv] = dom->min;
} else {
runtime::ThreadScope ts =
runtime::ThreadScope::make(bind_iv->thread_tag);
- if (stage->scope == "" || stage->scope == "warp" ||
+ if (stage->scope == "" ||
static_cast<int>(runtime::StorageScope::make(stage->scope).rank)
<= ts.rank) {
value_map[iv] = var;
+ } else if (stage->scope == "warp" && ts.rank == 1) {
+ // To determine whether a thread index is inside or outside a warp,
we need
+ // to know the thread extent. We leave a warning for now.
+ if (ts.dim_index == 0) {
+ value_map[iv] = var;
+ } else {
Review comment:
I modified the simplified example a little to bind threadIdx.y in the
warp stage to let threadIdx.y pass through the new code.
import tvm
import topi
import numpy as np
from tvm import te
n = 32
A = te.placeholder((2, n, n), name='A', dtype="float32")
C = te.compute((2, n, n), lambda x, i, j: A(x, i, (j + 1) % n), name='C')
s = te.create_schedule(C.op)
bk_x = te.thread_axis("blockIdx.x")
th_y = te.thread_axis("threadIdx.y")
th_x = te.thread_axis("threadIdx.x")
B = s.cache_read(A, "warp", [C])
cx, ci, cj = C.op.axis
bx, bi, bj = B.op.axis
# s[C].bind(ci, th_y)
s[C].bind(cj, th_x)
s[C].bind(cx, bk_x)
s[B].compute_at(s[C], cx)
s[B].bind(bi, th_y)
s[B].bind(bj, th_x)
print(tvm.lower(s, [A, C]))
func = tvm.build(s, [A, C], target="cuda", name='tid')
print(func.imported_modules[0].get_source())
The three situations make a good summary.
1st one already has at least one test in
tests/python/unittest/test_tir_transform_lower_warp_memory.py.
I hope the above code can lock down the 2nd situation and probably the error
for the 3rd one by reducing threadIdx.x's extent.
Warp has been special cased in several places, e.g. in bound.cc and here
before this PR. I tried to push back to add more special case code, but I am
Ok the accept the current change. Please try to add tests.
----------------------------------------------------------------
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]
