From 74805a311c160371b476be98736cff3b31919305 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 3 Apr 2023 12:29:50 -0500 Subject: [PATCH 01/10] [TIR] Flatten SeqStmt on constructions Previously, SeqStmt could be nested, making a distinction between the nested `SeqStmt({SeqStmt({a,b}), c})` and the flat `SeqStmt({a,b,c})`, even though the two are semantically equivalent. This also caused an issue with round-trips through TVMScript, which does not preserve this distinction. This commit updates the `SeqStmt` constructor and the `SeqStmt` visitor in `StmtMutator` to flatten nested sequential statements provided. --- src/tir/ir/stmt.cc | 12 +++++++++ src/tir/ir/stmt_functor.cc | 4 +-- .../unittest/test_tvmscript_roundtrip.py | 25 +++++++++++++++++++ 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/src/tir/ir/stmt.cc b/src/tir/ir/stmt.cc index e4569898f7ed..3b07bd943973 100644 --- a/src/tir/ir/stmt.cc +++ b/src/tir/ir/stmt.cc @@ -387,6 +387,18 @@ TVM_REGISTER_NODE_TYPE(PrefetchNode); // SeqStmt SeqStmt::SeqStmt(Array seq, Span span) { + bool requires_flattening = std::any_of( + seq.begin(), seq.end(), [](const Stmt& stmt) { return stmt->IsInstance(); }); + + if (requires_flattening) { + auto flattened = SeqStmt::Flatten(seq); + if (auto* ptr = flattened.as()) { + seq = ptr->seq; + } else { + seq = {flattened}; + } + } + auto node = make_object(); node->seq = std::move(seq); node->span = std::move(span); diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc index f5063b222b9b..7c693b7efcf7 100644 --- a/src/tir/ir/stmt_functor.cc +++ b/src/tir/ir/stmt_functor.cc @@ -439,9 +439,7 @@ Stmt StmtMutator::VisitStmt_(const SeqStmtNode* op) { if (seq.same_as(op->seq)) { return GetRef(op); } else { - auto n = CopyOnWrite(op); - n->seq = std::move(seq); - return Stmt(n); + return SeqStmt(seq); } } diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py index cd7f1726c9d9..f339ec287a66 100644 --- a/tests/python/unittest/test_tvmscript_roundtrip.py +++ b/tests/python/unittest/test_tvmscript_roundtrip.py @@ -3692,6 +3692,30 @@ def func( return func +def nested_seqstmt(): + """Nested SeqStmt should be normalized to flat SeqStmt + + Nested SeqStmt are representable in the TIR structures, but are + flattened when converted to TVMScript. Previously, this could + cause failures to round-trip through TVMScript, including + erroneous use of TVMScript's concise-scoping rules. This was + resolved by normalizing nested SeqStmt in TIR, such that the use + of `tir.SeqStmt` below results in a single flat `tir.SeqStmt` + containing the three `tir.Evaluate` calls. + """ + func = tvm.tir.PrimFunc( + params=[], + body=tvm.tir.SeqStmt( + [ + tvm.tir.SeqStmt([tvm.tir.Evaluate(0), tvm.tir.Evaluate(1)]), + tvm.tir.Evaluate(2), + ] + ), + ) + + return func + + ir_generator = tvm.testing.parameter( launch_env_thread, opt_gemm_normalize, @@ -3757,6 +3781,7 @@ def func( merge_shape_var_def, if_then_else_var, tvm_shfl_builtins, + nested_seqstmt, ) From 3f701a03392e5f77dd1cc0ae5f1eb4392fa7658c Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 4 Apr 2023 15:19:21 -0500 Subject: [PATCH 02/10] Add more aggressive normalization The flattening still can produce ambiguous structure for `SeqStmt` of size 0 and 1. Normalizing so that `SeqStmt::Flatten` will automatically unwrap when flattening produces a single statement, and will convert to `Evaluate(0)` when flattening removes all statements. --- include/tvm/tir/stmt.h | 100 +++++++++++--------- src/relay/backend/aot_executor_codegen.cc | 6 +- src/tir/ir/stmt.cc | 7 ++ src/tir/ir/stmt_functor.cc | 2 +- tests/python/relay/aot/test_c_device_api.py | 4 +- 5 files changed, 69 insertions(+), 50 deletions(-) diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index 9ed9973871d9..50d0ac316a41 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -690,6 +690,46 @@ class SeqStmtNode : public StmtNode { TVM_DECLARE_FINAL_OBJECT_INFO(SeqStmtNode, StmtNode); }; +/*! + * \brief Evaluates an expression. + * This is mostly used for putting a Call node into Stmt. + * + * If value do not have side-effect, this node can be safely removed. + */ +class EvaluateNode : public StmtNode { + public: + /*! \brief The expression to be evaluated. */ + PrimExpr value; + + void VisitAttrs(AttrVisitor* v) { + v->Visit("value", &value); + v->Visit("span", &span); + } + + bool SEqualReduce(const EvaluateNode* other, SEqualReducer equal) const { + return equal(value, other->value); + } + + void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(value); } + + static constexpr const char* _type_key = "tir.Evaluate"; + TVM_DECLARE_FINAL_OBJECT_INFO(EvaluateNode, StmtNode); +}; + +/*! + * \brief Managed reference to EvaluateNode. + * \sa EvaluateNode + */ +class Evaluate : public Stmt { + public: + TVM_DLL explicit Evaluate(PrimExpr value, Span span = Span()); + + explicit Evaluate(int value, Span span = Span()) : Evaluate(PrimExpr(value), span) {} + + TVM_DEFINE_OBJECT_REF_METHODS(Evaluate, Stmt, EvaluateNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(EvaluateNode); +}; + /*! \brief Sequence statement. */ class SeqStmt : public Stmt { public: @@ -726,8 +766,13 @@ class SeqStmt : public Stmt { static Stmt Flatten(Args&&... seq_args) { Array seq; runtime::detail::for_each(Flattener(&seq), std::forward(seq_args)...); - if (seq.size() == 1) return seq[0]; - return SeqStmt(seq); + if (seq.empty()) { + return Evaluate(0); + } else if (seq.size() == 1) { + return seq[0]; + } else { + return SeqStmt(seq); + } } /*! \brief Helper class to flatten sequence of arguments into Array. */ class Flattener { @@ -738,9 +783,16 @@ class SeqStmt : public Stmt { if (!stmt.defined()) return; if (auto* op = stmt.as()) { operator()(0, op->seq); - } else { - seq_->push_back(stmt); + return; + } + + if (auto* op = stmt.as()) { + if (auto* as_int = op->value.as(); as_int && as_int->value == 0) { + return; + } } + + seq_->push_back(stmt); } template @@ -805,46 +857,6 @@ class IfThenElse : public Stmt { TVM_DEFINE_OBJECT_REF_COW_METHOD(IfThenElseNode); }; -/*! - * \brief Evaluates an expression. - * This is mostly used for putting a Call node into Stmt. - * - * If value do not have side-effect, this node can be safely removed. - */ -class EvaluateNode : public StmtNode { - public: - /*! \brief The expression to be evaluated. */ - PrimExpr value; - - void VisitAttrs(AttrVisitor* v) { - v->Visit("value", &value); - v->Visit("span", &span); - } - - bool SEqualReduce(const EvaluateNode* other, SEqualReducer equal) const { - return equal(value, other->value); - } - - void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(value); } - - static constexpr const char* _type_key = "tir.Evaluate"; - TVM_DECLARE_FINAL_OBJECT_INFO(EvaluateNode, StmtNode); -}; - -/*! - * \brief Managed reference to EvaluateNode. - * \sa EvaluateNode - */ -class Evaluate : public Stmt { - public: - TVM_DLL explicit Evaluate(PrimExpr value, Span span = Span()); - - explicit Evaluate(int value, Span span = Span()) : Evaluate(PrimExpr(value), span) {} - - TVM_DEFINE_OBJECT_REF_METHODS(Evaluate, Stmt, EvaluateNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(EvaluateNode); -}; - /*! * \brief The kind of the loop. * diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc index 8f7098c24aea..37f317ccc4a8 100644 --- a/src/relay/backend/aot_executor_codegen.cc +++ b/src/relay/backend/aot_executor_codegen.cc @@ -494,7 +494,7 @@ class AOTExecutorCodegen : public MixedModeVisitor { })); } - tir::Stmt body = tir::SeqStmt({func_call}); + tir::Stmt body = tir::SeqStmt::Flatten(func_call); stmts_.push_back(body); } @@ -570,7 +570,7 @@ class AOTExecutorCodegen : public MixedModeVisitor { {tvm::tir::StringImm(device_hook_name), context}))); device_hooks.push_back(device_hook); } - return tir::SeqStmt(device_hooks); + return tir::SeqStmt::Flatten(device_hooks); } /** @@ -736,7 +736,7 @@ class AOTExecutorCodegen : public MixedModeVisitor { // the packed function calls don't pack their arguments. The AOT // runner function needs to be legalized by the LegalizePackedCalls pass. tir::PrimFunc CreateMainFunc(String mod_name, unsigned int relay_params) { - tir::Stmt body = tir::SeqStmt(stmts_); + tir::Stmt body = tir::SeqStmt::Flatten(stmts_); // Allocate the sids std::unordered_map allocated; diff --git a/src/tir/ir/stmt.cc b/src/tir/ir/stmt.cc index 3b07bd943973..b32b9b6c4584 100644 --- a/src/tir/ir/stmt.cc +++ b/src/tir/ir/stmt.cc @@ -399,6 +399,13 @@ SeqStmt::SeqStmt(Array seq, Span span) { } } + ICHECK_NE(seq.size(), 0) << "An empty SeqStmt is prohibited. " + << "To write a no-op, use Evaluate(0), " + << "or the result of SeqStmt::Flatten()"; + ICHECK_NE(seq.size(), 1) << "A SeqStmt of length 1 is prohibited. " + << "Use the node " << seq[0] << "directly, " + << "or for dynamic usage, normalize using SeqStmt::Flatten()"; + auto node = make_object(); node->seq = std::move(seq); node->span = std::move(span); diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc index 7c693b7efcf7..8f1bf0a6e2a8 100644 --- a/src/tir/ir/stmt_functor.cc +++ b/src/tir/ir/stmt_functor.cc @@ -439,7 +439,7 @@ Stmt StmtMutator::VisitStmt_(const SeqStmtNode* op) { if (seq.same_as(op->seq)) { return GetRef(op); } else { - return SeqStmt(seq); + return SeqStmt::Flatten(seq); } } diff --git a/tests/python/relay/aot/test_c_device_api.py b/tests/python/relay/aot/test_c_device_api.py index 247b22eac494..e317d059f053 100644 --- a/tests/python/relay/aot/test_c_device_api.py +++ b/tests/python/relay/aot/test_c_device_api.py @@ -238,7 +238,7 @@ def test_without_device_api_unpacked_api(non_device_api_main_func): """Test a graph without the Device API with the unpacked internal calls""" main_func = non_device_api_main_func(interface_api="c", use_unpacked_api=True) - body = main_func.body.seq[1].seq[0].seq[0].value + body = main_func.body.value assert ( repr(body) == 'T.tvm_check_return(0, -1, T.call_extern("int32", ' @@ -252,7 +252,7 @@ def test_without_device_api_packed_api(non_device_api_main_func): main_func = non_device_api_main_func(interface_api="packed", use_unpacked_api=False) - body = main_func.body.seq[1].seq[0].seq[0].value + body = main_func.body.value assert repr(body) == ( 'T.call_cpacked("tvmgen_default_fused_multiply", ' "T.tvm_stack_make_array(x_buffer_var, T.tvm_stack_make_shape(10, 10), " From 01db87fde12d905e4e8151f81c605c9e99c9819f Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Wed, 5 Apr 2023 17:09:33 -0500 Subject: [PATCH 03/10] Use SeqStmt::Flatten during TVMScript parsing --- src/script/ir_builder/tir/utils.h | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/script/ir_builder/tir/utils.h b/src/script/ir_builder/tir/utils.h index 7ccc132fa1fe..69db8f905f8f 100644 --- a/src/script/ir_builder/tir/utils.h +++ b/src/script/ir_builder/tir/utils.h @@ -51,14 +51,7 @@ inline void AddToParent(tvm::tir::Stmt stmt) { * \return The SeqStmt. */ inline tvm::tir::Stmt AsStmt(const Array& stmt) { - using namespace tvm::tir; - if (stmt.empty()) { - return tvm::tir::Evaluate(0); - } else if (stmt.size() == 1) { - return stmt[0]; - } else { - return SeqStmt(stmt); - } + return tvm::tir::SeqStmt::Flatten(stmt); } /*! From 35ab173489f43dcf496ad8236f1e4ca0b0ab2fbd Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 6 Apr 2023 09:55:42 -0500 Subject: [PATCH 04/10] Maintain SeqStmt identity through SeqStmt::Flatten if possible If neither flattening nor unwrapping is required, the returned SeqStmt should be the same object as the argument. --- include/tvm/tir/stmt.h | 80 +++++++++++++++++++++++------- src/tir/ir/stmt_functor.cc | 7 +-- src/tir/transforms/remove_no_op.cc | 25 ---------- 3 files changed, 63 insertions(+), 49 deletions(-) diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index 4d2a0d0c3229..7e40b329d290 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -758,6 +758,10 @@ class SeqStmt : public Stmt { * \note This function can directly return an element * if it is the only element in the sequence. * + * \note If the only argument to this function is a SeqStmt, and if + * no flattening of the SeqStmt is required, then the SeqStmt + * will be returned as-is. + * * \param seq_args The list of arguments to be flattened. * \tparam Args arguments * \return The constructed statement @@ -766,55 +770,93 @@ class SeqStmt : public Stmt { static Stmt Flatten(Args&&... seq_args) { Array seq; runtime::detail::for_each(Flattener(&seq), std::forward(seq_args)...); + if (seq.empty()) { return Evaluate(0); } else if (seq.size() == 1) { return seq[0]; - } else { - return SeqStmt(seq); } + + // If the argument is a single SeqStmt argument with no + // flattening or unwrapping required required, then we may + // return the SeqStmt as-is. + if constexpr (sizeof...(seq_args) == 1) { + if (auto opt = Flattener::AsSeqStmt(std::forward(seq_args)...)) { + SeqStmt original = opt.value(); + bool all_same = [&]() { + if (original->seq.size() != seq.size()) { + return false; + } + for (size_t i = 0; i < seq.size(); i++) { + if (!original->seq[i].same_as(seq[i])) { + return false; + } + } + return true; + }(); + if (all_same) { + return original; + } + } + } + + return SeqStmt(seq); } /*! \brief Helper class to flatten sequence of arguments into Array. */ class Flattener { public: explicit Flattener(Array* seq) : seq_(seq) {} + template + static Optional AsSeqStmt(const T& t) { + if constexpr (std::is_same_v) { + return t; + } else if constexpr (!std::is_base_of_v) { + return NullOpt; + } else if (auto* ptr = t.template as()) { + return GetRef(ptr); + } else { + return NullOpt; + } + } + template void operator()(size_t i, const T& stmt_or_seq) const { if constexpr (std::is_base_of_v) { // Early bail-out, applicable to any ObjectRef - if (!stmt_or_seq.defined()) return; + if (!stmt_or_seq.defined()) { + return; + } } if constexpr (std::is_same_v) { - // No need for dynamic type-checking if the static type is a - // SeqStmt. + // Static type-checking for a SeqStmt that could be flattened. (*this)(0, stmt_or_seq->seq); - } else if constexpr (std::is_base_of_v) { + return; + } + + if constexpr (std::is_base_of_v) { // Dynamic type-checking for a SeqStmt that could be // flattened. if (auto* op = stmt_or_seq.template as()) { operator()(0, op->seq); - } else { - seq_->push_back(stmt_or_seq); + return; } - } else if constexpr (std::is_base_of_v) { + } + + if constexpr (std::is_base_of_v) { // Evaluate(0) is used to represent a no-op, and may be // generated by previous calls to SeqStmt::Flatten(). These // should be removed to ensure that Flatten(a+b) is equivalent // to Flatten(Flatten(a), Flatten(b)). - bool is_canonical_no_op = [&stmt_or_seq]() { - if (auto* op = stmt_or_seq.template as()) { - if (auto* as_int = op->value.template as(); as_int && as_int->value == 0) { - return true; - } + if (auto* op = stmt_or_seq.template as()) { + if (auto* as_int = op->value.template as(); as_int && as_int->value == 0) { + return; } - return false; - }(); - if (!is_canonical_no_op) { - seq_->push_back(stmt_or_seq); } - } else if constexpr (std::is_base_of_v) { + } + + if constexpr (std::is_base_of_v) { // Any other Stmt type just gets appended. seq_->push_back(stmt_or_seq); } else { diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc index 8f1bf0a6e2a8..366808ea4f0e 100644 --- a/src/tir/ir/stmt_functor.cc +++ b/src/tir/ir/stmt_functor.cc @@ -436,11 +436,8 @@ Stmt StmtMutator::VisitStmt_(const PrefetchNode* op) { Stmt StmtMutator::VisitStmt_(const SeqStmtNode* op) { Array seq = Internal::Mutate(this, op->seq); - if (seq.same_as(op->seq)) { - return GetRef(op); - } else { - return SeqStmt::Flatten(seq); - } + Stmt stmt = seq.same_as(op->seq) ? GetRef(op) : SeqStmt(seq); + return SeqStmt::Flatten(seq); } // advanced visit function for seqstmt. diff --git a/src/tir/transforms/remove_no_op.cc b/src/tir/transforms/remove_no_op.cc index d35cf8b8d602..4179b00a3684 100644 --- a/src/tir/transforms/remove_no_op.cc +++ b/src/tir/transforms/remove_no_op.cc @@ -178,31 +178,6 @@ class NoOpRemover : public arith::IRMutatorWithAnalyzer { } } - Stmt VisitStmt_(const SeqStmtNode* op) final { - auto ret = Downcast(StmtMutator::VisitSeqStmt_(op, true)); - - bool need_compact = std::any_of(ret->seq.begin(), ret->seq.end(), - [](const auto& stmt) { return is_no_op(stmt); }); - - if (need_compact) { - Array filtered; - for (Stmt stmt : ret->seq) { - if (!is_no_op(stmt)) { - filtered.push_back(std::move(stmt)); - } - } - ret = SeqStmt(filtered); - } - - if (ret->size() == 0) { - return Evaluate(0); - } else if (ret->size() == 1) { - return ret->seq[0]; - } else { - return std::move(ret); - } - } - Stmt VisitStmt_(const BufferStoreNode* op) final { BufferStore store = GetRef(op); From 6519968ac266127daa67a7d62523f9c14b142e5e Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 6 Apr 2023 10:03:31 -0500 Subject: [PATCH 05/10] Use a different const int for testing RemoveNoOp Otherwise, the SeqStmt removes the "Evaluate(0)" before we can get to the part that we want to test. --- tests/python/unittest/test_tir_transform_remove_no_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/unittest/test_tir_transform_remove_no_op.py b/tests/python/unittest/test_tir_transform_remove_no_op.py index 15c5a577f9f5..133ef01ed001 100644 --- a/tests/python/unittest/test_tir_transform_remove_no_op.py +++ b/tests/python/unittest/test_tir_transform_remove_no_op.py @@ -23,7 +23,7 @@ def nop(): - return tvm.tir.Evaluate(0) + return tvm.tir.Evaluate(1) def test_remove_no_op(): From 743de24c993d447141e83aee5fcc0c84c68ca527 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Fri, 7 Apr 2023 14:32:16 -0500 Subject: [PATCH 06/10] Updated tests to avoid relying on no-op Evaluate(0) remaining --- tests/python/unittest/test_tir_stmt_functor_ir_transform.py | 4 ++-- tests/python/unittest/test_tvmscript_printer_annotation.py | 6 +++--- tests/python/unittest/test_tvmscript_printer_underlining.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/python/unittest/test_tir_stmt_functor_ir_transform.py b/tests/python/unittest/test_tir_stmt_functor_ir_transform.py index c79ba68481f6..a22e5bf466ee 100644 --- a/tests/python/unittest/test_tir_stmt_functor_ir_transform.py +++ b/tests/python/unittest/test_tir_stmt_functor_ir_transform.py @@ -31,7 +31,7 @@ def test_ir_transform(): def preorder(op): if op.op.same_as(builtin_call_extern) and op.args[0].value == "TestC": - return tvm.tir.const(0, "int32") + return tvm.tir.const(42, "int32") return None def postorder(op): @@ -43,7 +43,7 @@ def postorder(op): body = tvm.tir.stmt_functor.ir_transform(body, preorder, postorder, ["tir.Call"]) stmt_list = tvm.tir.stmt_list(body.body.body) assert stmt_list[0].value.args[1].args[0].value == "TestB" - assert stmt_list[1].value.value == 0 + assert stmt_list[1].value.value == 42 if __name__ == "__main__": diff --git a/tests/python/unittest/test_tvmscript_printer_annotation.py b/tests/python/unittest/test_tvmscript_printer_annotation.py index 70d5b655fb37..72d2238b2b63 100644 --- a/tests/python/unittest/test_tvmscript_printer_annotation.py +++ b/tests/python/unittest/test_tvmscript_printer_annotation.py @@ -24,7 +24,7 @@ @T.prim_func def _func(): - T.evaluate(0) + T.evaluate(-1) T.evaluate(1) T.evaluate(2) T.evaluate(3) @@ -49,7 +49,7 @@ def test_annotation_multi_object_paths(): @T.prim_func def main(): - T.evaluate(0) + T.evaluate(-1) T.evaluate(1) # annotation 1 T.evaluate(2) T.evaluate(3) # annotation 3 @@ -75,7 +75,7 @@ def test_annotate_from_multi_obj(): @T.prim_func def main(): - T.evaluate(0) + T.evaluate(-1) T.evaluate(1) # annotation 1 T.evaluate(2) T.evaluate(3) # annotation 3 diff --git a/tests/python/unittest/test_tvmscript_printer_underlining.py b/tests/python/unittest/test_tvmscript_printer_underlining.py index 4a4d17d0d89b..569f03d0f828 100644 --- a/tests/python/unittest/test_tvmscript_printer_underlining.py +++ b/tests/python/unittest/test_tvmscript_printer_underlining.py @@ -433,7 +433,7 @@ def main(a: T.int32, b: T.int32): def test_underline_from_multi_obj(): @T.prim_func def func(): - T.evaluate(0) + T.evaluate(-1) T.evaluate(1) T.evaluate(2) T.evaluate(3) @@ -456,7 +456,7 @@ def func(): @T.prim_func def main(): - T.evaluate(0) + T.evaluate(-1) T.evaluate(1) ^^^^^^^^^^^^^ T.evaluate(2) From d4ed8634e80f87fd67b3c93c41918af273894348 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 10 Apr 2023 09:17:36 -0500 Subject: [PATCH 07/10] Maintain CopyOnWrite handling in StmtMutator --- src/tir/ir/stmt_functor.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc index 366808ea4f0e..945d535a4aa2 100644 --- a/src/tir/ir/stmt_functor.cc +++ b/src/tir/ir/stmt_functor.cc @@ -436,8 +436,13 @@ Stmt StmtMutator::VisitStmt_(const PrefetchNode* op) { Stmt StmtMutator::VisitStmt_(const SeqStmtNode* op) { Array seq = Internal::Mutate(this, op->seq); - Stmt stmt = seq.same_as(op->seq) ? GetRef(op) : SeqStmt(seq); - return SeqStmt::Flatten(seq); + if (seq.same_as(op->seq)) { + return SeqStmt::Flatten(GetRef(op)); + } else { + auto node = CopyOnWrite(op); + node->seq = std::move(seq); + return SeqStmt::Flatten(SeqStmt(node)); + } } // advanced visit function for seqstmt. From 7fcf6ca57b4acb6433b82f4f74368aebb8946200 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 10 Apr 2023 09:20:43 -0500 Subject: [PATCH 08/10] Avoid length-1 SeqStmt in ExtractReductionUpdates --- src/tir/schedule/analysis/reducer.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/tir/schedule/analysis/reducer.cc b/src/tir/schedule/analysis/reducer.cc index b8ba7ac96f3c..e88fce160779 100644 --- a/src/tir/schedule/analysis/reducer.cc +++ b/src/tir/schedule/analysis/reducer.cc @@ -394,16 +394,15 @@ void ExtractReductionUpdates(const Optional& self, Block block, if (p_seq == nullptr && p_buf_store == nullptr) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/5); } - SeqStmt seq = - p_seq != nullptr ? GetRef(p_seq) : SeqStmt({GetRef(p_buf_store)}); - if (static_cast(seq->seq.size()) != n_buffers) { + Array seq = p_seq != nullptr ? p_seq->seq : Array{GetRef(p_buf_store)}; + if (static_cast(seq.size()) != n_buffers) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/6); } // Step 2. // - Create BufferStores according to the variables being stored. // - Construct the mapping from reduction buffers to the index. - for (const Stmt& stmt : seq->seq) { + for (const Stmt& stmt : seq) { const auto* buf_store = stmt.as(); if (buf_store == nullptr) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/5); From 0a3d3e0eba4bc93eff6f0065adbd34788b3046f5 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 10 Apr 2023 16:25:37 -0500 Subject: [PATCH 09/10] Additional CI fixes --- src/relay/backend/aot/aot_lower_main.cc | 2 +- tests/cpp/ir_functor_test.cc | 1 - tests/python/relay/aot/test_crt_aot.py | 2 +- tests/python/unittest/test_tvmscript_printer_tir.py | 4 ++-- 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/relay/backend/aot/aot_lower_main.cc b/src/relay/backend/aot/aot_lower_main.cc index fb13e8b66e5d..a97f7bcaa64a 100644 --- a/src/relay/backend/aot/aot_lower_main.cc +++ b/src/relay/backend/aot/aot_lower_main.cc @@ -674,7 +674,7 @@ class AOTMainLowerer : public MixedModeVisitor { })); } - tir::Stmt body = tir::SeqStmt({func_call}); + tir::Stmt body = tir::SeqStmt::Flatten(func_call); stmts_.push_back(body); } diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 2909915c3288..727b07c2e488 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -280,7 +280,6 @@ TEST(IRF, StmtMutator) { auto* ref2 = body2.get(); auto* extentptr = body.as()->extents.get(); // construct a recursive SeqStmt. - body = SeqStmt({body}); body = SeqStmt({body, body2}); body = SeqStmt({body, body2}); body = v(std::move(body)); diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py index 1eb34b07d7ab..f7e5af18d20e 100644 --- a/tests/python/relay/aot/test_crt_aot.py +++ b/tests/python/relay/aot/test_crt_aot.py @@ -1042,7 +1042,7 @@ def test_aot_codegen_checks_returns(): main_func = main_ir_module["__tvm_main__"] # Check operator call is wrapped properly - body = main_func.body[1].seq[0].seq[0].value + body = main_func.body.value assert ( repr(body) == 'T.tvm_check_return(0, -1, T.call_extern("int32", "tvmgen_default_fused_add",' diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py index 25272d912da2..8427754db71e 100644 --- a/tests/python/unittest/test_tvmscript_printer_tir.py +++ b/tests/python/unittest/test_tvmscript_printer_tir.py @@ -396,14 +396,14 @@ def test_prefetch(): def test_seq_stmt(): with IRBuilder() as ib: with T.serial(10): - T.evaluate(0) T.evaluate(1) + T.evaluate(2) obj = ib.get().body _assert_print( obj, """ -T.evaluate(0) T.evaluate(1) +T.evaluate(2) """, ) From a5eff130a76b629989c7f7851c5059e5a31d4f6a Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 11 Apr 2023 11:48:34 -0500 Subject: [PATCH 10/10] Avoid a couple more cases with flatten-able SeqStmt --- src/relay/backend/aot/aot_lower_main.cc | 4 ++-- src/tir/contrib/ethosu/passes.cc | 2 +- tests/cpp/ir_functor_test.cc | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/relay/backend/aot/aot_lower_main.cc b/src/relay/backend/aot/aot_lower_main.cc index a97f7bcaa64a..5bfe7cb8e749 100644 --- a/src/relay/backend/aot/aot_lower_main.cc +++ b/src/relay/backend/aot/aot_lower_main.cc @@ -417,7 +417,7 @@ class AOTMainLowerer : public MixedModeVisitor { * runner function needs to be legalized by the LegalizePackedCalls pass. */ tir::PrimFunc CreateMainFunc(String mod_name) { - tir::Stmt body = tir::SeqStmt(stmts_); + tir::Stmt body = tir::SeqStmt::Flatten(stmts_); // Allocate the sids std::unordered_map allocated; std::vector> sids_to_allocate; @@ -717,7 +717,7 @@ class AOTMainLowerer : public MixedModeVisitor { {tvm::tir::StringImm(device_hook_name), context}))); device_hooks.push_back(device_hook); } - return tir::SeqStmt(device_hooks); + return tir::SeqStmt::Flatten(device_hooks); } /*! diff --git a/src/tir/contrib/ethosu/passes.cc b/src/tir/contrib/ethosu/passes.cc index 369c4adc8536..5968febe9af0 100644 --- a/src/tir/contrib/ethosu/passes.cc +++ b/src/tir/contrib/ethosu/passes.cc @@ -144,7 +144,7 @@ class HoistAllocatesMutator : public StmtExprMutator { for (auto it = allocates_.rbegin(); it != allocates_.rend(); it++) { Allocate current_alloc = *it; if (it != allocates_.rbegin()) { - new_main_func_body = SeqStmt({new_main_func_body}); + new_main_func_body = SeqStmt::Flatten(new_main_func_body); } new_main_func_body = Allocate(current_alloc->buffer_var, current_alloc->dtype, current_alloc->extents, diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 727b07c2e488..30b1bc78247a 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -295,7 +295,7 @@ TEST(IRF, StmtMutator) { Stmt body2 = Evaluate(1); auto* extentptr = body.as()->extents.get(); // construct a recursive SeqStmt. - body = SeqStmt({body}); + body = SeqStmt({body, body2}); auto bref = body; body = SeqStmt({body, body2}); body = v(std::move(body));