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

[TIR][Schedule] Scoped CacheRead/Write producing compact region #15236

Merged

Conversation

MasterJH5574
Copy link
Contributor

This PR enhances CacheRead/Write so that when a cache operation is performed under an inner block, the generated cache buffer will have the shape as compact as possible, by region consumption analysis.

The motivation of this change comes from the needs of dynamic shape TIR scheduling, in which case we may isolate a "static shape" internal block using blockize, and do further scheduling inside the internal block. For such cases, the current CacheRead/Write inside the static-shape block will still produce dynamic-shape cache buffers, which is not ideal for analysis and subsequent scheduling.

One thing that worths noting is that, to ensure the IR correctness after inserting the cache block, we will only compact the cache buffer when all the consumer blocks of the read buffer (for CacheRead) or the write buffer (for CacheWrite) are children blocks of the cache block insertion location. Otherwise we will insist allocating the full-size cache buffer.

Co-authored-by: Bohan Hou spectrometerh@gmail.com

@tvm-bot
Copy link
Collaborator

tvm-bot commented Jul 5, 2023

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

@@ -105,7 +105,7 @@ TVM_DLL std::vector<int64_t> SamplePerfectTile(
* The sampled tile size will be partitioned into two parts. The second part has a guarantee
* that their extent's product have a factor of `innerpart_factor`. The first part is loops at
* [0, partition_pos); the second part is loops at [partition_pos, n) and we will have
* `innerpart_factor` | \prod_{l=partition_pos}^{n-1} l.extent
* `innerpart_factor` | prod_{l=partition_pos}^{n-1} l.extent
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updating this since Clang will warn

1 warning generated.
In file included from /home/ruihangl/tvm/src/tir/schedule/primitive/for_kind.cc:19:
In file included from /home/ruihangl/tvm/src/tir/schedule/primitive/../utils.h:49:
/home/ruihangl/tvm/src/tir/schedule/primitive/.././primitive.h:108:26: warning: unknown command tag name [-Wdocumentation-unknown-command]
 *  `innerpart_factor` | \prod_{l=partition_pos}^{n-1} l.extent
                         ^~~~~~

Copy link
Member

@junrushao junrushao left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM! Likely it won't break anything as scopes don't happen a lot in existing TIR/MS pipeline

@MasterJH5574 MasterJH5574 force-pushed the tvm-dev/2023-07-05-cache-read-write branch 2 times, most recently from 3ce17c3 to 6c35f3f Compare July 5, 2023 21:11
This PR enhances CacheRead/Write so that when a cache operation is
performed under an inner block, the generated cache buffer will have
the shape as compact as possible, by region consumption analysis.

The motivation of this change comes from the needs of dynamic shape TIR
scheduling, in which case we may isolate a "static shape" internal block
using blockize, and do further scheduling inside the internal block. For
such cases, the current CacheRead/Write inside the static-shape block
will still produce dynamic-shape cache buffers, which is not ideal for
analysis and subsequent scheduling.

One thing that worths noting is that, to ensure the IR correctness after
inserting the cache block, we will only compact the cache buffer when
all the consumer blocks of the read buffer (for CacheRead) or the write
buffer (for CacheWrite) are children blocks of the cache block insertion
location. Otherwise we will insist allocating the full-size cache
buffer.

Co-authored-by: Bohan Hou <spectrometerh@gmail.com>
@MasterJH5574 MasterJH5574 force-pushed the tvm-dev/2023-07-05-cache-read-write branch from 6c35f3f to 1d432c5 Compare July 6, 2023 00:31
@tqchen tqchen merged commit 81463d7 into apache:main Jul 6, 2023
5 checks passed
junrushao pushed a commit to junrushao/tvm that referenced this pull request Jul 15, 2023
…he#15236)

This PR enhances CacheRead/Write so that when a cache operation is
performed under an inner block, the generated cache buffer will have
the shape as compact as possible, by region consumption analysis.

The motivation of this change comes from the needs of dynamic shape TIR
scheduling, in which case we may isolate a "static shape" internal block
using blockize, and do further scheduling inside the internal block. For
such cases, the current CacheRead/Write inside the static-shape block
will still produce dynamic-shape cache buffers, which is not ideal for
analysis and subsequent scheduling.

One thing that worths noting is that, to ensure the IR correctness after
inserting the cache block, we will only compact the cache buffer when
all the consumer blocks of the read buffer (for CacheRead) or the write
buffer (for CacheWrite) are children blocks of the cache block insertion
location. Otherwise we will insist allocating the full-size cache
buffer.

Co-authored-by: Bohan Hou <spectrometerh@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants