Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
LeiWang1999 opened this issue Sep 14, 2024 · 2 comments
Labels
needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it type: bug

Comments

@LeiWang1999
Copy link
Contributor

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

@LeiWang1999 LeiWang1999 added needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it type: bug labels Sep 14, 2024
@spectrometerHBH
Copy link
Contributor

LGTM

@LeiWang1999
Copy link
Contributor Author

I made a simple fix, CC @spectrometerHBH

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it type: bug
Projects
None yet
Development

No branches or pull requests

2 participants