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

[LeiWang1999]

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]

LGTM

[LeiWang1999]

I made a simple fix, CC @spectrometerHBH

[LeiWang1999]

Closed, as it doesn’t need to be addressed for now.