While trying to fix https://github.com/dmlc/tvm/issues/3072 , I ran into another issue that has been raised a few months ago: https://github.com/dmlc/tvm/issues/1139
I’ve changed the bounds returned by PassUpDomain to be tighter (currently the full extents are used, which can be inefficient). The problem is that in some cases allocate size ends up depending on the loop variable. StorageRewrite pass moves allocate-s to the outermost scope, where the variable doesn’t exist yet, and we end up with something like this:
// attr [A1] storage_scope = "global"
allocate A1[float32 * (((((i.j.fused.outer*10) % 6) + 15)/6)*6)]
produce A2 {
for (i.j.fused.outer, 0, 8) {
produce A1 {
for (i, 0, ((((i.j.fused.outer*10) % 6) + 15)/6)) {
for (j, 0, 6) {
if (likely((((i.j.fused.outer*10)/6) < (12 - i)))) {
A1[((i*6) + j)] = A[(((((i.j.fused.outer*10)/6) + i)*6) + j)]
}
}
}
}
for (i.j.fused.inner, 0, 10) {
if (likely(((i.j.fused.outer*10) < (72 - i.j.fused.inner)))) {
A2[((i.j.fused.outer*10) + i.j.fused.inner)] = (A1[(((i.j.fused.outer*10) + i.j.fused.inner) - (((i.j.fused.outer*10)/6)*6))] + 3.000000f)
}
}
}
}
To resolve this, I need to replace the expression for allocate size with the maximum for the ranges of loop variables appearing in that expression.
I see 2 possible solutions:
- attempt to make interval extents constant with respect to loop variables in PassUpDomain
- compute maximum in StorageRewrite when allocate is being moved out of the loops.
The first one seems like it may be easier, but the second one has an advantage of not throwing away information that may be needed by some other passes between InferBounds and StorageRewrite.
Please share thoughts, suggestions, better solutions.