Hi all,
I defined a toy computation and scheduled it in TVM. I am having some
difficulty in understanding how the lowered code that TVM produces corresponds
to the schedule. I have reproduced both the Python and the lowered IR below.
Python code:
```
import tvm
from tvm import te
batch_size = 24
hidden_size = 256
length = 100
scan_state = te.placeholder((length, batch_size, hidden_size))
scan_init = te.placeholder((1, batch_size, hidden_size))
C = te.compute((length, batch_size, hidden_size), lambda t, b, i: scan_state[t
- 1, b, i] * 4)
D = te.compute((length, batch_size, hidden_size), lambda t, b, i: C[t, b, i] +
7892)
scan_update = D
scan = tvm.te.scan(scan_init, scan_update, scan_state)
s = te.create_schedule([scan.op])
bx = te.thread_axis((0, batch_size), "blockIdx.x")
tx = te.thread_axis((0, hidden_size), "threadIdx.x")
s[scan].env_threads([bx])
xo, xi = s[C].split(s[C].op.axis[1], factor=24)
s[C].bind(xi, bx)
s[C].bind(s[C].op.axis[2], tx)
xo, xi = s[D].split(s[D].op.axis[1], factor=24)
s[D].bind(xi, bx)
s[D].bind(s[D].op.axis[2], tx)
print(tvm.lower(s, [scan_init, scan_state, scan_update, scan], simple_mode =
True))
```
Lowered IR:
```
produce scan {
// attr [iter_var(blockIdx.x, range(min=0, ext=24), blockIdx.x)]
thread_extent = 24
// attr [compute] storage_scope = "shared"
allocate compute[float32 * 256]
for (scan.idx, 0, 99) {
produce compute {
// attr [iter_var(threadIdx.x, range(min=0, ext=256), threadIdx.x)]
thread_extent = 256
if (likely((blockIdx.x < 1))) {
if (likely((blockIdx.x < 12))) {
compute[((blockIdx.x*256) + threadIdx.x)] = (scan[(((scan.idx*6144)
+ (blockIdx.x*512)) + threadIdx.x)]*4f)
}
}
}
// attr [iter_var(threadIdx.x, range(min=0, ext=256), threadIdx.x)]
thread_extent = 256
scan[((((scan.idx*6144) + (blockIdx.x*256)) + threadIdx.x) + 6144)] =
(compute[threadIdx.x] + 7892f)
}
}
```
Specifically I do not understand the predicates (`blockIdx.x < 1` and
`blockIdx.x < 12`) in the computation on the first compute op.
---
[Visit Topic](https://discuss.tvm.ai/t/unintuitive-lowered-code-in-tvm/6081/1)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.ai/email/unsubscribe/72db99ccd96b211b3074773d19cd09088950d64098d3afe4ff335faa2d03e904).
