[MetaSchedule] Support schedules with cache read in RewriteLayout #13384
Conversation
|
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.
Generated by tvm-bot |
masahi
force-pushed
the
ms-rewrite-layout-fix
branch
from
c924fc1 to
29aa4ee
Compare
| n->annotations.erase(topi_attr); | ||
| } else { | ||
| n->annotations.Set(topi_attr, new_buffers); | ||
| } |
There was a problem hiding this comment.
This is a fix for RewriteLayout to work when link-params = True, the breakage discussed in #13195. With this, I confirmed that the new RewriteLayout implementation works when there is CacheRead acting on AllocateConst. I'm not adding a test since TVMScript parsing is broken when there is AllocateConst with a large constant (can't print the whole array).
There was a problem hiding this comment.
Continuing the discussion from this thread.
I am convinced that we need to store buffers in block attributes for hexagon-specific usecases, and the only thing that I feel less natural is generally storing IR nodes in attributes. As an example, in layout_free_placeholders, we didn't store any IR nodes in the attribute, but instead use a list of integers which is simpler.
In the meantime, I completely understand that we need to get around this quickly, so in this particular case, how about adding or reusing a pass config flag, e.g. the one we are using "link-params", and only add enable topi_attr to be set when the flag is on?
else if (link_param is on) {
n->annotations.Set(topi_attr, new_buffers);
}There was a problem hiding this comment.
An alternative idea: having an extra flag in CreatePrimFunc that instructs the translator to retain buffers in block annotation, and the flag is on only if link-param is detected as "on" in TECompiler.
There was a problem hiding this comment.
This else branch is only hit when there is a layout-free buffer that's not passed as a parameter to the prim func. Currently that can only happen when link-params = True (this is exactly what link-params is meant for). So isn't having another flag redundant?
There was a problem hiding this comment.
I see. That makes perfect sense to me :-) As long as it's not affecting normal lowering flow, it looks like a good idea!
…ache#13384) Currently when `CacheRead` and `RewriteLayout` are used together, the index map is derived based on the cache read block, which leads to weird result. This is because the current implementation assumes that the "layout-free" buffer is directly consumed by an "anchor" op such as conv2d or dense. When `CacheRead` is involved, we need to find the index map for the cache-read buffer as it is consumed by an anchor op, and apply the same transformation to the layout-free buffer. My solution supports more general cases where there are multiple cache reads forming a "chain" of blocks, starting from the one that directly consumes the layout-free buffer passed as a parameter. So the layout transformation is back propagated over such chain.
Currently when
CacheReadandRewriteLayoutare used together, the index map is derived based on the cache read block, which leads to weird result. This is because the current implementation assumes that the "layout-free" buffer is directly consumed by an "anchor" op such as conv2d or dense.When
CacheReadis involved, we need to find the index map for the cache-read buffer as it is consumed by an anchor op, and apply the same transformation to the layout-free buffer. My solution supports more general cases where there are multiple cache reads forming a "chain" of blocks, starting from the one that directly consumes the layout-free buffer passed as a parameter. So the layout transformation is back propagated over such chain.cc @vinx13 @junrushao @nverke