From e5c27e1f84177f58a2ac7e7d3c1a6e3e67b5e136 Mon Sep 17 00:00:00 2001 From: Zihao Date: Mon, 28 Feb 2022 17:19:18 -0800 Subject: [PATCH 1/4] upd upd sanity re-trigger CI fix reorg upd docstring upd --- src/tir/schedule/analysis/analysis.cc | 83 +++++++++--- .../unittest/test_tir_schedule_for_kind.py | 120 ++++++++++++++++++ 2 files changed, 184 insertions(+), 19 deletions(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index c7ed67187793..e510bdaacaf4 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -167,6 +167,24 @@ bool IsDominantBlock(const BlockScope& scope, const StmtSRef& block_sref) { return true; } +bool ContainsOnlyDataParBlockIter(const Array& iters) { + for (const IterVar& iter_var : iters) { + if (iter_var->iter_type != kDataPar) { + return false; + } + } + return true; +} + +bool ContainsOnlyDataParAndReductionBlockIter(const Array& iters) { + for (const IterVar& iter_var : iters) { + if (iter_var->iter_type != kDataPar && iter_var->iter_type != kCommReduce) { + return false; + } + } + return true; +} + /*! * \brief A helper function that checks whether a given block is a complete block under the scope, * or return the condition it violates if it is not a complete block @@ -181,10 +199,8 @@ int CheckCompleteBlockErrorCode(const ScheduleState& self, const StmtSRef& block BlockScope scope = self->GetBlockScope(scope_root_sref); // Cond 1. All block vars are data parallel const BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); - for (const IterVar& iter_var : block->iter_vars) { - if (iter_var->iter_type != kDataPar) { - return 1; - } + if (!ContainsOnlyDataParBlockIter(block->iter_vars)) { + return 1; } // Cond 2. Dominant: the block is the only writer of its output, // dominating the reader of its output buffers @@ -211,7 +227,11 @@ static const char* kCompleteBlockDefinition = R"(Definition of a complete block: 3) No overlap between the buffers the block reads and writes)"; static const char* kReductionBlockDefinition = R"(Definition of a reduction block: -1) The block has the `init` statement +1) The block has the `init` statement, or + - all block iter vars are data-parallel + - there are sub-blocks + - each sub-block is complete/reduction, + - there is a least a reduction sub-block. 2) All the block bindings are quasi-affine expressions 3) All block vars are either data parallel block vars or reduction block vars 4) Dominant: the block is the only writer of its output, dominating the reader of its output buffers @@ -262,10 +282,6 @@ int CheckReductionBlockErrorCode(const ScheduleState& self, const StmtSRef& bloc const StmtSRef& scope_root_sref) { BlockScope scope = self->GetBlockScope(scope_root_sref); const BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); - // Cond 1. The block has the `init` statement. - if (!block->init.defined()) { - return 1; - } // Cond 2. All the block bindings are quasi-affine expressions. if (!self->IsAffineBlockBinding(block_sref)) { return 2; @@ -281,7 +297,45 @@ int CheckReductionBlockErrorCode(const ScheduleState& self, const StmtSRef& bloc return 4; } // Cond 5. The reduction block vars are not used to index the output buffers. - return ReductionIterNotIndexOutputBuffer(GetRef(block)) ? 0 : 5; + if (!ReductionIterNotIndexOutputBuffer(GetRef(block))) { + return 5; + } + // Cond 1. The block has the `init` statement, or + // - all block iter vars are data-parallel, + // - there are sub-blocks, + // - each sub-block is complete/reduction + // - there is a least a reduction sub-block. + if (!block->init.defined()) { + Array child_block_srefs = GetChildBlockSRefOnSRefTree(self, block_sref); + if (child_block_srefs.empty()) { + return 1; + } else { // have sub-blocks. + // all block iter vars are data-parallel + if (!ContainsOnlyDataParBlockIter(block->iter_vars)) { + return 1; + } + // all sub-blocks are complete + bool has_reduction = false; + bool all_complete_reduction = true; + for (const StmtSRef& child_block_sref : child_block_srefs) { + int complete_code = CheckCompleteBlockErrorCode(self, child_block_sref, scope_root_sref); + int reduction_code = CheckReductionBlockErrorCode(self, child_block_sref, scope_root_sref); + if (complete_code != 0 && reduction_code != 0) { + all_complete_reduction = false; + break; + } + if (reduction_code == 0) { + has_reduction = true; + } + } + if (has_reduction && all_complete_reduction) { + return 0; + } else { + return 1; + } + } + } + return 0; } bool IsReductionBlock(const ScheduleState& self, const StmtSRef& block_sref, @@ -1417,15 +1471,6 @@ std::pair GetBufferStoresFromReductionBlock( return std::make_pair(GetRef(init), GetRef(body)); } -bool ContainsOnlyDataParAndReductionBlockIter(const Array& iters) { - for (const IterVar& iter_var : iters) { - if (iter_var->iter_type != kDataPar && iter_var->iter_type != kCommReduce) { - return false; - } - } - return true; -} - bool ReductionIterNotIndexOutputBuffer(const Block& block) { // Step 1. Collect the reduction block iters. std::unordered_set reduction_block_iters; diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index caecde05b40f..b80a52d2852c 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring +from os import get_blocking import sys import pytest @@ -22,6 +23,7 @@ import tvm.testing from tvm import tir from tvm.script import tir as T +from tvm.tir.schedule.schedule import Schedule from tvm.tir.schedule.testing import verify_trace_roundtrip # pylint: disable=no-member,invalid-name,unused-variable @@ -330,6 +332,100 @@ 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]) + T.reads(B[vi, vj, 0:16], A[vi, vj, 0:16, 0:16]) + T.writes(B[vi, vj, 0:16]) + for k, l in T.grid(16, 16): + with T.block("inner"): + vk, vl = T.axis.remap("SR", [k, l]) + T.reads(B[vi, vj, vk], A[vi, vj, vk, vl]) + T.writes(B[vi, vj, vk]) + 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]) + T.reads(B[vi, vj, 0:16], A[vi, vj, 0:16, 0:16]) + T.writes(B[vi, vj, 0:16]) + 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]) + T.reads(B[vi, vj, vk], A[vi, vj, vk, vl]) + T.writes(B[vi, vj, vk]) + 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") + T.reads(B[vi], A[vi, 0:16], A_shared[0, 0:16]) + T.writes(B[vi], A_shared[0, 0:16]) + 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) + T.reads(A[v0, v1]) + T.writes(A_shared[v0, v1]) + A_shared[v0, v1] = A[v0, v1] + for j in T.serial(16): + with T.block("inner"): + vj = T.axis.reduce(16, j) + T.reads(B[vi], A_shared[vi, vj]) + T.writes(B[vi]) + 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") + T.reads(B[vi], A[vi, 0:16], A_shared[0, 0:16]) + T.writes(B[vi], A_shared[0, 0:16]) + 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) + T.reads(A[v0, v1]) + T.writes(A_shared[v0, v1]) + 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) + T.reads(B[vi], A_shared[vi, vj]) + T.writes(B[vi]) + 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 +564,29 @@ 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") + 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 a2a406fbdeb07b0302258b24fbc2231227601820 Mon Sep 17 00:00:00 2001 From: Zihao Date: Fri, 11 Mar 2022 18:28:50 -0800 Subject: [PATCH 2/4] lint --- src/tir/schedule/analysis/analysis.cc | 4 +--- tests/python/unittest/test_tir_schedule_for_kind.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index e510bdaacaf4..eb2bb877beb0 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -328,9 +328,7 @@ int CheckReductionBlockErrorCode(const ScheduleState& self, const StmtSRef& bloc has_reduction = true; } } - if (has_reduction && all_complete_reduction) { - return 0; - } else { + if (!has_reduction || !all_complete_reduction) { return 1; } } diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index b80a52d2852c..d9b41d05057c 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -15,7 +15,6 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring,missing-module-docstring -from os import get_blocking import sys import pytest From d64d2338619b3bbeb25a5bc92df68bfc33f91209 Mon Sep 17 00:00:00 2001 From: Zihao Date: Fri, 11 Mar 2022 18:29:49 -0800 Subject: [PATCH 3/4] lint --- tests/python/unittest/test_tir_schedule_for_kind.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index d9b41d05057c..3fbdff22938d 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -22,7 +22,6 @@ import tvm.testing from tvm import tir from tvm.script import tir as T -from tvm.tir.schedule.schedule import Schedule from tvm.tir.schedule.testing import verify_trace_roundtrip # pylint: disable=no-member,invalid-name,unused-variable From 4f4236f7d8cd74554b5ce4989cefcd3846404212 Mon Sep 17 00:00:00 2001 From: Zihao Date: Sat, 12 Mar 2022 01:17:52 -0800 Subject: [PATCH 4/4] upd --- src/tir/schedule/analysis/analysis.cc | 13 ++++++++++--- tests/python/unittest/test_tir_schedule_for_kind.py | 9 +++++---- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index eb2bb877beb0..9a4cf411f68b 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -318,8 +318,8 @@ int CheckReductionBlockErrorCode(const ScheduleState& self, const StmtSRef& bloc bool has_reduction = false; bool all_complete_reduction = true; for (const StmtSRef& child_block_sref : child_block_srefs) { - int complete_code = CheckCompleteBlockErrorCode(self, child_block_sref, scope_root_sref); - int reduction_code = CheckReductionBlockErrorCode(self, child_block_sref, scope_root_sref); + int complete_code = CheckCompleteBlockErrorCode(self, child_block_sref, block_sref); + int reduction_code = CheckReductionBlockErrorCode(self, child_block_sref, block_sref); if (complete_code != 0 && reduction_code != 0) { all_complete_reduction = false; break; @@ -1498,7 +1498,14 @@ bool ReductionIterNotIndexOutputBuffer(const Block& block) { if (!store) { return true; } - ICHECK(buffer_written.count(store->buffer.get())) + // whether the buffer is allocated inside block. + bool is_block_allocated_buf = false; + for (const Buffer& alloc_buf : block->alloc_buffers) { + if (store->buffer == alloc_buf) { + return true; + } + } + ICHECK(buffer_written.count(store->buffer.get()) || is_block_allocated_buf) << "ValueError: The buffer \"" << store->buffer << "\" is written in the block but is not in the block's signature"; for (const PrimExpr& index : store->indices) { diff --git a/tests/python/unittest/test_tir_schedule_for_kind.py b/tests/python/unittest/test_tir_schedule_for_kind.py index 3fbdff22938d..1abd5cdac86c 100644 --- a/tests/python/unittest/test_tir_schedule_for_kind.py +++ b/tests/python/unittest/test_tir_schedule_for_kind.py @@ -377,9 +377,9 @@ def nested_block_bind_after_cache_read( for i in T.serial(16): with T.block("outer"): vi = T.axis.spatial(16, i) + T.reads(B[vi], A[vi, 0:16]) + T.writes(B[vi]) A_shared = T.alloc_buffer([1, 16], dtype="float32", scope="shared") - T.reads(B[vi], A[vi, 0:16], A_shared[0, 0:16]) - T.writes(B[vi], A_shared[0, 0:16]) for ax0, ax1 in T.grid(1, 16): with T.block("A_shared"): v0 = T.axis.spatial(16, vi + ax0) @@ -404,9 +404,9 @@ def thread_bound_nested_block_after_cache_read( for i in T.thread_binding(16, thread="blockIdx.x"): with T.block("outer"): vi = T.axis.spatial(16, i) + T.reads(B[vi], A[vi, 0:16]) + T.writes(B[vi]) A_shared = T.alloc_buffer([1, 16], dtype="float32", scope="shared") - T.reads(B[vi], A[vi, 0:16], A_shared[0, 0:16]) - T.writes(B[vi], A_shared[0, 0:16]) for ax0, ax1 in T.grid(1, 16): with T.block("A_shared"): v0 = T.axis.spatial(16, vi + ax0) @@ -582,6 +582,7 @@ 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)