[PASS] Update StorageRewrite to Handle Allocate Depending on Loop Variable

[derisavi-huawei]

Consider the following hypothetical example:

import tvm
N = tvm.var("N")
y = tvm.reduce_axis((0,N), 'y')
A = tvm.placeholder((2*N,2*N), "float32", "A")
C = tvm.compute((N, N), lambda i,j: A[2*i,j] + A[i+1,j], name='C')
s = tvm.create_schedule(C.op)
AA = s.cache_read(A, "local", [C])
s[AA].compute_at(s[C], s[C].op.axis[0])
print(tvm.lower(s, [A,C], simple_mode=True))

The resulting lowered code is as follows. See how the size of A.local depends on i which is a loop variable. I think the problem is that the two references to A, i.e., A[2*i,j] and A[i+1,j] have different functions of i (i.e., 2*i vs i) in their first dimension. That makes calculating buffer size more complicated, and TVM probably doesn't check for that. We should either handle it properly or emit an error message that such indexing is not supported.

// attr [A.local] storage_scope = "local"
allocate A.local[float32 * ((max((i*2), (i + 1)) - min((i*2), (i + 1))) + 1) * N]
produce C {
  for (i, 0, N) {
    produce A.local {
      for (ax0, 0, ((max((i*2), (i + 1)) - min((i*2), (i + 1))) + 1)) {
        for (ax1, 0, N) {
          if (likely((min((i*2), (i + 1)) < ((N*2) - ax0)))) {
            if (likely((ax1 < (N*2)))) {
              A.local[((ax0*N) + ax1)] = A[((((min((i*2), (i + 1)) + ax0)*N)*2) + ax1)]
            }
          }
        }
      }
    }
    for (j, 0, N) {
      C[((i*N) + j)] = (A.local[((max((i + -1), 0)*N) + j)] + A.local[(((1 - min(i, 1))*N) + j)])
    }
  }
}

[tqchen]

For general CPU code, this code is still valid, but the allocate should not be lifted into the outer most scope. We should update StorageRewrite to check for this case(and possibly report an error since local usually need static size)

[wweic]

In [8]: import tvm
   ...: N = tvm.var("N")
   ...: A = tvm.placeholder((N,N), "float32", "A")
   ...: C = tvm.compute((N, N), lambda i,j: A[i,j] + 1, name='C')
   ...: s = tvm.create_schedule(C.op)
   ...: AA = s.cache_read(A, "local", [C])
   ...: s[AA].compute_at(s[C], s[C].op.axis[0])
   ...: print(tvm.lower(s, [A,C], simple_mode=True))
   ...:
   ...:
// attr [A.local] storage_scope = "local"
allocate A.local[float32 * 1 * N]
produce C {
  for (i, 0, N) {
    produce A.local {
      for (ax1, 0, N) {
        A.local[ax1] = A[((i*N) + ax1)]
      }
    }
    for (j, 0, N) {
      C[((i*N) + j)] = (A.local[j] + 1.000000f)
    }
  }
}

Looks like even for allocate that doesn't depend on loop variable, the allocate stmt is in global scope. Is this expected? (@tqchen). Are there other IR passes that might possibly change the AttStmt's position after lower? I can take a look if it's not by design.

[tqchen]

The "global scope" is a bit misleading, as it means global memory in GPU, or the memory that can actually be seen by the user. In CPU, everything is "global memory", although llvm may promote constant arrays into register.

The special memory "local, shared" is only enabled when we tag them to be so

[wweic]

Ah thanks for catching that! I meant the outer most scope in your first reply. Back to the example, do you still think the allocate statement is not in the right scope? Or it should be something similar to below?

produce C {
  for (i, 0, N) {
    // attr [A.local] storage_scope = "local"
    allocate A.local[float32 * 1 * N]

    produce A.local {
      for (ax1, 0, N) {
        A.local[ax1] = A[((i*N) + ax1)]
      }
    }
    for (j, 0, N) {
      C[((i*N) + j)] = (A.local[j] + 1.000000f)
    }
  }
}

[tqchen]

Yes, in this case, the allocate should be in the for loop(as shown in your code). We probably want to send an warning as well, because usually non-fixed size temporal memory means dynamic memory allocation and we can only use DRAM(CPU memory) and global scope

[wweic]

OK. The reason the allocate was put into outer most scope is because the code only call PlanNewScope for ForType::Parallel:
https://github.com/dmlc/tvm/blob/7bafca4e550d451925abfa2f0bc897d080cbbf35/src/pass/storage_rewrite.cc#L754-L757

Is it still the case?

Also we enter/exit scope AFTER we call FindAlloc
https://github.com/dmlc/tvm/blob/7bafca4e550d451925abfa2f0bc897d080cbbf35/src/pass/storage_rewrite.cc#L736

It makes more sense to me that we move enter/exit scope code right after:
https://github.com/dmlc/tvm/blob/7bafca4e550d451925abfa2f0bc897d080cbbf35/src/pass/storage_rewrite.cc#L693

After making the fixes, it does generate the expected IR. While 6 python tests failed in test_pass_storage_rewrite.py. I'll dig, but want to check that I'm heading to the right direction @tqchen

[tqchen]

The original assumption was indeed that most memory allocation should appear in the outermost scope of the threading boundary. This is true and useful for most memory optimizations, as dynamic allocating and de-allocating is not good for performance and do not work for on-device memory.

Maybe in cases where there is dependent memory, we should either shoot an error, or relax the memory boundary to its maximum value then run the lifting, thoughts?

[tqchen]

The reasoning was not about correctly implementing something, but more about if we can use the simplest method to guarantee the correct and efficient usage. The original decision on restricting allocation site is made with that reasoning.

[wweic]

I see. I would vote to throw error to start, and if we see more occurrence of such pattern, we can consider handling it properly.

[wweic]

simple_mode skips the sanity check by IRUseDefAnalysis. So I added a test to show that we can detect such patterns.