pr #9341 introduced liveness analysis to merge the shared memory allocations , places touched buffer records at the outermost scope (e.g., outer loops) rather than at the innermost possible scope (e.g., inner loops or conditional branches). This can lead to incorrect liveness analysis. This approach luckily works well for some cases, such as the GEMM kernel, it fails in more complex scenarios, like the batched GEMM case or more complex algos. (as the outermost loop is always the single for loop node, that lead to incorrect gen kill point for each buffer).
One solutions I'm applying is to replace:
if (it != alloc_info_.end() && it->second.alloc) {
ICHECK_LT(it->second.level, scope_.size());
if (IsAppropriateSharedMemory(GetRef<Var>(buf))) {
scope_[it->second.level].touched.push_back(buf);
}
}
into
if (IsAppropriateSharedMemory(GetRef<Var>(buf))) {
scope_[scope_.size() - 1].touched.push_back(buf);
}
Lead to Suboptimal Shared Memory Reuse.
pr #9341 introduced liveness analysis to merge the shared memory allocations , places touched buffer records at the outermost scope (e.g., outer loops) rather than at the innermost possible scope (e.g., inner loops or conditional branches). This can lead to incorrect liveness analysis. This approach luckily works well for some cases, such as the GEMM kernel, it fails in more complex scenarios, like the batched GEMM case or more complex algos. (as the outermost loop is always the single for loop node, that lead to incorrect gen kill point for each buffer).
One solutions I'm applying is to replace:
into
more detailed analysis can be found at TVM Shared Memory Reuse Analysis
If you think this analysis is correct, I can submit a PR then. :)