diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index 9ed9973871d9..51c99fb375f3 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -734,19 +734,33 @@ class SeqStmt : public Stmt { public: explicit Flattener(Array* seq) : seq_(seq) {} - void operator()(size_t i, const Stmt& stmt) const { - if (!stmt.defined()) return; - if (auto* op = stmt.as()) { - operator()(0, op->seq); - } else { - seq_->push_back(stmt); + 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; } - } - template - void operator()(size_t i, const T& seq) const { - for (auto v : seq) { - this->operator()(0, v); + if constexpr (std::is_same_v) { + // No need for dynamic type-checking if the static type is a + // SeqStmt. + (*this)(0, stmt_or_seq->seq); + } else 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); + } + } else if constexpr (std::is_base_of_v) { + // Any other Stmt type just gets appended. + seq_->push_back(stmt_or_seq); + } else { + // Anything else is treated as an iterable of Stmt. + for (auto v : stmt_or_seq) { + this->operator()(0, v); + } } }