Hi All: I start to study TVM recently, and we offen use the compute as the following:
... ...
rc = tvm.reduce_axis((0, in_channel), name='rc')
ry = tvm.reduce_axis((0, kernel), name='ry')
rx = tvm.reduce_axis((0, kernel), name='rx')
# Compute the convolution
B = tvm.compute(
(out_size, out_size, out_channel, batch),
lambda yy, xx, ff, nn: tvm.sum(
A[yy * stride + ry, xx * stride + rx, rc, nn] * W[ry, rx, rc, ff],
axis=[ry, rx, rc]),
name='B')
s = tvm.create_schedule(B.op)
print(tvm.lower(s,[A,W], simple_mode=True))
And we got some code like:
// attr [B] storage_scope = "global"
allocate B[float32 * 25690112]
produce B {
for (yy, 0, 14) {
for (xx, 0, 14) {
for (ff, 0, 512) {
for (nn, 0, 256) {
B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] = 0f
for (ry, 0, 3) {
for (rx, 0, 3) {
for (rc, 0, 256) {
B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] = (B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] + (A[((((((yy*917504) + (ry*917504)) + (xx*65536)) + (rx*65536)) + (rc*256)) + nn)]*W[((((ry*393216) + (rx*131072)) + (rc*512)) + ff)]))
}
}
}
}
}
}
}
}
And my question is if I want to add some expression between the for loop, and the pseudocode such as:
... ...
for (yy, 0, 14) {
for (xx, 0, 14) {
for (ff, 0, 512) {
for (nn, 0, 256) {
B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] = 0f
for (ry, 0, 3) {
for (rx, 0, 3) {
for (rc, 0, 256) {
B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] = (B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] + (A[((((((yy*917504) + (ry*917504)) + (xx*65536)) + (rx*65536)) + (rc*256)) + nn)]*W[((((ry*393216) + (rx*131072)) + (rc*512)) + ff)]))
}
}
}
B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)] = tvm.if_else_then(B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)]>0, B[((((yy*1835008) + (xx*131072)) + (ff*256)) + nn)], 0.0)
}
}
}
}
}
How can I realize it? Could any one give me an example？ Thanks a lot.