apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.76k stars 3.47k forks source link

[Bug] Improper touched buffer assignment of Pass `MergeSharedMemoryAllocations` #17375

Open LeiWang1999 opened 1 month ago

LeiWang1999 commented 1 month ago

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:

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);
}

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. :)

spectrometerHBH commented 1 month ago

LGTM

LeiWang1999 commented 1 month ago

I made a simple fix, CC @spectrometerHBH