From d5ef43032a8c02a7f3cd4797ea2c3a617f5918ed Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 15 Mar 2022 19:19:47 -0700 Subject: [PATCH 01/10] upd --- .../analysis/block_access_region_detector.cc | 76 +++++++------- .../unittest/test_tir_schedule_for_kind.py | 99 +++++++++++++++++++ 2 files changed, 140 insertions(+), 35 deletions(-) diff --git a/src/tir/analysis/block_access_region_detector.cc b/src/tir/analysis/block_access_region_detector.cc index 974f6ecd644f..8843b3ee062b 100644 --- a/src/tir/analysis/block_access_region_detector.cc +++ b/src/tir/analysis/block_access_region_detector.cc @@ -41,9 +41,11 @@ class BlockReadWriteDetector : public StmtExprVisitor { : buffer_var_map_(buffer_var_map) {} /*! \brief Return read regions of the block */ - Array CollectReads(); + Array CollectReads( + const std::unordered_set* excluded_buffers = nullptr); /*! \brief Return write regions of the block */ - Array CollectWrites(); + Array CollectWrites( + const std::unordered_set* excluded_buffers = nullptr); /*! * \brief Return opaque buffer regions of the block * \note The buffer accessed by load/store or call with buffer.data will @@ -88,8 +90,10 @@ class BlockReadWriteDetector : public StmtExprVisitor { Buffer buffer, std::vector region); /*! \brief Helper function to collect access regions. */ - Array CollectRegions(const std::vector& buffers, - const std::vector>& regions); + Array CollectRegions( + const std::vector& buffers, + const std::vector>& regions, + const std::unordered_set* excluded_buffers = nullptr); /*! \brief Helper function to convert matched access region to source region. */ std::vector ConvertMatchedRegion(const MatchBufferRegion& match_buffer, @@ -126,12 +130,14 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) { StmtExprVisitor::operator()(stmt); } -Array BlockReadWriteDetector::CollectReads() { - return CollectRegions(read_buffers_, read_regions_); +Array BlockReadWriteDetector::CollectReads( + const std::unordered_set* excluded_buffers) { + return CollectRegions(read_buffers_, read_regions_, excluded_buffers); } -Array BlockReadWriteDetector::CollectWrites() { - return CollectRegions(writes_buffers_, write_regions_); +Array BlockReadWriteDetector::CollectWrites( + const std::unordered_set* excluded_buffers) { + return CollectRegions(writes_buffers_, write_regions_, excluded_buffers); } Array BlockReadWriteDetector::CollectOpaques() { @@ -282,12 +288,15 @@ void BlockReadWriteDetector::Update(std::vector* buffers, } Array BlockReadWriteDetector::CollectRegions( - const std::vector& buffers, - const std::vector>& regions) { + const std::vector& buffers, const std::vector>& regions, + const std::unordered_set* excluded_buffers) { ICHECK_EQ(buffers.size(), regions.size()); Array 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 region; region.reserve(regions[i].size()); ICHECK_EQ(buffers[i]->shape.size(), regions[i].size()); @@ -319,38 +328,35 @@ Array> GetBlockAccessRegion(const Block& block, const Map& buffer_var_map) { BlockReadWriteDetector detector(buffer_var_map); detector(block); - return {detector.CollectReads(), detector.CollectWrites(), detector.CollectOpaques()}; + Array writes = detector.CollectWrites(); + std::unordered_set excluded_buffers; + if (block->init.defined()) { + for (const BufferRegion& write_access : writes) { + excluded_buffers.insert(write_access->buffer.get()); + } + } + Array reads = detector.CollectReads(&excluded_buffers); + Array opaques = detector.CollectOpaques(); + return {reads, writes, opaques}; } Array> GetBlockReadWriteRegion(const Block& block, const Map& buffer_var_map) { - // Step 1. Get all the read/write/opaque accesses in the input block. - Array> access_regions = GetBlockAccessRegion(block, buffer_var_map); - // Step 2. Collect all the buffers that are opaquely accessed. - std::unordered_set 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 new_read_regions; - Array 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 opaques = detector.CollectOpaques(); + std::unordered_set 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 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); - } - return {new_read_regions, new_write_regions}; + Array reads = detector.CollectReads(&excluded_buffers); + return {reads, writes}; } TVM_REGISTER_GLOBAL("tir.analysis.GetBlockAccessRegion").set_body_typed(GetBlockAccessRegion); diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index caecde05b40f..3ec206f1d16e 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -330,6 +330,80 @@ def decomposed_gemm_after_vectorize( C[vi, vj] = local[vi, vj] +@T.prim_func +def nested_block_bind( + A: T.Buffer[(16, 16, 16, 16), "float32"], B: T.Buffer[(16, 16, 16), "float32"] +): + for i, j in T.grid(16, 16): + with T.block("outer"): + vi, vj = T.axis.remap("SS", [i, j]) + for k, l in T.grid(16, 16): + with T.block("inner"): + vk, vl = T.axis.remap("SR", [k, l]) + with T.init(): + B[vi, vj, vk] = 0.0 + B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl] + + +@T.prim_func +def thread_bound_nested_block( + A: T.Buffer[(16, 16, 16, 16), "float32"], B: T.Buffer[(16, 16, 16), "float32"] +) -> None: + for i in T.serial(16): + for j in T.thread_binding(16, thread="blockIdx.x"): + with T.block("outer"): + vi, vj = T.axis.remap("SS", [i, j]) + for k in T.serial(16): + for l in T.thread_binding(16, thread="threadIdx.x"): + with T.block("inner"): + vk, vl = T.axis.remap("SR", [k, l]) + with T.init(): + B[vi, vj, vk] = T.float32(0) + B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl] + + +@T.prim_func +def nested_block_bind_after_cache_read( + A: T.Buffer[(16, 16), "float32"], B: T.Buffer[(16,), "float32"] +) -> None: + for i in T.serial(16): + with T.block("outer"): + vi = T.axis.spatial(16, i) + A_shared = T.alloc_buffer([1, 16], dtype="float32", scope="shared") + for ax0, ax1 in T.grid(1, 16): + with T.block("A_shared"): + v0 = T.axis.spatial(16, vi + ax0) + v1 = T.axis.spatial(16, ax1) + A_shared[v0, v1] = A[v0, v1] + for j in T.serial(16): + with T.block("inner"): + vj = T.axis.reduce(16, j) + with T.init(): + B[vi] = T.float32(0) + B[vi] = B[vi] + A_shared[vi, vj] + + +@T.prim_func +def thread_bound_nested_block_after_cache_read( + A: T.Buffer[(16, 16), "float32"], B: T.Buffer[(16,), "float32"] +) -> None: + for i in T.thread_binding(16, thread="blockIdx.x"): + with T.block("outer"): + vi = T.axis.spatial(16, i) + A_shared = T.alloc_buffer([1, 16], dtype="float32", scope="shared") + for ax0, ax1 in T.grid(1, 16): + with T.block("A_shared"): + v0 = T.axis.spatial(16, vi + ax0) + v1 = T.axis.spatial(16, ax1) + A_shared[v0, v1] = A[v0, v1] + for j in T.thread_binding(16, thread="threadIdx.x"): + with T.block("inner"): + vj = T.axis.reduce(16, j) + with T.init(): + B[vi] = T.float32(0) + B[vi] = B[vi] + A_shared[vi, vj] + + # pylint: enable=no-member,invalid-name,unused-variable @@ -468,5 +542,30 @@ def test_vectorize_after_decompose(): verify_trace_roundtrip(s, mod=decomposed_gemm) +def test_nested_block_bind(): + s = tir.Schedule(nested_block_bind) + block_outer = s.get_block("outer") + block_inner = s.get_block("inner") + _, j = s.get_loops(block_outer) + _, l = s.get_loops(block_inner) + s.bind(l, "threadIdx.x") + s.bind(j, "blockIdx.x") + tvm.ir.assert_structural_equal(s.mod["main"], thread_bound_nested_block) + verify_trace_roundtrip(s, mod=nested_block_bind) + + +def test_nexted_block_bind_after_cache_read(): + s = tir.Schedule(nested_block_bind_after_cache_read) + block_outer = s.get_block("outer") + block_inner = s.get_block("inner") + (i,) = s.get_loops(block_outer) + (j,) = s.get_loops(block_inner) + s.bind(i, "blockIdx.x") + s.bind(j, "threadIdx.x") + print(s.mod["main"].script()) + tvm.ir.assert_structural_equal(s.mod["main"], thread_bound_nested_block_after_cache_read) + verify_trace_roundtrip(s, mod=nested_block_bind_after_cache_read) + + if __name__ == "__main__": sys.exit(pytest.main([__file__] + sys.argv[1:])) From 7fe6228a0eed548019f8de7ee5871391c2e56302 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 15 Mar 2022 23:36:54 -0700 Subject: [PATCH 02/10] fix --- src/tir/analysis/block_access_region_detector.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/tir/analysis/block_access_region_detector.cc b/src/tir/analysis/block_access_region_detector.cc index 8843b3ee062b..b6cfdd99eef5 100644 --- a/src/tir/analysis/block_access_region_detector.cc +++ b/src/tir/analysis/block_access_region_detector.cc @@ -356,6 +356,10 @@ Array> GetBlockReadWriteRegion(const Block& block, } } Array reads = detector.CollectReads(&excluded_buffers); + for (const BufferRegion& opaque_access : opaques) { + reads.push_back(opaque_access); + writes.push_back(opaque_access); + } return {reads, writes}; } From d826300bffb4237d8effb95e039d521efad53334 Mon Sep 17 00:00:00 2001 From: Zihao Date: Wed, 16 Mar 2022 01:25:50 -0700 Subject: [PATCH 03/10] fix rfactor --- src/tir/schedule/primitive/reduction.cc | 3 +-- tests/python/unittest/test_tir_schedule_rfactor.py | 4 ++-- tests/python/unittest/test_tvmscript_complete.py | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/tir/schedule/primitive/reduction.cc b/src/tir/schedule/primitive/reduction.cc index 4b9b78e3b299..a02e5fb70f15 100644 --- a/src/tir/schedule/primitive/reduction.cc +++ b/src/tir/schedule/primitive/reduction.cc @@ -826,9 +826,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) { diff --git a/tests/python/unittest/test_tir_schedule_rfactor.py b/tests/python/unittest/test_tir_schedule_rfactor.py index dc60fcd9b303..b2885404c51e 100644 --- a/tests/python/unittest/test_tir_schedule_rfactor.py +++ b/tests/python/unittest/test_tir_schedule_rfactor.py @@ -37,7 +37,7 @@ def transformed_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: with T.block("update"): vi, vj = T.axis.remap("SS", [i0, i1]) vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner) - T.reads([C[vi, vj], A[vi, vk], B[vj, vk]]) + T.reads([A[vi, vk], B[vj, vk]]) T.writes([C[vi, vj]]) with T.init(): C[vi, vj] = 0.0 @@ -172,7 +172,7 @@ def transformed_square_sum_square_root(a: T.handle, d: T.handle) -> None: b = T.axis.S(16, i0) i = T.axis.R(256, T.floordiv(i1_i2_fused_outer, 256)) j = T.axis.R(256, T.floormod(i1_i2_fused_outer, 256)) - T.reads([C[b], A[b, i, j]]) + T.reads([A[b, i, j]]) T.writes([C[b]]) with T.init(): C[b] = 0.0 diff --git a/tests/python/unittest/test_tvmscript_complete.py b/tests/python/unittest/test_tvmscript_complete.py index 429f54809929..17e6d94e6744 100644 --- a/tests/python/unittest/test_tvmscript_complete.py +++ b/tests/python/unittest/test_tvmscript_complete.py @@ -117,7 +117,7 @@ def test_complete_matmul(): access_A = tvm.tir.BufferRegion(A, [Range.from_min_extent(vi, 1), Range.from_min_extent(vk, 1)]) access_B = tvm.tir.BufferRegion(B, [Range.from_min_extent(vj, 1), Range.from_min_extent(vk, 1)]) access_C = tvm.tir.BufferRegion(C, [Range.from_min_extent(vi, 1), Range.from_min_extent(vj, 1)]) - tvm.ir.assert_structural_equal(block.reads, [access_C, access_A, access_B]) + tvm.ir.assert_structural_equal(block.reads, [access_A, access_B]) tvm.ir.assert_structural_equal(block.writes, [access_C]) From 26843649900e861e9cd719f913e19192a4e29b0a Mon Sep 17 00:00:00 2001 From: Zihao Date: Wed, 16 Mar 2022 01:35:53 -0700 Subject: [PATCH 04/10] fix blockize --- src/tir/analysis/block_access_region_detector.cc | 9 ++++++++- tests/python/unittest/test_tir_schedule_blockize.py | 1 + 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/tir/analysis/block_access_region_detector.cc b/src/tir/analysis/block_access_region_detector.cc index b6cfdd99eef5..69519ee85ba1 100644 --- a/src/tir/analysis/block_access_region_detector.cc +++ b/src/tir/analysis/block_access_region_detector.cc @@ -330,7 +330,14 @@ Array> GetBlockAccessRegion(const Block& block, detector(block); Array writes = detector.CollectWrites(); std::unordered_set excluded_buffers; - if (block->init.defined()) { + // exclude write buffers from read regions for reductions. + bool has_reduction_var = false; + for (const IterVar var: block->iter_vars) { + if (var->iter_type == kCommReduce) { + has_reduction_var = true; + } + } + if (has_reduction_var || block->init.defined()) { for (const BufferRegion& write_access : writes) { excluded_buffers.insert(write_access->buffer.get()); } diff --git a/tests/python/unittest/test_tir_schedule_blockize.py b/tests/python/unittest/test_tir_schedule_blockize.py index b4a16a8231b8..bc8585d78e11 100644 --- a/tests/python/unittest/test_tir_schedule_blockize.py +++ b/tests/python/unittest/test_tir_schedule_blockize.py @@ -202,6 +202,7 @@ def test_blockize_init_loops(): s = tir.Schedule(rowsum, debug_mask="all") k, _ = s.get_loops(s.get_block("B")) s.blockize(k) + print(s.mod["main"].script(), rowsum_blockized.script()) tvm.ir.assert_structural_equal(s.mod["main"], rowsum_blockized) verify_trace_roundtrip(sch=s, mod=rowsum) From 8805c1e669f343e103da8ec142df0ca02550376c Mon Sep 17 00:00:00 2001 From: Zihao Date: Wed, 16 Mar 2022 01:43:01 -0700 Subject: [PATCH 05/10] tensorize --- .../unittest/test_tir_schedule_tensorize.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tests/python/unittest/test_tir_schedule_tensorize.py b/tests/python/unittest/test_tir_schedule_tensorize.py index 5cef8d63587d..31e45ce3ede7 100644 --- a/tests/python/unittest/test_tir_schedule_tensorize.py +++ b/tests/python/unittest/test_tir_schedule_tensorize.py @@ -33,7 +33,7 @@ def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) for i, j, k in T.grid(16, 16, 16): with T.block("update"): @@ -48,7 +48,7 @@ def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) T.evaluate( T.tvm_mma_sync( @@ -72,7 +72,7 @@ def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, ()) with T.block("root"): - T.reads(C[()], A[0 : 4], B[0 : 4]) + T.reads(A[0 : 4], B[0 : 4]) T.writes(C[()]) for i in range(0, 4): with T.block("update"): @@ -87,7 +87,7 @@ def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (), offset_factor=1) with T.block("root"): - T.reads(C[()], A[0 : 4], B[0 : 4]) + T.reads(A[0 : 4], B[0 : 4]) T.writes(C[()]) T.evaluate( T.call_extern( @@ -111,7 +111,6 @@ def outer_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None: with T.block("root"): T.reads( - C[0 : 16, 0 : 16], A[0 : 16, 0 : 1], B[0 : 16, 0 : 1], ) @@ -130,7 +129,6 @@ def outer_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: with T.block("root"): T.reads( - C[0 : 16, 0 : 16], A[0 : 16, 0 : 1], B[0 : 16, 0 : 1], ) @@ -180,7 +178,6 @@ def tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer]) T.reads( [ - C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ] @@ -253,7 +250,6 @@ def tensorized_batch_matmul_mma( with T.block("update"): vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k]) T.reads( - C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ) @@ -307,7 +303,7 @@ def tensorized_batch_matmul_dot_product( with T.block("blockized_update"): vn, vi, vj, vko = T.axis.remap("SSSR", [n, i, j, k_0]) T.reads( - C[vn, vi, vj], A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4] + A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4] ) T.writes(C[vn, vi, vj]) A_1 = T.match_buffer( @@ -347,7 +343,6 @@ def tensorized_batch_matmul_outer_product( with T.block("blockized_update"): vn, vio, vjo, vk = T.axis.remap("SSSR", [n, i_0, j_0, k]) T.reads( - C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16], A[vn, vio * 16 : vio * 16 + 16, vk], B[vn, vjo * 16 : vjo * 16 + 16, vk], ) @@ -372,7 +367,7 @@ def annotated_mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) for i, j, k in T.grid(16, 16, 16): with T.block("update"): @@ -414,7 +409,6 @@ def annotated_tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer]) T.reads( [ - C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ] From e1a7c506a287236e43c7146d85131dbb37d439b9 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 22 Mar 2022 00:30:28 -0700 Subject: [PATCH 06/10] upd --- .../analysis/block_access_region_detector.cc | 10 ++-------- .../schedule/primitive/blockize_tensorize.cc | 20 +++++++++++++++++++ src/tir/schedule/primitive/reduction.cc | 17 +++++++++++++++- .../unittest/test_tir_schedule_blockize.py | 1 - .../unittest/test_tir_schedule_for_kind.py | 1 - .../unittest/test_tir_schedule_tensorize.py | 18 +++++++++++------ 6 files changed, 50 insertions(+), 17 deletions(-) diff --git a/src/tir/analysis/block_access_region_detector.cc b/src/tir/analysis/block_access_region_detector.cc index 69519ee85ba1..ffe0c7529400 100644 --- a/src/tir/analysis/block_access_region_detector.cc +++ b/src/tir/analysis/block_access_region_detector.cc @@ -330,14 +330,8 @@ Array> GetBlockAccessRegion(const Block& block, detector(block); Array writes = detector.CollectWrites(); std::unordered_set excluded_buffers; - // exclude write buffers from read regions for reductions. - bool has_reduction_var = false; - for (const IterVar var: block->iter_vars) { - if (var->iter_type == kCommReduce) { - has_reduction_var = true; - } - } - if (has_reduction_var || block->init.defined()) { + // 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()); } diff --git a/src/tir/schedule/primitive/blockize_tensorize.cc b/src/tir/schedule/primitive/blockize_tensorize.cc index 2cecbf1ba2ae..6daea391b918 100644 --- a/src/tir/schedule/primitive/blockize_tensorize.cc +++ b/src/tir/schedule/primitive/blockize_tensorize.cc @@ -496,6 +496,12 @@ StmtSRef Blockize(ScheduleState self, const StmtSRef& loop_sref) { Block new_block = Downcast(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; @@ -503,6 +509,20 @@ StmtSRef Blockize(ScheduleState self, const StmtSRef& loop_sref) { 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 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. diff --git a/src/tir/schedule/primitive/reduction.cc b/src/tir/schedule/primitive/reduction.cc index a02e5fb70f15..4baff106096c 100644 --- a/src/tir/schedule/primitive/reduction.cc +++ b/src/tir/schedule/primitive/reduction.cc @@ -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 new_reads; + std::unordered_set 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(block))->reads = std::move(new_reads); + // 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}; diff --git a/tests/python/unittest/test_tir_schedule_blockize.py b/tests/python/unittest/test_tir_schedule_blockize.py index bc8585d78e11..b4a16a8231b8 100644 --- a/tests/python/unittest/test_tir_schedule_blockize.py +++ b/tests/python/unittest/test_tir_schedule_blockize.py @@ -202,7 +202,6 @@ def test_blockize_init_loops(): s = tir.Schedule(rowsum, debug_mask="all") k, _ = s.get_loops(s.get_block("B")) s.blockize(k) - print(s.mod["main"].script(), rowsum_blockized.script()) tvm.ir.assert_structural_equal(s.mod["main"], rowsum_blockized) verify_trace_roundtrip(sch=s, mod=rowsum) diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index 3ec206f1d16e..e75d5583deea 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -562,7 +562,6 @@ def test_nexted_block_bind_after_cache_read(): (j,) = s.get_loops(block_inner) s.bind(i, "blockIdx.x") s.bind(j, "threadIdx.x") - print(s.mod["main"].script()) tvm.ir.assert_structural_equal(s.mod["main"], thread_bound_nested_block_after_cache_read) verify_trace_roundtrip(s, mod=nested_block_bind_after_cache_read) diff --git a/tests/python/unittest/test_tir_schedule_tensorize.py b/tests/python/unittest/test_tir_schedule_tensorize.py index 31e45ce3ede7..5cef8d63587d 100644 --- a/tests/python/unittest/test_tir_schedule_tensorize.py +++ b/tests/python/unittest/test_tir_schedule_tensorize.py @@ -33,7 +33,7 @@ def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) for i, j, k in T.grid(16, 16, 16): with T.block("update"): @@ -48,7 +48,7 @@ def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) T.evaluate( T.tvm_mma_sync( @@ -72,7 +72,7 @@ def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, ()) with T.block("root"): - T.reads(A[0 : 4], B[0 : 4]) + T.reads(C[()], A[0 : 4], B[0 : 4]) T.writes(C[()]) for i in range(0, 4): with T.block("update"): @@ -87,7 +87,7 @@ def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (), offset_factor=1) with T.block("root"): - T.reads(A[0 : 4], B[0 : 4]) + T.reads(C[()], A[0 : 4], B[0 : 4]) T.writes(C[()]) T.evaluate( T.call_extern( @@ -111,6 +111,7 @@ def outer_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None: with T.block("root"): T.reads( + C[0 : 16, 0 : 16], A[0 : 16, 0 : 1], B[0 : 16, 0 : 1], ) @@ -129,6 +130,7 @@ def outer_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: with T.block("root"): T.reads( + C[0 : 16, 0 : 16], A[0 : 16, 0 : 1], B[0 : 16, 0 : 1], ) @@ -178,6 +180,7 @@ def tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer]) T.reads( [ + C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ] @@ -250,6 +253,7 @@ def tensorized_batch_matmul_mma( with T.block("update"): vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k]) T.reads( + C[vn : vn + 1, vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vn : vn + 1, vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vn : vn + 1, vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ) @@ -303,7 +307,7 @@ def tensorized_batch_matmul_dot_product( with T.block("blockized_update"): vn, vi, vj, vko = T.axis.remap("SSSR", [n, i, j, k_0]) T.reads( - A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4] + C[vn, vi, vj], A[vn, vi, vko * 4 : vko * 4 + 4], B[vn, vj, vko * 4 : vko * 4 + 4] ) T.writes(C[vn, vi, vj]) A_1 = T.match_buffer( @@ -343,6 +347,7 @@ def tensorized_batch_matmul_outer_product( with T.block("blockized_update"): vn, vio, vjo, vk = T.axis.remap("SSSR", [n, i_0, j_0, k]) T.reads( + C[vn, vio * 16 : vio * 16 + 16, vjo * 16 : vjo * 16 + 16], A[vn, vio * 16 : vio * 16 + 16, vk], B[vn, vjo * 16 : vjo * 16 + 16, vk], ) @@ -367,7 +372,7 @@ def annotated_mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None: C = T.match_buffer(c, (16, 16), align=128, offset_factor=1) with T.block("root"): - T.reads(A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) + T.reads(C[0 : 16, 0 : 16], A[0 : 16, 0 : 16], B[0 : 16, 0 : 16]) T.writes(C[0 : 16, 0 : 16]) for i, j, k in T.grid(16, 16, 16): with T.block("update"): @@ -409,6 +414,7 @@ def annotated_tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: vi, vj, vk = T.axis.remap("SSR", [i_outer, j_outer, k_outer]) T.reads( [ + C[vi * 16 : vi * 16 + 16, vj * 16 : vj * 16 + 16], A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], B[vj * 16 : vj * 16 + 16, vk * 16 : vk * 16 + 16], ] From 9dc82735876a9c001f588781bf0320b457228ab6 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 22 Mar 2022 02:17:25 -0700 Subject: [PATCH 07/10] fix meta-schedules --- ...st_meta_schedule_postproc_rewrite_cooperative_fetch.py | 6 +++--- ...test_meta_schedule_schedule_rule_multi_level_tiling.py | 8 ++++---- tests/python/unittest/test_meta_schedule_tune_relay.py | 2 +- .../unittest/test_tir_transform_lower_init_block.py | 5 ++++- 4 files changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_postproc_rewrite_cooperative_fetch.py b/tests/python/unittest/test_meta_schedule_postproc_rewrite_cooperative_fetch.py index 38847b6dba4c..e4dff51cf9d4 100644 --- a/tests/python/unittest/test_meta_schedule_postproc_rewrite_cooperative_fetch.py +++ b/tests/python/unittest/test_meta_schedule_postproc_rewrite_cooperative_fetch.py @@ -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) @@ -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) diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py b/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py index 8882ed625bf1..52218e6c2104 100644 --- a/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py +++ b/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py @@ -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)", @@ -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)", diff --git a/tests/python/unittest/test_meta_schedule_tune_relay.py b/tests/python/unittest/test_meta_schedule_tune_relay.py index e065fd048a1e..071734f68dcb 100644 --- a/tests/python/unittest/test_meta_schedule_tune_relay.py +++ b/tests/python/unittest/test_meta_schedule_tune_relay.py @@ -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(): diff --git a/tests/python/unittest/test_tir_transform_lower_init_block.py b/tests/python/unittest/test_tir_transform_lower_init_block.py index a4fd9404eee4..3ada747f6915 100644 --- a/tests/python/unittest/test_tir_transform_lower_init_block.py +++ b/tests/python/unittest/test_tir_transform_lower_init_block.py @@ -48,6 +48,8 @@ def main(a: T.handle, b: T.handle) -> None: for k0 in T.serial(32, 64): with T.block(): i, j, k = T.axis.remap("SRR", [i0, j0, k0]) + T.reads(A[i, j, k]) + T.writes(B[i]) if (j == 0) and (k == 32): B[i] = T.float32(0) B[i] += A[i, j, k] @@ -82,6 +84,8 @@ def main(a: T.handle, b: T.handle) -> None: for k0 in T.serial(32, 64): with T.block(): i, j, k = T.axis.remap("SRR", [i0, j0, k0]) + T.reads(A[i, j, k]) + T.writes(B[i]) BB = T.match_buffer(B[i], ()) AA = T.match_buffer(A[i, 0:64, 0:64], (64, 64)) if (j == 0) and (k == 32): @@ -92,7 +96,6 @@ def main(a: T.handle, b: T.handle) -> None: def test_lower_reduction(): origin_mod = WithInit mod = tvm.tir.transform.LowerInitBlock()(origin_mod) - print(mod.script()) tvm.ir.assert_structural_equal(mod, WithBranch, True) From c641c30151bed9b74b3323281c593c75036bb8f6 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 22 Mar 2022 02:32:52 -0700 Subject: [PATCH 08/10] upd --- .../unittest/test_tvmscript_roundtrip.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py index 95e5837c5349..148e1c80a5c1 100644 --- a/tests/python/unittest/test_tvmscript_roundtrip.py +++ b/tests/python/unittest/test_tvmscript_roundtrip.py @@ -3205,6 +3205,56 @@ def segment_sum( return segment_sum +def decompose_reduction_read_write(): + @T.prim_func + def decomposed_gemm( + A: T.Buffer[(16, 16), "float32"], + B: T.Buffer[(16, 16), "float32"], + C: T.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] + + return decomposed_gemm + + +def reduction_read_write(): + @T.prim_func + def gemm( + A: T.Buffer[(16, 16), "float32"], + B: T.Buffer[(16, 16), "float32"], + C: T.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] + + return gemm + + ir_generator = tvm.testing.parameter( opt_gemm_normalize, opt_gemm_lower, @@ -3237,6 +3287,8 @@ def segment_sum( func_T_ptr_allocate, llvm_intrin_call, parse_bufferslice_as_range_bound, + reduction_read_write, + decompose_reduction_read_write, ) From d268ba0796aff280d3580af8ca76657c54218da3 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 22 Mar 2022 18:49:06 -0700 Subject: [PATCH 09/10] fix argmax --- tests/python/unittest/test_te_create_primfunc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_te_create_primfunc.py b/tests/python/unittest/test_te_create_primfunc.py index 48082c44a4ab..a65c5d8a0bd8 100644 --- a/tests/python/unittest/test_te_create_primfunc.py +++ b/tests/python/unittest/test_te_create_primfunc.py @@ -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) @@ -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") From 9a59bdce5de966ad44874eda1a58e8716ba72729 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 22 Mar 2022 22:41:39 -0700 Subject: [PATCH 10/10] manally call region analysis --- ...st_tir_analysis_get_block_access_region.py | 64 +++++++++++++++++++ .../unittest/test_tvmscript_roundtrip.py | 52 --------------- 2 files changed, 64 insertions(+), 52 deletions(-) diff --git a/tests/python/unittest/test_tir_analysis_get_block_access_region.py b/tests/python/unittest/test_tir_analysis_get_block_access_region.py index 12e069085b41..f5d701ea7187 100644 --- a/tests/python/unittest/test_tir_analysis_get_block_access_region.py +++ b/tests/python/unittest/test_tir_analysis_get_block_access_region.py @@ -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]) @@ -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() @@ -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() diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py index 148e1c80a5c1..95e5837c5349 100644 --- a/tests/python/unittest/test_tvmscript_roundtrip.py +++ b/tests/python/unittest/test_tvmscript_roundtrip.py @@ -3205,56 +3205,6 @@ def segment_sum( return segment_sum -def decompose_reduction_read_write(): - @T.prim_func - def decomposed_gemm( - A: T.Buffer[(16, 16), "float32"], - B: T.Buffer[(16, 16), "float32"], - C: T.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] - - return decomposed_gemm - - -def reduction_read_write(): - @T.prim_func - def gemm( - A: T.Buffer[(16, 16), "float32"], - B: T.Buffer[(16, 16), "float32"], - C: T.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] - - return gemm - - ir_generator = tvm.testing.parameter( opt_gemm_normalize, opt_gemm_lower, @@ -3287,8 +3237,6 @@ def gemm( func_T_ptr_allocate, llvm_intrin_call, parse_bufferslice_as_range_bound, - reduction_read_write, - decompose_reduction_read_write, )