Skip to content

Commit

Permalink
[TIR][Schedule] Scoped CacheRead/Write producing compact region (#15236)
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
MasterJH5574 and spectrometerHBH authored Jul 6, 2023
1 parent 73a62f6 commit 81463d7
Show file tree
Hide file tree
Showing 3 changed files with 403 additions and 73 deletions.
4 changes: 2 additions & 2 deletions src/tir/schedule/primitive.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
*
* \param rand_state The random state
* \param extent The loop extent to be tiled
Expand All @@ -123,7 +123,7 @@ TVM_DLL std::vector<int64_t> SamplePartitionedTile(
* 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
*
* \param rand_state The random state
* \param loop_sref The loop to be tiled
Expand Down
Loading

0 comments on commit 81463d7

Please sign in to comment.