Skip to content

[TIR] Change the behavior of read/write region analysis for reduction blocks. #10638

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.

Sign up for GitHub

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

Merged
junrushao merged 11 commits into from
Mar 23, 2022
Merged

[TIR] Change the behavior of read/write region analysis for reduction blocks. #10638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
d5ef430
upd
yzh119 Mar 16, 2022
7fe6228
fix
yzh119 Mar 16, 2022
d826300
fix rfactor
yzh119 Mar 16, 2022
2684364
fix blockize
yzh119 Mar 16, 2022
8805c1e
tensorize
yzh119 Mar 16, 2022
e1a7c50
upd
yzh119 Mar 22, 2022
9dc8273
fix meta-schedules
yzh119 Mar 22, 2022
c641c30
upd
yzh119 Mar 22, 2022
d268ba0
fix argmax
yzh119 Mar 23, 2022
9a59bdc
manally call region analysis
yzh119 Mar 23, 2022
cd42ab1
Merge branch 'main' into reduction-block-reads
yzh119 Mar 23, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 45 additions & 34 deletions src/tir/analysis/block_access_region_detector.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,9 +41,11 @@ class BlockReadWriteDetector : public StmtExprVisitor {
: buffer_var_map_(buffer_var_map) {}

/*! \brief Return read regions of the block */
Array<BufferRegion> CollectReads();
Array<BufferRegion> CollectReads(
const std::unordered_set<const BufferNode*>* excluded_buffers = nullptr);
/*! \brief Return write regions of the block */
Array<BufferRegion> CollectWrites();
Array<BufferRegion> CollectWrites(
const std::unordered_set<const BufferNode*>* excluded_buffers = nullptr);
/*!
* \brief Return opaque buffer regions of the block
* \note The buffer accessed by load/store or call with buffer.data will
Expand Down Expand Up @@ -88,8 +90,10 @@ class BlockReadWriteDetector : public StmtExprVisitor {
Buffer buffer, std::vector<arith::IntSet> region);

/*! \brief Helper function to collect access regions. */
Array<BufferRegion> CollectRegions(const std::vector<Buffer>& buffers,
const std::vector<std::vector<tvm::arith::IntSet>>& regions);
Array<BufferRegion> CollectRegions(
const std::vector<Buffer>& buffers,
const std::vector<std::vector<tvm::arith::IntSet>>& regions,
const std::unordered_set<const BufferNode*>* excluded_buffers = nullptr);

/*! \brief Helper function to convert matched access region to source region. */
std::vector<arith::IntSet> ConvertMatchedRegion(const MatchBufferRegion& match_buffer,
Expand Down Expand Up @@ -126,12 +130,14 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) {
StmtExprVisitor::operator()(stmt);
}

Array<BufferRegion> BlockReadWriteDetector::CollectReads() {
return CollectRegions(read_buffers_, read_regions_);
Array<BufferRegion> BlockReadWriteDetector::CollectReads(
const std::unordered_set<const BufferNode*>* excluded_buffers) {
return CollectRegions(read_buffers_, read_regions_, excluded_buffers);
}

Array<BufferRegion> BlockReadWriteDetector::CollectWrites() {
return CollectRegions(writes_buffers_, write_regions_);
Array<BufferRegion> BlockReadWriteDetector::CollectWrites(
const std::unordered_set<const BufferNode*>* excluded_buffers) {
return CollectRegions(writes_buffers_, write_regions_, excluded_buffers);
}

Array<BufferRegion> BlockReadWriteDetector::CollectOpaques() {
Expand Down Expand Up @@ -282,12 +288,15 @@ void BlockReadWriteDetector::Update(std::vector<Buffer>* buffers,
}

Array<BufferRegion> BlockReadWriteDetector::CollectRegions(
const std::vector<Buffer>& buffers,
const std::vector<std::vector<tvm::arith::IntSet>>& regions) {
const std::vector<Buffer>& buffers, const std::vector<std::vector<tvm::arith::IntSet>>& regions,
const std::unordered_set<const BufferNode*>* excluded_buffers) {
ICHECK_EQ(buffers.size(), regions.size());
Array<BufferRegion> res;
res.reserve(buffers.size());
for (size_t i = 0; i < regions.size(); ++i) {
if (excluded_buffers != nullptr && excluded_buffers->count(buffers[i].get())) {
continue;
}
Array<Range> region;
region.reserve(regions[i].size());
ICHECK_EQ(buffers[i]->shape.size(), regions[i].size());
Expand Down Expand Up @@ -319,38 +328,40 @@ Array<Array<BufferRegion>> GetBlockAccessRegion(const Block& block,
const Map<Var, Buffer>& buffer_var_map) {
BlockReadWriteDetector detector(buffer_var_map);
detector(block);
return {detector.CollectReads(), detector.CollectWrites(), detector.CollectOpaques()};
Array<BufferRegion> writes = detector.CollectWrites();
std::unordered_set<const BufferNode*> excluded_buffers;
// exclude write buffers from read regions for reductions if init block is defined.
if (block->init.defined()) {
for (const BufferRegion& write_access : writes) {
excluded_buffers.insert(write_access->buffer.get());
}
}
Array<BufferRegion> reads = detector.CollectReads(&excluded_buffers);
Array<BufferRegion> opaques = detector.CollectOpaques();
return {reads, writes, opaques};
}

Array<Array<BufferRegion>> GetBlockReadWriteRegion(const Block& block,
const Map<Var, Buffer>& buffer_var_map) {
// Step 1. Get all the read/write/opaque accesses in the input block.
Array<Array<BufferRegion>> access_regions = GetBlockAccessRegion(block, buffer_var_map);
// Step 2. Collect all the buffers that are opaquely accessed.
std::unordered_set<const BufferNode*> opaque_accessed_buffers;
for (const BufferRegion& opaque_access : access_regions[2]) {
opaque_accessed_buffers.insert(opaque_access->buffer.get());
}
// Step 3. Create new arrays of read/write regions.
Array<BufferRegion> new_read_regions;
Array<BufferRegion> new_write_regions;
new_read_regions.reserve(access_regions[0].size() + access_regions[2].size());
new_write_regions.reserve(access_regions[1].size() + access_regions[2].size());
for (const BufferRegion& read_access : access_regions[0]) {
if (!opaque_accessed_buffers.count(read_access->buffer.get())) {
new_read_regions.push_back(read_access);
}
BlockReadWriteDetector detector(buffer_var_map);
detector(block);
Array<BufferRegion> opaques = detector.CollectOpaques();
std::unordered_set<const BufferNode*> excluded_buffers;
for (const BufferRegion& opaque_access : opaques) {
excluded_buffers.insert(opaque_access->buffer.get());
}
for (const BufferRegion& write_access : access_regions[1]) {
if (!opaque_accessed_buffers.count(write_access->buffer.get())) {
new_write_regions.push_back(write_access);
Array<BufferRegion> writes = detector.CollectWrites(&excluded_buffers);
if (block->init.defined()) {
for (const BufferRegion& write_access : writes) {
excluded_buffers.insert(write_access->buffer.get());
}
}
for (const BufferRegion& opaque_access : access_regions[2]) {
new_read_regions.push_back(opaque_access);
new_write_regions.push_back(opaque_access);
Array<BufferRegion> reads = detector.CollectReads(&excluded_buffers);
for (const BufferRegion& opaque_access : opaques) {
reads.push_back(opaque_access);
writes.push_back(opaque_access);
}
return {new_read_regions, new_write_regions};
return {reads, writes};
}

TVM_REGISTER_GLOBAL("tir.analysis.GetBlockAccessRegion").set_body_typed(GetBlockAccessRegion);
Expand Down
20 changes: 20 additions & 0 deletions src/tir/schedule/primitive/blockize_tensorize.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -496,13 +496,33 @@ StmtSRef Blockize(ScheduleState self, const StmtSRef& loop_sref) {
Block new_block = Downcast<Block>(replacer(block));

// Step 6: Generate the inner block.
bool outer_reduction = false; // whether there are outer reduction iter vars.
for (const IterVar& iter_var : extractor.outer_iter_vars) {
if (iter_var->iter_type == kCommReduce) {
outer_reduction = true;
}
}
BlockRealizeNode* inner_block_realize = block_realize.CopyOnWrite();
inner_block_realize->iter_values = extractor.inner_bindings;
inner_block_realize->predicate = inner_pred;
inner_block_realize->block = new_block;
BlockNode* inner_block = inner_block_realize->block.CopyOnWrite();
inner_block->iter_vars = extractor.inner_iter_vars;
inner_block->init = NullOpt;
/* Add write regions to read regions if
* 1. there are outer reduction iter vars.
* 2. the init block is defined for current block.
*/
if (outer_reduction && block->init.defined()) {
Array<BufferRegion> new_reads;
for (const BufferRegion& write_access : inner_block->writes) {
new_reads.push_back(write_access);
}
for (const BufferRegion& read_access : inner_block->reads) {
new_reads.push_back(read_access);
}
inner_block->reads = std::move(new_reads);
}
block_sref_reuse.Set(block, inner_block_realize->block);

// Step 6: Generate the outer block.
Expand Down
20 changes: 17 additions & 3 deletions src/tir/schedule/primitive/reduction.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,22 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
/*body=*/body);
}
body = Substitute(body, loop_var_map);
// Step 6. Mutate IR
// Step 6. Add write regions back to read regions in update block.
Array<BufferRegion> new_reads;
std::unordered_set<const BufferNode*> read_bufs;
for (const BufferRegion& read_access : block->reads) {
read_bufs.insert(read_access->buffer.get());
}
for (const BufferRegion& write_access : block->writes) {
if (read_bufs.find(write_access->buffer.get()) == read_bufs.end()) {
new_reads.push_back(write_access);
}
}
for (const BufferRegion& read_access : block->reads) {
new_reads.push_back(read_access);
}
(const_cast<BlockNode*>(block))->reads = std::move(new_reads);

@masahi masahi Mar 27, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

After this PR, I'm getting a structural hash mismatch in meta schedule ApplyHistoryBest

tvm/src/meta_schedule/integration.cc

Line 134 in ce335c3

if (database->HasWorkload(prim_mod)) {

I found that if I do decompose_reduction, IRModules that have been committed to the database are modified. Specifically, the query module has

reads([placeholder[b, i, k], T_layout_trans[b, floordiv(j, 16), floordiv(k, 4), floormod(j, 16), floormod(k, 4)]])
writes([compute[b, i, j]])
...
with init() {
  compute[b, i, j] = 0
}
compute[b, i, j] = (compute[b, i, j] + (int32(placeholder[b, i, k])*int32(T_layout_trans[b, floordiv(j, 16), floordiv(k, 4), floormod(j, 16), floormod(k, 4)])))

but the corresponding mod in the database has different reads:

reads(compute[b, i, j], [placeholder[b, i, k], T_layout_trans[b, floordiv(j, 16), floordiv(k, 4), floormod(j, 16), floormod(k, 4)]])
...

I haven't looked at what this PR does, but I'm pretty sure this line is the offending bug...

@junrushao junrushao Mar 28, 2022

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

@masahi thanks for reporting this! I'm not sure why the IRModule is mutated in place after decompose-reduction (not sure if i understand correctly), but would be great to have a minimal reproducible example after you returning from vacation :-)

@yzh119 would you mind taking a look at the particular case Masa mentions and perhaps add a regression test?

@masahi masahi Mar 28, 2022 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

yeah, I can prepare a repro. It's a bit complicated since it involves scheduling, database stuff, and ApplyHistoryBest. I know the issue comes from this PR since I bisected it.

So when decompose_reduction is called, we are now executing

https://github.com/yzh119/tvm/blob/cd42ab101ab993edae2956fdc0415561be5f9718/src/tir/schedule/primitive/reduction.cc#L301

and I'm guessing that this in-place update propagates all the way up to the committed prim_mod in the tuning database.

@masahi masahi Mar 28, 2022 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

ok the repro is ready: https://gist.github.com/masahi/591078723d26f09ece3430af95835c99

The test works with the current main. When I run it, I get this output:

With decompose_reduction=True
One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_nn_contrib_dense_pack
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_expand_dims
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast_1
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul

With decompose_reduction=False
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_expand_dims
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast
[07:20:13] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast_1

This means, if decompose_reduction is applied during scheduling, ApplyHistoryBest fails to match the structural hash of the query mod, corresponding to nn_contrib_dense and nn_batch_matmul, against the ones in database.

If I revert this commit, I get the following expected output

With decompose_reduction=True
One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_expand_dims
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast_1

With decompose_reduction=False
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_expand_dims
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast
[07:24:21] /home/masa/projects/dev/tvm/src/meta_schedule/integration.cc:146: Warning: Cannot find workload: tvmgen_default_fused_cast_1

You can dump the contents of the database at

tvm/src/meta_schedule/database/json_database.cc

Lines 72 to 74 in b08e8e4

bool HasWorkload(const IRModule& mod) {
return workloads2idx_.find(Workload(mod, tvm::StructuralHash()(mod))) != workloads2idx_.end();
}
to see that reads region for dense and bmm is modified if decompose_reduction is enabled. This causes the structual hash mismatch between the modules in the database and the query mod.

// Step 7. Mutate IR
const BlockNode* old_scope_root = TVM_SREF_TO_BLOCK(old_scope_root, scope_root_sref);
Block new_scope_root{nullptr};
Block new_reduction_block{nullptr};
Expand Down Expand Up @@ -826,9 +841,8 @@ class WriteBackBlockCreator : public BaseBlockCreator {
}

void CreateReadWriteRegions() final {
read_regions_.push_back(CreateRegion(wb_lhs_));
read_regions_.push_back(CreateRegion(wb_rhs_));
write_regions_.push_back(read_regions_[0]);
write_regions_.push_back(CreateRegion(wb_lhs_));
}

static BufferRegion CreateRegion(const BufferLoad& load) {
Expand Down
6 changes: 3 additions & 3 deletions tests/python/unittest/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ def main(var_A: T.handle, var_B: T.handle, var_C: T.handle) -> None:
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3 * 16 + i0_4)
j = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3 * 2 + i1_4)
k = T.axis.reduce(512, i2_1 * 32 + i2_2)
T.reads([C_local[i, j], A_shared[i, k], B_shared[k, j]])
T.reads([A_shared[i, k], B_shared[k, j]])
T.writes([C_local[i, j]])
with T.init():
C_local[i, j] = T.float32(0)
Expand Down Expand Up @@ -129,13 +129,13 @@ def test_rewrite_cooperative_fetch():
sch.bind(loop=l32, thread_axis="vthread.x")
l33 = sch.fuse(l12, l22)
sch.bind(loop=l33, thread_axis="threadIdx.x")
b34 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")
b34 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")
sch.compute_at(block=b34, loop=l28, preserve_unit_loops=True)
_, _, _, _, l39, l40 = sch.get_loops(block=b34)
l41 = sch.fuse(l39, l40)
_, v43 = sch.sample_perfect_tile(loop=l41, n=2, max_innermost_factor=4, decision=[262144, 1])
sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v43)
b44 = sch.cache_read(block=b0, read_buffer_index=2, storage_scope="shared")
b44 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")
sch.compute_at(block=b44, loop=l28, preserve_unit_loops=True)
_, _, _, _, l49, l50 = sch.get_loops(block=b44)
l51 = sch.fuse(l49, l50)
Expand Down
8 changes: 4 additions & 4 deletions tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,13 +187,13 @@ def test_cuda_matmul():
'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)',
'b33 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")',
"sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=True)",
'b34 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
'b34 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")',
"sch.compute_at(block=b34, loop=l27, preserve_unit_loops=True)",
"l35, l36, l37, l38, l39, l40 = sch.get_loops(block=b34)",
"l41 = sch.fuse(l39, l40)",
"v42 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
'sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v42)',
'b43 = sch.cache_read(block=b0, read_buffer_index=2, storage_scope="shared")',
'b43 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
"sch.compute_at(block=b43, loop=l27, preserve_unit_loops=True)",
"l44, l45, l46, l47, l48, l49 = sch.get_loops(block=b43)",
"l50 = sch.fuse(l48, l49)",
Expand Down Expand Up @@ -241,13 +241,13 @@ def test_cuda_matmul_relu():
'sch.bind(loop=l32, thread_axis="threadIdx.x")',
'b33 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")',
"sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=True)",
'b34 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
'b34 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")',
"sch.compute_at(block=b34, loop=l27, preserve_unit_loops=True)",
"l35, l36, l37, l38, l39, l40 = sch.get_loops(block=b34)",
"l41 = sch.fuse(l39, l40)",
"v42 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
'sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v42)',
'b43 = sch.cache_read(block=b0, read_buffer_index=2, storage_scope="shared")',
'b43 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
"sch.compute_at(block=b43, loop=l27, preserve_unit_loops=True)",
"l44, l45, l46, l47, l48, l49 = sch.get_loops(block=b43)",
"l50 = sch.fuse(l48, l49)",
Expand Down
2 changes: 1 addition & 1 deletion tests/python/unittest/test_meta_schedule_tune_relay.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ def main(placeholder: T.Buffer[(1, 1, 16, 16, 3), "float32"], placeholder_1: T.B
for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(1, 2, 16, 16, 4, 3, 5, 5):
with T.block("conv2d_NCHWc"):
n, oc_chunk, oh, ow, oc_block, ic, kh, kw = T.axis.remap("SSSSSRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
T.reads(conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block], data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3], placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block])
T.reads(data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3], placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block])
T.writes(conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block])
T.block_attr({"workload":["conv2d_NCHWc.x86", ["TENSOR", [1, 1, 16, 16, 3], "float32"], ["TENSOR", [2, 1, 5, 5, 3, 4], "float32"], [1, 1], [2, 2, 2, 2], [1, 1], "NCHW3c", "NCHW4c", "float32"]})
with T.init():
Expand Down
4 changes: 2 additions & 2 deletions tests/python/unittest/test_te_create_primfunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ def tir_argmax_idx_val(
for i0, i1 in T.grid(m, n):
with T.block("argmax"):
i, k = T.axis.remap("SR", [i0, i1])
T.reads(argmax_v1[i], val[i, k], argmax_v0[i], idx[i, k])
T.reads(val[i, k], idx[i, k])
T.writes(argmax_v0[i], argmax_v1[i])
with T.init():
argmax_v0[i] = T.int32(-1)
Expand Down Expand Up @@ -442,7 +442,7 @@ def tir_argmax_val_idx(
for i0, i1 in T.grid(m, n):
with T.block("argmax"):
i, k = T.axis.remap("SR", [i0, i1])
T.reads(argmax_v0[i], val[i, k], argmax_v1[i], idx[i, k])
T.reads(val[i, k], idx[i, k])
T.writes(argmax_v0[i], argmax_v1[i])
with T.init():
argmax_v0[i] = T.min_value("float32")
Expand Down
64 changes: 64 additions & 0 deletions tests/python/unittest/test_tir_analysis_get_block_access_region.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,48 @@ def access_in_branch_func() -> None:
B[i] = A[i - 1]


@T.prim_func
def gemm() -> None:
A = T.alloc_buffer([16, 16], "float32")
B = T.alloc_buffer([16, 16], "float32")
C = T.alloc_buffer([16, 16], "float32")
for i, j, k, ii, jj in T.grid(4, 4, 16, 4, 4):
with T.block("update"):
vi = T.axis.S(16, i * 4 + ii)
vj = T.axis.S(16, j * 4 + jj)
vk = T.axis.R(16, k)
T.reads(A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
with T.init():
T.reads([])
T.writes(C[vi, vj])
C[vi, vj] = 0
C[vi, vj] += A[vi, vk] * B[vj, vk]


@T.prim_func
def decomposed_gemm() -> None:
A = T.alloc_buffer([16, 16], "float32")
B = T.alloc_buffer([16, 16], "float32")
C = T.alloc_buffer([16, 16], "float32")
for i, j in T.grid(4, 4):
for ii, jj in T.grid(4, 4):
with T.block("init"):
vi = T.axis.S(16, i * 4 + ii)
vj = T.axis.S(16, j * 4 + jj)
T.reads([])
T.writes(C[vi, vj])
C[vi, vj] = 0
for k, ii, jj in T.grid(16, 4, 4):
with T.block("update"):
vi = T.axis.S(16, i * 4 + ii)
vj = T.axis.S(16, j * 4 + jj)
vk = T.axis.R(16, k)
T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
C[vi, vj] += A[vi, vk] * B[vj, vk]


@T.prim_func
def access_of_padding_pattern() -> None:
X = T.alloc_buffer([28, 28])
Expand Down Expand Up @@ -271,6 +313,26 @@ def do_check_block(block_name):
do_check_block("padding_reverse")


def test_access_of_reduction():
block = gemm.body.block.body.body.body.body.body.body.block
alloc_buffers = gemm.body.block.alloc_buffers
buffer_var_map = {buf.data: buf for buf in alloc_buffers}
ret = tir.analysis.get_block_access_region(block, buffer_var_map)
tvm.ir.assert_structural_equal(block.reads, ret[0])
tvm.ir.assert_structural_equal(block.writes, ret[1])


def test_access_of_decompose_reduction():
init = decomposed_gemm.body.block.body.body.body[0].body.body.block
update = decomposed_gemm.body.block.body.body.body[1].body.body.body.block
alloc_buffers = decomposed_gemm.body.block.alloc_buffers
buffer_var_map = {buf.data: buf for buf in alloc_buffers}
for block in [init, update]:
ret = tir.analysis.get_block_access_region(block, buffer_var_map)
tvm.ir.assert_structural_equal(block.reads, ret[0])
tvm.ir.assert_structural_equal(block.writes, ret[1])


if __name__ == "__main__":
test_block_access_region_detector()
test_opaque_block()
Expand All @@ -279,3 +341,5 @@ def do_check_block(block_name):
test_access_in_if_then_else_func()
test_access_in_branch_func()
test_access_of_padding_pattern()
test_access_of_reduction()
test_access_of_decompose_reduction()
Loading