From db90652b949f731216f989321862d9f66d7014fd Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Wed, 29 Dec 2021 11:12:49 +0800 Subject: [PATCH 01/26] Rewrite CollectComputeLocations --- src/tir/schedule/analysis.h | 15 ++++----- src/tir/schedule/analysis/analysis.cc | 44 ++++++++++++++++++++------- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 2f5f03f29140..1547c82714a7 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -333,7 +333,7 @@ bool HasSingleChild(const StmtSRef& loop_or_block_sref); /*! * \brief Check if a block is the direct children of the root block - * \param self The TIR schedule class + * \param self The schedule state * \param block_sref The block to be analyzed * \return A boolean flag indicating if the block is the subroot block */ @@ -348,12 +348,13 @@ bool IsSubrootBlock(const tir::ScheduleState& self, const tir::StmtSRef& block_s StmtSRef GetSRefLowestCommonAncestor(const Array& srefs); /*! - * \brief Collect all the feasible compute locations among the loops above the block - * \param self The TIR schedule class - * \param block_sref The input block - * \return All the feasible compute locations among the loops above the block + * \brief Collect all the feasible compute locations among the loops above the input blocks + * \param self The schedule state + * \param block_srefs The input blocks whose compute locations are to be collected + * \return All the feasible compute locations among the loops above the blocks */ -Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref); +Array CollectComputeLocation(const ScheduleState& self, + const Array& block_srefs); /******** Tensorization ********/ @@ -452,7 +453,7 @@ class TensorizeInfo : public ObjectRef { /*! * \brief Check if the given block can be tensorized, and in the meantime gather the necessary * information for tensorization - * \param self The TIR schedule + * \param self The schedule state * \param block_sref The block to be analyzed * \param desc_func The target function for tensorization * \return The necessary information used for tensorization, or NullOpt if the block cannot be diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 03fd1dd7ca55..8f4b6482a65c 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -822,24 +822,46 @@ StmtSRef GetSRefLowestCommonAncestor(const Array& srefs) { return GetRef(p); } -Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref) { - Array loop_srefs = GetLoops(block_sref); - Array result; - result.reserve(loop_srefs.size()); +Array CollectComputeLocation(const ScheduleState& self, + const Array& block_srefs) { + StmtSRef lca_sref = GetSRefLowestCommonAncestor(block_srefs); + if (lca_sref->StmtAs() != nullptr) { + return {}; + } + + Array loop_srefs = GetLoops(block_srefs[0]); + int lca_pos = std::find(loop_srefs.begin(), loop_srefs.end(), lca_sref) - loop_srefs.begin(); + ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); + + std::vector loop_iter_types; + loop_iter_types.reserve(lca_pos + 1); + int i_last_datapar = -1; bool visited_reduce = false; - for (const StmtSRef& loop_sref : loop_srefs) { - const ForNode* loop = TVM_SREF_TO_FOR(loop, loop_sref); - IterVarType iter_type = GetLoopIterType(loop_sref); + + for (int i = 0; i <= lca_pos; ++i) { + IterVarType iter_type = GetLoopIterType(loop_srefs[i]); + loop_iter_types.push_back(iter_type); if (iter_type == IterVarType::kDataPar) { + i_last_datapar = i; + } + } + + for (int i = 0; i <= lca_pos; ++i) { + if (loop_iter_types[i] == IterVarType::kDataPar) { if (visited_reduce) { - break; + return {loop_srefs.begin(), loop_srefs.begin() + i}; } - } else { + } else if (loop_iter_types[i] == IterVarType::kCommReduce) { visited_reduce = true; + if (i > i_last_datapar) { + return {loop_srefs.begin(), loop_srefs.begin() + i}; + } + } else { + return {loop_srefs.begin(), loop_srefs.begin() + i}; } - result.push_back(loop_sref); } - return result; + + return {}; } /******** Tensorization ********/ From b7bbedde95337fe76f3618e97b4ed6c5c775f15b Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Wed, 29 Dec 2021 11:28:22 +0800 Subject: [PATCH 02/26] rewrite SampleComputeLocation --- include/tvm/tir/schedule/schedule.h | 9 ++-- python/tvm/tir/schedule/schedule.py | 13 +++--- src/tir/schedule/concrete_schedule.cc | 4 +- src/tir/schedule/concrete_schedule.h | 3 +- src/tir/schedule/primitive.h | 9 ++-- src/tir/schedule/primitive/sampling.cc | 64 ++++++++++---------------- src/tir/schedule/traced_schedule.cc | 8 ++-- src/tir/schedule/traced_schedule.h | 3 +- 8 files changed, 52 insertions(+), 61 deletions(-) diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index b170b191b00d..44dde3a323d6 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -211,12 +211,13 @@ class ScheduleNode : public runtime::Object { virtual Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) = 0; /*! - * \brief Sample a compute-at location on a BlockRV so that its producer can compute at that loop - * \param block_rv The consumer block to be computed at + * \brief Sample a compute-at location on a list of BlockRVs so that their common producer can + * compute at that loop + * \param block_rvs The consumer blocks that are used to gather the compute-at candidate locations * \param decision The sampling decision - * \return The sampled loop to be computed at + * \return The sampled loop where the producer is to be computed at */ - virtual LoopRV SampleComputeLocation(const BlockRV& block_rv, + virtual LoopRV SampleComputeLocation(const Array& block_rvs, Optional decision = NullOpt) = 0; /******** Schedule: Get blocks & loops ********/ diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 06806a8e1b5c..c6b4a244a4cc 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -372,26 +372,27 @@ def sample_perfect_tile( def sample_compute_location( self, - block: BlockRV, + blocks: List[BlockRV], decision: Optional[int] = None, ) -> LoopRV: - """Sample a compute-at location on a BlockRV so that its producer can compute at that loop + """Sample a compute-at location on a list of blocks so that its producer can compute at that + loop Parameters ---------- - block : BlockRV - The consumer block to be computed at + blocks : List[BlockRV] + The consumer blocks that are used to gather the compute-at candidate locations decision : Optional[int] The sampling decision Returns ------- result : LoopRV - The sampled loop to be computed at + The sampled loop where the producer is to be computed at """ return _ffi_api.ScheduleSampleComputeLocation( # pylint: disable=no-member self, - block, + blocks, decision, ) diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index afaa998e6c8d..e333f1cfaf86 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -240,11 +240,11 @@ Array ConcreteScheduleNode::SamplePerfectTile(const LoopRV& loop_rv, int throw; } -LoopRV ConcreteScheduleNode::SampleComputeLocation(const BlockRV& block_rv, +LoopRV ConcreteScheduleNode::SampleComputeLocation(const Array& block_rvs, Optional decision) { TVM_TIR_SCHEDULE_BEGIN(); return CreateRV( - tir::SampleComputeLocation(state_, &this->rand_state_, this->GetSRef(block_rv), &decision)); + tir::SampleComputeLocation(state_, &this->rand_state_, this->GetSRefs(block_rvs), &decision)); TVM_TIR_SCHEDULE_END("sample-compute-location", this->error_render_level_); throw; } diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index cacd8e389dff..2acf6ac947c8 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -86,7 +86,7 @@ class ConcreteScheduleNode : public ScheduleNode { Optional decision = NullOpt) override; Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) override; - LoopRV SampleComputeLocation(const BlockRV& block_rv, + LoopRV SampleComputeLocation(const Array& block_rvs, Optional decision = NullOpt) override; /******** Schedule: Get blocks & loops ********/ BlockRV GetBlock(const String& name, const String& func_name = "main") override; @@ -284,6 +284,7 @@ inline Array GetSRefsHelper(const ConcreteScheduleNode* sch, const Arr return result; } +// Todo: expose it to schedule.h? inline Array ConcreteScheduleNode::GetSRefs(const Array& rvs) const { return GetSRefsHelper(this, rvs); } diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index 6ca748e28573..562cb01f2e4f 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -110,15 +110,16 @@ TVM_DLL std::vector SamplePerfectTile( const tir::StmtSRef& loop_sref, int32_t n_split, int32_t max_innermost_factor, Optional>* decision); /*! - * \brief Sample a compute-at location on a BlockRV so that its producer can compute at that loop + * \brief Sample a compute-at location on a list of blocks so that its producer can compute at that + * loop * \param self The schedule state * \param rand_state The random state - * \param block_rv The consumer block to be computed at - * \return The sampled loop to be computed at + * \param block_srefs The consumer blocks that are used to gather the compute-at candidate locations + * \return The sampled loop where the producer is to be computed at */ TVM_DLL tir::StmtSRef SampleComputeLocation( tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, - const tir::StmtSRef& block_sref, Optional* decision); + const Array& block_srefs, Optional* decision); /******** Schedule: Get blocks & loops ********/ /*! diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index 4c2d3d835dd2..1150b675aebc 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -347,51 +347,29 @@ std::vector SamplePerfectTile( tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, - const tir::StmtSRef& block_sref, Optional* decision) { + const Array& block_srefs, + Optional* decision) { // Find all possible compute-at locations - Array loop_srefs = tir::CollectComputeLocation(self, block_sref); + Array loop_srefs = tir::CollectComputeLocation(self, block_srefs); int n = loop_srefs.size(); - // Extract non-unit loops - std::vector choices; - choices.reserve(n); - for (int i = 0; i < n; ++i) { - const int64_t* extent = tir::GetLoopIntExtent(loop_srefs[i]); - if (extent != nullptr) { - choices.push_back(i); - } - } // The decision made, by default it is -1 int i = -1; if (decision->defined()) { // Handle existing decision - const auto* int_imm = decision->as(); - int64_t decided = int_imm->value; - if (decided == -2 || decided == -1) { - i = decided; + int64_t val_decision = decision->as()->value; + if (val_decision >= n) { + LOG(WARNING) << "old decision is " << val_decision << " while current candidate count is " + << n << ". Hence cannot reapply the old decision"; + i = n - 1; } else { - for (int choice : choices) { - if (choice <= decided) { - i = choice; - } else { - break; - } - } + i = val_decision; } } else { // Sample possible combinations - i = SampleInt(rand_state, -2, choices.size()); - if (i >= 0) { - i = choices[i]; - } + i = SampleInt(rand_state, -2, n); } *decision = Integer(i); - if (i == -2) { - return tir::StmtSRef::InlineMark(); - } - if (i == -1) { - return tir::StmtSRef::RootMark(); - } - return loop_srefs[i]; + return i >= 0 ? loop_srefs[i] : i == -1 ? tir::StmtSRef::RootMark() : tir::StmtSRef::InlineMark(); } /******** InstructionKind Registration ********/ @@ -467,17 +445,25 @@ struct SampleComputeLocationTraits : public UnpackedInstTraits + static TVM_ALWAYS_INLINE void _SetInputs(const runtime::TVMArgsSetter& setter, + const Array& inputs) { + setter(delta, inputs); + } + + static LoopRV UnpackedApplyToSchedule(Schedule sch, // + Array block_rvs, // Optional decision) { - return sch->SampleComputeLocation(block_rv, decision); + return sch->SampleComputeLocation(block_rvs, decision); } - static String UnpackedAsPython(Array outputs, // - String block_rv, // + static String UnpackedAsPython(Array outputs, // + Array block_rvs, // Optional decision) { PythonAPICall py("sample_compute_location"); - py.Input("block", block_rv); + for (const String& block_rv : block_rvs) { + py.Input("", block_rv); + } py.Decision(decision); py.SingleOutput(outputs); return py.Str(); diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc index fc92306354fe..d4f0853a983b 100644 --- a/src/tir/schedule/traced_schedule.cc +++ b/src/tir/schedule/traced_schedule.cc @@ -73,14 +73,14 @@ Array TracedScheduleNode::SamplePerfectTile(const LoopRV& loop_rv, int n return results; } -LoopRV TracedScheduleNode::SampleComputeLocation(const BlockRV& block_rv, +LoopRV TracedScheduleNode::SampleComputeLocation(const Array& block_rvs, Optional decision) { - LoopRV result = CreateRV(tir::SampleComputeLocation(this->state_, &this->rand_state_, - this->GetSRef(block_rv), &decision)); + LoopRV result = CreateRV(tir::SampleComputeLocation( + this->state_, &this->rand_state_, this->GetSRefs(block_rvs), &decision)); static const InstructionKind& kind = InstructionKind::Get("SampleComputeLocation"); trace_->Append(/*inst=*/Instruction(/*kind=*/kind, // - /*inputs=*/{block_rv}, + /*inputs=*/{block_rvs.begin(), block_rvs.end()}, /*attrs=*/{}, /*outputs=*/{result}), /*decision=*/decision); diff --git a/src/tir/schedule/traced_schedule.h b/src/tir/schedule/traced_schedule.h index 12696567816a..3a23cbad85ce 100644 --- a/src/tir/schedule/traced_schedule.h +++ b/src/tir/schedule/traced_schedule.h @@ -51,7 +51,8 @@ class TracedScheduleNode : public ConcreteScheduleNode { Optional decision = NullOpt) final; Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) final; - LoopRV SampleComputeLocation(const BlockRV& block_rv, Optional decision = NullOpt) final; + LoopRV SampleComputeLocation(const Array& block_rvs, + Optional decision = NullOpt) final; /******** Schedule: Get blocks & loops ********/ BlockRV GetBlock(const String& name, const String& func_name = "main") final; Array GetLoops(const BlockRV& block_rv) final; From cb517631c722dbf0bfd5dfb5c1239091a1cb3df8 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Wed, 29 Dec 2021 13:22:10 +0800 Subject: [PATCH 03/26] Fix CollectComputeLocation/SampleComputeLocation --- src/tir/schedule/analysis/analysis.cc | 16 +++++++++++----- src/tir/schedule/primitive/sampling.cc | 6 ++---- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 8f4b6482a65c..333cf27077a0 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -822,15 +822,21 @@ StmtSRef GetSRefLowestCommonAncestor(const Array& srefs) { return GetRef(p); } -Array CollectComputeLocation(const ScheduleState& self, +Array CollectComputeLocation(const ScheduleState& self, // Todo: improve (producer input) const Array& block_srefs) { - StmtSRef lca_sref = GetSRefLowestCommonAncestor(block_srefs); - if (lca_sref->StmtAs() != nullptr) { + if (block_srefs.empty()) { + return {}; + } + + StmtSRef loop_boundary_sref = block_srefs.size() > 1 ? GetSRefLowestCommonAncestor(block_srefs) + : GetRef(block_srefs[0]->parent); + if (loop_boundary_sref->StmtAs() != nullptr) { return {}; } Array loop_srefs = GetLoops(block_srefs[0]); - int lca_pos = std::find(loop_srefs.begin(), loop_srefs.end(), lca_sref) - loop_srefs.begin(); + int lca_pos = + std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary_sref) - loop_srefs.begin(); ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); std::vector loop_iter_types; @@ -861,7 +867,7 @@ Array CollectComputeLocation(const ScheduleState& self, } } - return {}; + return {loop_srefs.begin(), loop_srefs.begin() + lca_pos + 1}; } /******** Tensorization ********/ diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index 1150b675aebc..fc2b0b735a8d 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -366,7 +366,7 @@ tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, } } else { // Sample possible combinations - i = SampleInt(rand_state, -2, n); + i = SampleInt(rand_state, -1, n); // Todo: Temporarily disable Inline } *decision = Integer(i); return i >= 0 ? loop_srefs[i] : i == -1 ? tir::StmtSRef::RootMark() : tir::StmtSRef::InlineMark(); @@ -461,9 +461,7 @@ struct SampleComputeLocationTraits : public UnpackedInstTraits block_rvs, // Optional decision) { PythonAPICall py("sample_compute_location"); - for (const String& block_rv : block_rvs) { - py.Input("", block_rv); - } + py.Input("blocks", block_rvs); py.Decision(decision); py.SingleOutput(outputs); return py.Str(); From a445316f393813782df4d604e471c63a78052d7e Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Wed, 29 Dec 2021 21:48:23 +0800 Subject: [PATCH 04/26] Rewrite rule RandomComputeLocation; aligned --- python/tvm/meta_schedule/tune.py | 1 + .../mutator/mutate_compute_location.cc | 38 +++++++------ .../schedule_rule/random_compute_location.cc | 54 +++++++++++-------- .../space_generator/post_order_apply.cc | 3 +- 4 files changed, 51 insertions(+), 45 deletions(-) diff --git a/python/tvm/meta_schedule/tune.py b/python/tvm/meta_schedule/tune.py index 21d7a2614261..245de28744e0 100644 --- a/python/tvm/meta_schedule/tune.py +++ b/python/tvm/meta_schedule/tune.py @@ -133,6 +133,7 @@ def _mutator_probs() -> Dict[Mutator, float]: return { M.MutateTileSize(): 0.9, + M.MutateComputeLocation(): 0.05, M.MutateUnroll(): 0.03, M.MutateParallel(max_jobs_per_core=16): 0.02, } diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc index 1af03f073671..d6779f2402a5 100644 --- a/src/meta_schedule/mutator/mutate_compute_location.cc +++ b/src/meta_schedule/mutator/mutate_compute_location.cc @@ -25,7 +25,7 @@ using tir::Instruction; using tir::InstructionKind; using tir::Trace; -/*! \brief Create a Mutator that mutates auto unroll step */ +/*! \brief Create a Mutator that mutates auto unroll step */ // Todo class MutateComputeLocationNode : public MutatorNode { public: /*! \brief JSON representation of the workload */ @@ -73,32 +73,30 @@ std::vector FindCandidates(const Trace& tr const ObjectRef& decision) -> ObjectRef { if (inst->kind.same_as(inst_sample_compute_location)) { // The decision made - int decided = Downcast(decision)->value; + int old_decision = Downcast(decision)->value; // Extract the inputs - ICHECK_EQ(inputs.size(), 1); - tir::BlockRV block_rv = Downcast(inputs[0]); - tir::StmtSRef block_sref = sch->GetSRef(block_rv); + Array block_srefs; + block_srefs.reserve(inputs.size()); + for (const ObjectRef& obj : inputs) { + block_srefs.push_back(sch->GetSRef(Downcast(obj))); + } // Extract locations that can be computed at - Array loop_srefs = CollectComputeLocation(sch->state(), block_sref); - std::vector locs{-2, -1}; - { - int i = 0; - for (const tir::StmtSRef& loop_sref : loop_srefs) { - int64_t extent = *tir::GetLoopIntExtent(loop_sref); - if (extent != 1 && extent != -1) { - locs.push_back(i); - } - ++i; - } + Array loop_srefs = CollectComputeLocation(sch->state(), block_srefs); + // std::vector locs{-2, -1}; // Todo: temporarily disable inline + std::vector locs{-1}; + for (int i = 0; i < static_cast(loop_srefs.size()); ++i) { + locs.push_back(i); } - // Remove `decided` - std::vector::iterator rm = std::find(locs.begin(), locs.end(), decided); + // Remove `old_decision` + std::vector::iterator rm = std::find(locs.begin(), locs.end(), old_decision); if (rm != locs.end()) { locs.erase(rm); } + // Add the candidate - ICHECK(!locs.empty()); - candidates.emplace_back(inst, std::move(locs)); + if (!locs.empty()) { + candidates.emplace_back(inst, std::move(locs)); + } } return decision; }; diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index 1757f650aabb..c02c90019ac2 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -23,30 +23,34 @@ namespace meta_schedule { class RandomComputeLocationNode : public ScheduleRuleNode { public: - bool IsFreeBlock(const tir::Schedule sch, const tir::StmtSRef& block_sref) const { + bool CheckConditions(const tir::Schedule sch, const tir::BlockRV& block_rv, + Array* consumers) const { + const tir::StmtSRef& block_sref = sch->GetSRef(block_rv); + const tir::BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); + + // Cond 1. The block is not the root block. if (block_sref->parent == nullptr) { return false; } - if (!tir::IsSubrootBlock(sch->state(), block_sref)) { + // Cond 2. The block should be the direct child block of the root block. + if (!tir::IsSubrootBlock(sch->state(), block_sref)) { // Todo return false; } - tir::ScheduleState state = sch->state(); - if (!tir::IsCompleteBlock(state, block_sref, - tir::GetScopeRoot(state, block_sref, false, false))) { + // Cond 3 & 4. The block has at least one outer loop, and the outermost loop has only one child + // block. + Array loop_srefs = tir::GetLoops(block_sref); + if (loop_srefs.empty()) { return false; } - Array loop_srefs = tir::GetLoops(block_sref); - for (const tir::StmtSRef& loop_sref : loop_srefs) { - if (!tir::HasSingleChild(loop_sref)) { - return false; - } + if (tir::GetChildBlockSRefOnSRefTree(sch->state(), loop_srefs[0]).size() > 1) { + return false; } - Array binds = tir::GetBlockRealize(state, block_sref)->iter_values; - for (const PrimExpr& bind : binds) { - if (!bind->IsInstance() && !bind->IsInstance()) { - return false; - } + // Cond 5. The block has at lease one consumer. + *consumers = sch->GetConsumers(block_rv); + if (consumers->empty()) { + return false; } + return true; } @@ -55,24 +59,28 @@ class RandomComputeLocationNode : public ScheduleRuleNode { // Inherited from ScheduleRuleNode Array Apply(const tir::Schedule& sch, const tir::BlockRV& block_rv) final { - tir::StmtSRef block_sref = sch->GetSRef(block_rv); - if (!IsFreeBlock(sch, block_sref)) { - return {sch}; - } - Array consumers = sch->GetConsumers(block_rv); - if (consumers.size() != 1) { + Array consumers{nullptr}; + if (!CheckConditions(sch, block_rv, &consumers)) { return {sch}; } - tir::BlockRV consumer = consumers[0]; + ICHECK(consumers.defined()); + // Try to compute `block_rv` at `consumer` + int err_cnt = 0; for (;;) { - tir::LoopRV compute_at_loc = sch->SampleComputeLocation(consumer); + if (err_cnt == 100) { + LOG(WARNING) << "err_cnt = 100, force quit"; + break; + } + tir::LoopRV compute_at_loc = sch->SampleComputeLocation(consumers); try { sch->ComputeAt(block_rv, compute_at_loc, true); } catch (const dmlc::Error& e) { // ComputeAt fails, cleanup the following before re-try: // 1) trace: instruction & decisions // 2) sym_tab + LOG(INFO) << "fail"; + ++err_cnt; sch->trace().value()->Pop(); sch->RemoveRV(compute_at_loc); continue; diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index 3f685407817b..324422aa80b3 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -105,10 +105,9 @@ class PostOrderApplyNode : public SpaceGeneratorNode { Array result{sch}; // Enumerate the schedule rules first because you can // always concat multiple schedule rules as one - Array all_blocks = BlockCollector::Collect(sch); for (ScheduleRule sch_rule : sch_rules_) { for (const tir::Schedule& sch : result) { - stack.emplace_back(sch, all_blocks); + stack.emplace_back(sch, BlockCollector::Collect(sch)); } result.clear(); From 557bf4bbe9713b0a635cfcab1924222a41366f75 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Wed, 29 Dec 2021 23:58:40 +0800 Subject: [PATCH 05/26] Fix AddRFactor to avoid unit-loop template --- src/meta_schedule/schedule_rule/add_rfactor.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/meta_schedule/schedule_rule/add_rfactor.cc b/src/meta_schedule/schedule_rule/add_rfactor.cc index 7be0dcff5a07..c4c05b226b5c 100644 --- a/src/meta_schedule/schedule_rule/add_rfactor.cc +++ b/src/meta_schedule/schedule_rule/add_rfactor.cc @@ -89,7 +89,11 @@ Array AddRFactorNode::Apply(const tir::Schedule& sch, const tir:: ReorderAndFuseReductionLoops(sch, block_rv, &fused_reduce_loop, &num_spatial_loops); // Split the fused reduction loop. - Array factors = sch->SamplePerfectTile(fused_reduce_loop, 2, max_innermost_factor); + Array factors; + do { + factors = sch->SamplePerfectTile(fused_reduce_loop, 2, max_innermost_factor); + } while (*tir::as_const_int(sch->Get(factors[0])) == 1 || + *tir::as_const_int(sch->Get(factors[1])) == 1); const Array& split_loops = sch->Split(fused_reduce_loop, {factors.begin(), factors.end()}); From 1327f7dcafd3d00e9f20bbee4019d65b061c0cc0 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 00:51:51 +0800 Subject: [PATCH 06/26] Change the API of SampleComputeLocation --- include/tvm/tir/schedule/schedule.h | 9 ++++---- python/tvm/tir/schedule/schedule.py | 15 ++++++------ .../mutator/mutate_compute_location.cc | 9 +++----- .../schedule_rule/random_compute_location.cc | 15 ++++-------- src/tir/schedule/analysis.h | 9 ++++---- src/tir/schedule/analysis/analysis.cc | 15 +++++++----- src/tir/schedule/concrete_schedule.cc | 4 ++-- src/tir/schedule/concrete_schedule.h | 2 +- src/tir/schedule/primitive.h | 10 ++++---- src/tir/schedule/primitive/sampling.cc | 23 +++++++------------ src/tir/schedule/traced_schedule.cc | 8 +++---- src/tir/schedule/traced_schedule.h | 3 +-- 12 files changed, 53 insertions(+), 69 deletions(-) diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index 44dde3a323d6..dc5e99faccb3 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -211,13 +211,12 @@ class ScheduleNode : public runtime::Object { virtual Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) = 0; /*! - * \brief Sample a compute-at location on a list of BlockRVs so that their common producer can - * compute at that loop - * \param block_rvs The consumer blocks that are used to gather the compute-at candidate locations + * \brief Sample a compute-at location of the given block + * \param block_rv The block whose compute-at location is to be sampled * \param decision The sampling decision - * \return The sampled loop where the producer is to be computed at + * \return The sampled loop where the input block is to be computed at */ - virtual LoopRV SampleComputeLocation(const Array& block_rvs, + virtual LoopRV SampleComputeLocation(const BlockRV& block_rv, Optional decision = NullOpt) = 0; /******** Schedule: Get blocks & loops ********/ diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index c6b4a244a4cc..cb207fea6ab7 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -370,29 +370,28 @@ def sample_perfect_tile( ) ) - def sample_compute_location( + def sample_compute_location( # Todo: add some unittests self, - blocks: List[BlockRV], + block: BlockRV, decision: Optional[int] = None, ) -> LoopRV: - """Sample a compute-at location on a list of blocks so that its producer can compute at that - loop + """Sample a compute-at location of the given block Parameters ---------- - blocks : List[BlockRV] - The consumer blocks that are used to gather the compute-at candidate locations + block : BlockRV + The block whose compute-at location is to be sampled decision : Optional[int] The sampling decision Returns ------- result : LoopRV - The sampled loop where the producer is to be computed at + The sampled loop where the input block is to be computed at """ return _ffi_api.ScheduleSampleComputeLocation( # pylint: disable=no-member self, - blocks, + block, decision, ) diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc index d6779f2402a5..1e387a053c20 100644 --- a/src/meta_schedule/mutator/mutate_compute_location.cc +++ b/src/meta_schedule/mutator/mutate_compute_location.cc @@ -75,13 +75,10 @@ std::vector FindCandidates(const Trace& tr // The decision made int old_decision = Downcast(decision)->value; // Extract the inputs - Array block_srefs; - block_srefs.reserve(inputs.size()); - for (const ObjectRef& obj : inputs) { - block_srefs.push_back(sch->GetSRef(Downcast(obj))); - } + ICHECK_EQ(inputs.size(), 1); + tir::StmtSRef block_sref = sch->GetSRef(Downcast(inputs[0])); // Extract locations that can be computed at - Array loop_srefs = CollectComputeLocation(sch->state(), block_srefs); + Array loop_srefs = CollectComputeLocation(sch->state(), block_sref); // std::vector locs{-2, -1}; // Todo: temporarily disable inline std::vector locs{-1}; for (int i = 0; i < static_cast(loop_srefs.size()); ++i) { diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index c02c90019ac2..5dc6f66e3e07 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -23,8 +23,7 @@ namespace meta_schedule { class RandomComputeLocationNode : public ScheduleRuleNode { public: - bool CheckConditions(const tir::Schedule sch, const tir::BlockRV& block_rv, - Array* consumers) const { + bool CheckConditions(const tir::Schedule sch, const tir::BlockRV& block_rv) const { const tir::StmtSRef& block_sref = sch->GetSRef(block_rv); const tir::BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); @@ -46,8 +45,7 @@ class RandomComputeLocationNode : public ScheduleRuleNode { return false; } // Cond 5. The block has at lease one consumer. - *consumers = sch->GetConsumers(block_rv); - if (consumers->empty()) { + if (tir::GetConsumers(sch->state(), sch->GetSRef(block_rv)).empty()) { return false; } @@ -59,20 +57,17 @@ class RandomComputeLocationNode : public ScheduleRuleNode { // Inherited from ScheduleRuleNode Array Apply(const tir::Schedule& sch, const tir::BlockRV& block_rv) final { - Array consumers{nullptr}; - if (!CheckConditions(sch, block_rv, &consumers)) { + if (!CheckConditions(sch, block_rv)) { return {sch}; } - ICHECK(consumers.defined()); - // Try to compute `block_rv` at `consumer` int err_cnt = 0; for (;;) { if (err_cnt == 100) { - LOG(WARNING) << "err_cnt = 100, force quit"; + LOG(INFO) << "err_cnt = 100, force quit"; break; } - tir::LoopRV compute_at_loc = sch->SampleComputeLocation(consumers); + tir::LoopRV compute_at_loc = sch->SampleComputeLocation(block_rv); try { sch->ComputeAt(block_rv, compute_at_loc, true); } catch (const dmlc::Error& e) { diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 1547c82714a7..ecb7daa3d6ce 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -348,13 +348,12 @@ bool IsSubrootBlock(const tir::ScheduleState& self, const tir::StmtSRef& block_s StmtSRef GetSRefLowestCommonAncestor(const Array& srefs); /*! - * \brief Collect all the feasible compute locations among the loops above the input blocks + * \brief Collect all the feasible compute-at locations of the input block * \param self The schedule state - * \param block_srefs The input blocks whose compute locations are to be collected - * \return All the feasible compute locations among the loops above the blocks + * \param block_sref The block whose compute-at locations are to be collected + * \return All the feasible compute-at locations of the input block, given as an array of loop srefs */ -Array CollectComputeLocation(const ScheduleState& self, - const Array& block_srefs); +Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref); /******** Tensorization ********/ diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 333cf27077a0..c1065b250be8 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -822,19 +822,19 @@ StmtSRef GetSRefLowestCommonAncestor(const Array& srefs) { return GetRef(p); } -Array CollectComputeLocation(const ScheduleState& self, // Todo: improve (producer input) - const Array& block_srefs) { - if (block_srefs.empty()) { +Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref) { + Array consumers = GetConsumers(self, block_sref); + if (consumers.empty()) { return {}; } - StmtSRef loop_boundary_sref = block_srefs.size() > 1 ? GetSRefLowestCommonAncestor(block_srefs) - : GetRef(block_srefs[0]->parent); + StmtSRef loop_boundary_sref = consumers.size() > 1 ? GetSRefLowestCommonAncestor(consumers) + : GetRef(consumers[0]->parent); if (loop_boundary_sref->StmtAs() != nullptr) { return {}; } - Array loop_srefs = GetLoops(block_srefs[0]); + Array loop_srefs = GetLoops(consumers[0]); int lca_pos = std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary_sref) - loop_srefs.begin(); ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); @@ -852,6 +852,9 @@ Array CollectComputeLocation(const ScheduleState& self, // Todo: impr } } + // Todo 1: take the reduction iterators of the input block into considertion + // Todo 2: skip unit loops + for (int i = 0; i <= lca_pos; ++i) { if (loop_iter_types[i] == IterVarType::kDataPar) { if (visited_reduce) { diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index e333f1cfaf86..afaa998e6c8d 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -240,11 +240,11 @@ Array ConcreteScheduleNode::SamplePerfectTile(const LoopRV& loop_rv, int throw; } -LoopRV ConcreteScheduleNode::SampleComputeLocation(const Array& block_rvs, +LoopRV ConcreteScheduleNode::SampleComputeLocation(const BlockRV& block_rv, Optional decision) { TVM_TIR_SCHEDULE_BEGIN(); return CreateRV( - tir::SampleComputeLocation(state_, &this->rand_state_, this->GetSRefs(block_rvs), &decision)); + tir::SampleComputeLocation(state_, &this->rand_state_, this->GetSRef(block_rv), &decision)); TVM_TIR_SCHEDULE_END("sample-compute-location", this->error_render_level_); throw; } diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index 2acf6ac947c8..d625c295079e 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -86,7 +86,7 @@ class ConcreteScheduleNode : public ScheduleNode { Optional decision = NullOpt) override; Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) override; - LoopRV SampleComputeLocation(const Array& block_rvs, + LoopRV SampleComputeLocation(const BlockRV& block_rv, Optional decision = NullOpt) override; /******** Schedule: Get blocks & loops ********/ BlockRV GetBlock(const String& name, const String& func_name = "main") override; diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index 562cb01f2e4f..c93ed6a1bf27 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -110,16 +110,16 @@ TVM_DLL std::vector SamplePerfectTile( const tir::StmtSRef& loop_sref, int32_t n_split, int32_t max_innermost_factor, Optional>* decision); /*! - * \brief Sample a compute-at location on a list of blocks so that its producer can compute at that - * loop + * \brief Sample a compute-at location of the given block * \param self The schedule state * \param rand_state The random state - * \param block_srefs The consumer blocks that are used to gather the compute-at candidate locations - * \return The sampled loop where the producer is to be computed at + * \param block_sref The sref of the block whose compute-at location is to be sampled + * \param decision The sampling decision + * \return The sampled loop where the input block is to be computed at */ TVM_DLL tir::StmtSRef SampleComputeLocation( tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, - const Array& block_srefs, Optional* decision); + const tir::StmtSRef& block_sref, Optional* decision); /******** Schedule: Get blocks & loops ********/ /*! diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index fc2b0b735a8d..86989097a6c0 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -347,10 +347,9 @@ std::vector SamplePerfectTile( tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, - const Array& block_srefs, - Optional* decision) { + const StmtSRef& block_sref, Optional* decision) { // Find all possible compute-at locations - Array loop_srefs = tir::CollectComputeLocation(self, block_srefs); + Array loop_srefs = CollectComputeLocation(self, block_sref); int n = loop_srefs.size(); // The decision made, by default it is -1 int i = -1; @@ -445,23 +444,17 @@ struct SampleComputeLocationTraits : public UnpackedInstTraits - static TVM_ALWAYS_INLINE void _SetInputs(const runtime::TVMArgsSetter& setter, - const Array& inputs) { - setter(delta, inputs); - } - - static LoopRV UnpackedApplyToSchedule(Schedule sch, // - Array block_rvs, // + static LoopRV UnpackedApplyToSchedule(Schedule sch, // + BlockRV block_rv, // Optional decision) { - return sch->SampleComputeLocation(block_rvs, decision); + return sch->SampleComputeLocation(block_rv, decision); } - static String UnpackedAsPython(Array outputs, // - Array block_rvs, // + static String UnpackedAsPython(Array outputs, // + String block_rv, // Optional decision) { PythonAPICall py("sample_compute_location"); - py.Input("blocks", block_rvs); + py.Input("block", block_rv); py.Decision(decision); py.SingleOutput(outputs); return py.Str(); diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc index d4f0853a983b..fc92306354fe 100644 --- a/src/tir/schedule/traced_schedule.cc +++ b/src/tir/schedule/traced_schedule.cc @@ -73,14 +73,14 @@ Array TracedScheduleNode::SamplePerfectTile(const LoopRV& loop_rv, int n return results; } -LoopRV TracedScheduleNode::SampleComputeLocation(const Array& block_rvs, +LoopRV TracedScheduleNode::SampleComputeLocation(const BlockRV& block_rv, Optional decision) { - LoopRV result = CreateRV(tir::SampleComputeLocation( - this->state_, &this->rand_state_, this->GetSRefs(block_rvs), &decision)); + LoopRV result = CreateRV(tir::SampleComputeLocation(this->state_, &this->rand_state_, + this->GetSRef(block_rv), &decision)); static const InstructionKind& kind = InstructionKind::Get("SampleComputeLocation"); trace_->Append(/*inst=*/Instruction(/*kind=*/kind, // - /*inputs=*/{block_rvs.begin(), block_rvs.end()}, + /*inputs=*/{block_rv}, /*attrs=*/{}, /*outputs=*/{result}), /*decision=*/decision); diff --git a/src/tir/schedule/traced_schedule.h b/src/tir/schedule/traced_schedule.h index 3a23cbad85ce..12696567816a 100644 --- a/src/tir/schedule/traced_schedule.h +++ b/src/tir/schedule/traced_schedule.h @@ -51,8 +51,7 @@ class TracedScheduleNode : public ConcreteScheduleNode { Optional decision = NullOpt) final; Array SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor, Optional> decision = NullOpt) final; - LoopRV SampleComputeLocation(const Array& block_rvs, - Optional decision = NullOpt) final; + LoopRV SampleComputeLocation(const BlockRV& block_rv, Optional decision = NullOpt) final; /******** Schedule: Get blocks & loops ********/ BlockRV GetBlock(const String& name, const String& func_name = "main") final; Array GetLoops(const BlockRV& block_rv) final; From 5933e52e881398d1153de71a40b14ebf828dfae2 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 10:51:28 +0800 Subject: [PATCH 07/26] Enable inlining --- .../mutator/mutate_compute_location.cc | 23 ++++++------ src/tir/schedule/analysis.h | 4 ++- src/tir/schedule/analysis/analysis.cc | 32 ++++++++++++----- src/tir/schedule/primitive/sampling.cc | 36 ++++++++++++------- 4 files changed, 62 insertions(+), 33 deletions(-) diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc index 1e387a053c20..3600ba23da85 100644 --- a/src/meta_schedule/mutator/mutate_compute_location.cc +++ b/src/meta_schedule/mutator/mutate_compute_location.cc @@ -78,21 +78,22 @@ std::vector FindCandidates(const Trace& tr ICHECK_EQ(inputs.size(), 1); tir::StmtSRef block_sref = sch->GetSRef(Downcast(inputs[0])); // Extract locations that can be computed at - Array loop_srefs = CollectComputeLocation(sch->state(), block_sref); - // std::vector locs{-2, -1}; // Todo: temporarily disable inline - std::vector locs{-1}; - for (int i = 0; i < static_cast(loop_srefs.size()); ++i) { - locs.push_back(i); - } + + Array location_srefs; + std::vector location_indices; + std::tie(location_srefs, location_indices) = CollectComputeLocation(sch->state(), block_sref); + // Remove `old_decision` - std::vector::iterator rm = std::find(locs.begin(), locs.end(), old_decision); - if (rm != locs.end()) { - locs.erase(rm); + auto it = std::find(location_indices.begin(), location_indices.end(), old_decision); + if (it != location_indices.end()) { + location_srefs.erase(location_srefs.begin() + (it - location_indices.begin())); + location_indices.erase(it); } + ICHECK_EQ(location_srefs.size(), location_indices.size()); // Add the candidate - if (!locs.empty()) { - candidates.emplace_back(inst, std::move(locs)); + if (!location_srefs.empty()) { + candidates.emplace_back(inst, std::move(location_indices)); } } return decision; diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index ecb7daa3d6ce..6d2fd3943813 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -352,8 +352,10 @@ StmtSRef GetSRefLowestCommonAncestor(const Array& srefs); * \param self The schedule state * \param block_sref The block whose compute-at locations are to be collected * \return All the feasible compute-at locations of the input block, given as an array of loop srefs + * and an array of their indices among the outer loops of the input block */ -Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref); +std::pair, std::vector> CollectComputeLocation(const ScheduleState& self, + const StmtSRef& block_sref); /******** Tensorization ********/ diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index c1065b250be8..4b2ff6ee386e 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -822,16 +822,26 @@ StmtSRef GetSRefLowestCommonAncestor(const Array& srefs) { return GetRef(p); } -Array CollectComputeLocation(const ScheduleState& self, const StmtSRef& block_sref) { +std::pair, std::vector> CollectComputeLocation(const ScheduleState& self, + const StmtSRef& block_sref) { + Array location_srefs; + std::vector location_indices; + if (CanComputeInline(self, block_sref)) { + location_srefs.push_back(StmtSRef::InlineMark()); + location_indices.push_back(-2); + } + location_srefs.push_back(StmtSRef::RootMark()); + location_indices.push_back(-1); + Array consumers = GetConsumers(self, block_sref); if (consumers.empty()) { - return {}; + return std::make_pair(location_srefs, location_indices); } StmtSRef loop_boundary_sref = consumers.size() > 1 ? GetSRefLowestCommonAncestor(consumers) : GetRef(consumers[0]->parent); if (loop_boundary_sref->StmtAs() != nullptr) { - return {}; + return std::make_pair(location_srefs, location_indices); } Array loop_srefs = GetLoops(consumers[0]); @@ -842,8 +852,6 @@ Array CollectComputeLocation(const ScheduleState& self, const StmtSRef std::vector loop_iter_types; loop_iter_types.reserve(lca_pos + 1); int i_last_datapar = -1; - bool visited_reduce = false; - for (int i = 0; i <= lca_pos; ++i) { IterVarType iter_type = GetLoopIterType(loop_srefs[i]); loop_iter_types.push_back(iter_type); @@ -852,25 +860,31 @@ Array CollectComputeLocation(const ScheduleState& self, const StmtSRef } } + location_srefs.reserve(lca_pos + 3); + location_indices.reserve(lca_pos + 3); + bool visited_reduce = false; + // Todo 1: take the reduction iterators of the input block into considertion // Todo 2: skip unit loops for (int i = 0; i <= lca_pos; ++i) { if (loop_iter_types[i] == IterVarType::kDataPar) { if (visited_reduce) { - return {loop_srefs.begin(), loop_srefs.begin() + i}; + break; } } else if (loop_iter_types[i] == IterVarType::kCommReduce) { visited_reduce = true; if (i > i_last_datapar) { - return {loop_srefs.begin(), loop_srefs.begin() + i}; + break; } } else { - return {loop_srefs.begin(), loop_srefs.begin() + i}; + break; } + location_srefs.push_back(loop_srefs[i]); + location_indices.push_back(i); } - return {loop_srefs.begin(), loop_srefs.begin() + lca_pos + 1}; + return std::make_pair(location_srefs, location_indices); } /******** Tensorization ********/ diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index 86989097a6c0..45f27dcaf0f7 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -349,26 +349,38 @@ tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, const StmtSRef& block_sref, Optional* decision) { // Find all possible compute-at locations - Array loop_srefs = CollectComputeLocation(self, block_sref); - int n = loop_srefs.size(); + Array location_srefs; + std::vector location_indices; + std::tie(location_srefs, location_indices) = CollectComputeLocation(self, block_sref); + ICHECK_EQ(location_srefs.size(), location_indices.size()); + // The decision made, by default it is -1 - int i = -1; if (decision->defined()) { // Handle existing decision - int64_t val_decision = decision->as()->value; - if (val_decision >= n) { - LOG(WARNING) << "old decision is " << val_decision << " while current candidate count is " - << n << ". Hence cannot reapply the old decision"; - i = n - 1; + int64_t old_decision = Downcast(*decision)->value; + auto it = std::lower_bound(location_indices.begin(), location_indices.end(), old_decision); + int idx = it - location_indices.begin(); + + if (it != location_indices.end() && *it == old_decision) { + *decision = Integer(old_decision); + return location_srefs[idx]; + } else if (it != location_indices.begin()) { + *decision = Integer(*--it); + LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " + << *decision; + return location_srefs[idx - 1]; } else { - i = val_decision; + *decision = Integer(-1); + LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " + << *decision; + return StmtSRef::RootMark(); } } else { // Sample possible combinations - i = SampleInt(rand_state, -1, n); // Todo: Temporarily disable Inline + int sampled_idx = SampleInt(rand_state, 0, location_indices.size()); + *decision = Integer(location_indices[sampled_idx]); + return location_srefs[sampled_idx]; } - *decision = Integer(i); - return i >= 0 ? loop_srefs[i] : i == -1 ? tir::StmtSRef::RootMark() : tir::StmtSRef::InlineMark(); } /******** InstructionKind Registration ********/ From f409d7a2d39fd66ecced37572671e2395ca8b685 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 11:26:12 +0800 Subject: [PATCH 08/26] Skip unit loops --- src/tir/schedule/analysis/analysis.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 4b2ff6ee386e..95b0bb4ac795 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -865,9 +865,13 @@ std::pair, std::vector> CollectComputeLocation(const Schedu bool visited_reduce = false; // Todo 1: take the reduction iterators of the input block into considertion - // Todo 2: skip unit loops for (int i = 0; i <= lca_pos; ++i) { + const int64_t* loop_extent = GetLoopIntExtent(loop_srefs[i]); + if (loop_extent != nullptr && *loop_extent == 1) { + continue; + } + if (loop_iter_types[i] == IterVarType::kDataPar) { if (visited_reduce) { break; From 50d0f3a776c839d7a85a716607dc0b929098cf2f Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 14:02:46 +0800 Subject: [PATCH 09/26] Take reduction block iterators into consideration --- src/tir/schedule/analysis.h | 7 +++++++ src/tir/schedule/analysis/analysis.cc | 27 ++++++++++++++++++++------- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 6d2fd3943813..9083dd9e15a1 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -667,6 +667,13 @@ bool HasOp(const Stmt& stmt, const Array& ops); */ bool HasIfThenElse(const Stmt& stmt); +/*! + * \brief Get the number of continuously leading data-parallel block iterator of the given block + * \param block_sref The sref of the block to be queried + * \return The number of continuously leading data-parallel block iterator of the input block + */ +int GetNumOfLeadingDataParIter(const StmtSRef& block_sref); + /******** Storage Scope ********/ /*! diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 95b0bb4ac795..3f904d2f3af9 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -845,14 +845,17 @@ std::pair, std::vector> CollectComputeLocation(const Schedu } Array loop_srefs = GetLoops(consumers[0]); + int lca_pos = std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary_sref) - loop_srefs.begin(); ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); + int n_leading_datapar_iter = GetNumOfLeadingDataParIter(block_sref); + int n_candidate = std::min(lca_pos + 1, n_leading_datapar_iter); std::vector loop_iter_types; - loop_iter_types.reserve(lca_pos + 1); + loop_iter_types.reserve(n_candidate); int i_last_datapar = -1; - for (int i = 0; i <= lca_pos; ++i) { + for (int i = 0; i < n_candidate; ++i) { IterVarType iter_type = GetLoopIterType(loop_srefs[i]); loop_iter_types.push_back(iter_type); if (iter_type == IterVarType::kDataPar) { @@ -860,13 +863,11 @@ std::pair, std::vector> CollectComputeLocation(const Schedu } } - location_srefs.reserve(lca_pos + 3); - location_indices.reserve(lca_pos + 3); + location_srefs.reserve(n_candidate + 2); + location_indices.reserve(n_candidate + 2); bool visited_reduce = false; - // Todo 1: take the reduction iterators of the input block into considertion - - for (int i = 0; i <= lca_pos; ++i) { + for (int i = 0; i < n_candidate; ++i) { const int64_t* loop_extent = GetLoopIntExtent(loop_srefs[i]); if (loop_extent != nullptr && *loop_extent == 1) { continue; @@ -2085,6 +2086,18 @@ bool HasIfThenElse(const Stmt& stmt) { return has_branch; } +int GetNumOfLeadingDataParIter(const StmtSRef& block_sref) { + const BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); + int n_iter = static_cast(block->iter_vars.size()); + + for (int i = 0; i < n_iter; ++i) { + if (block->iter_vars[i]->iter_type != kDataPar) { + return i; + } + } + return n_iter; +} + /******** Storage Scope ********/ void CheckStorageScope(const ScheduleState& self, String storage_scope) { From 14c279f5a5b249a6ac486e38c9837438ba1d1cde Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 20:11:41 +0800 Subject: [PATCH 10/26] [TIR] For-kind inheritance in decompose-reduction --- src/tir/schedule/primitive/reduction.cc | 2 +- .../unittest/test_tir_schedule_reduction.py | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/tir/schedule/primitive/reduction.cc b/src/tir/schedule/primitive/reduction.cc index f3d398c06d92..6c174560e954 100644 --- a/src/tir/schedule/primitive/reduction.cc +++ b/src/tir/schedule/primitive/reduction.cc @@ -281,7 +281,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, body = For(/*loop_var=*/new_loop_var, /*min=*/old_loop->min, /*extent=*/old_loop->extent, - /*kind=*/ForKind::kSerial, + /*kind=*/old_loop->kind, /*body=*/body); } body = Substitute(body, loop_var_map); diff --git a/tests/python/unittest/test_tir_schedule_reduction.py b/tests/python/unittest/test_tir_schedule_reduction.py index 5f5daa144e96..5ad366b2fa02 100644 --- a/tests/python/unittest/test_tir_schedule_reduction.py +++ b/tests/python/unittest/test_tir_schedule_reduction.py @@ -185,6 +185,34 @@ def matmul_decompose_with_annotation(a: T.handle, b: T.handle, c: T.handle) -> N C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk] +@T.prim_func +def colsum_with_vectorization(a: T.handle, b: T.handle) -> None: + A = T.match_buffer(a, [128, 32], dtype="float32") + B = T.match_buffer(b, [32], dtype="float32") + for k in T.serial(0, 128): + for i in T.vectorized(0, 32): + with T.block("B"): + vk, vi = T.axis.remap("RS", [k, i]) + with T.init(): + B[vi] = T.float32(0) + B[vi] = B[vi] + A[vk, vi] + + +@T.prim_func +def colsum_decompose_with_vectorization(a: T.handle, b: T.handle) -> None: + A = T.match_buffer(a, [128, 32], dtype="float32") + B = T.match_buffer(b, [32], dtype="float32") + for i in T.vectorized(0, 32): + with T.block("B_init"): + vi = T.axis.S(32, i) + B[vi] = T.float32(0) + for k in T.serial(0, 128): + for i in T.vectorized(0, 32): + with T.block("B"): + vk, vi = T.axis.remap("RS", [k, i]) + B[vi] = B[vi] + A[vk, vi] + + # pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg @@ -243,5 +271,16 @@ def test_reduction_decompose_with_annotation(): verify_trace_roundtrip(s, mod=matmul_with_annotation) +def test_reduction_decompose_with_different_for_kind(): + s = tir.Schedule(colsum_with_vectorization, debug_mask="all") + B = s.get_block("B") + k, _ = s.get_loops(B) + B_init = s.decompose_reduction(B, k) + tvm.ir.assert_structural_equal(s.mod["main"], colsum_decompose_with_vectorization) + assert s.get(B).same_as(s.get(s.get_block("B_update"))) + assert s.get(B_init).same_as(s.get(s.get_block("B_init"))) + verify_trace_roundtrip(s, mod=colsum_with_vectorization) + + if __name__ == "__main__": sys.exit(pytest.main([__file__] + sys.argv[1:])) From 70714c824ac030e7c2f9de63bdbb72232188839c Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 30 Dec 2021 22:22:53 +0800 Subject: [PATCH 11/26] Complete MutatorComputeLocation with test --- .../mutator/mutate_compute_location.py | 5 +- .../mutator/mutate_compute_location.cc | 67 +++++++++---------- ...chedule_mutator_mutate_compute_location.py | 5 +- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/python/tvm/meta_schedule/mutator/mutate_compute_location.py b/python/tvm/meta_schedule/mutator/mutate_compute_location.py index 5223d2a65b1a..bb361247bf62 100644 --- a/python/tvm/meta_schedule/mutator/mutate_compute_location.py +++ b/python/tvm/meta_schedule/mutator/mutate_compute_location.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Mutator that mutates the outcome of SampleComputeLocation""" +"""A mutator that mutates the compute-at location decision of SampleComputeLocation""" from tvm._ffi.registry import register_object from .. import _ffi_api @@ -23,10 +23,9 @@ @register_object("meta_schedule.MutateComputeLocation") class MutateComputeLocation(Mutator): - """Mutator thatmutates the outcome of SampleComputeLocation""" + """A mutator that mutates the compute-at location decision of SampleComputeLocation""" def __init__(self) -> None: - """Mutator that mutates the outcome of SampleComputeLocation""" self.__init_handle_by_constructor__( _ffi_api.MutatorMutateComputeLocation, # type: ignore # pylint: disable=no-member ) diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc index 3600ba23da85..9b3f2265171e 100644 --- a/src/meta_schedule/mutator/mutate_compute_location.cc +++ b/src/meta_schedule/mutator/mutate_compute_location.cc @@ -25,7 +25,7 @@ using tir::Instruction; using tir::InstructionKind; using tir::Trace; -/*! \brief Create a Mutator that mutates auto unroll step */ // Todo +/*! \brief A mutator that mutates the compute-at location decision of SampleComputeLocation */ class MutateComputeLocationNode : public MutatorNode { public: /*! \brief JSON representation of the workload */ @@ -36,7 +36,18 @@ class MutateComputeLocationNode : public MutatorNode { TVM_DECLARE_FINAL_OBJECT_INFO(MutateComputeLocationNode, MutatorNode); public: - struct Candidate; + struct Candidate { + /*! \brief The SampleComputeLocation instruction */ + Instruction inst; + /*! \brief The candidate compute locations */ + std::vector locs; + + explicit Candidate(Instruction inst, std::vector locs) + : inst(std::move(inst)), locs(std::move(locs)) {} + }; + + std::vector FindCandidates(const Trace& trace, TRandState* rand_state); + // Inherit from `MutatorNode` void InitializeWithTuneContext(const TuneContext& context) final { this->json_mod_ = SaveJSON(context->mod.value()); @@ -45,53 +56,44 @@ class MutateComputeLocationNode : public MutatorNode { Optional Apply(const Trace& trace, TRandState* rand_state) final; }; -/*! \brief The candidate to be mutated */ -struct MutateComputeLocationNode::Candidate { - /*! \brief The SampleComputeLocation instruction */ - Instruction inst; - /*! \brief The candidate compute locations */ - std::vector locs; - - explicit Candidate(Instruction inst, std::vector locs) - : inst(std::move(inst)), locs(std::move(locs)) {} -}; - /*! - * \brief Find instruction `SampleComputeLocation` + * \brief Find all appearances of instruction `SampleComputeLocation` whose decision can be mutated + * to at lease one other value * \param trace The trace from which to find the instructions - * \param workload The workload - * \return All the candidate instructions together with the candidate compute locations + * \return All the candidate instructions together with the candidate compute-at locations */ -std::vector FindCandidates(const Trace& trace, - const tir::Schedule& sch) { +std::vector MutateComputeLocationNode::FindCandidates( + const Trace& trace, TRandState* rand_state) { + tir::Schedule sch = tir::Schedule::Traced( // + /*mod=*/Downcast(LoadJSON(this->json_mod_)), // + /*rand_state=*/ForkSeed(rand_state), // + /*debug_mode=*/0, // + /*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone); static InstructionKind inst_sample_compute_location = InstructionKind::Get("SampleComputeLocation"); std::vector candidates; - auto f_provide_decision = [&](const tir::Instruction& inst, + + auto f_provide_decision = [&](const tir::Instruction& inst, // const Array& inputs, // - const Array& attrs, + const Array& attrs, // const ObjectRef& decision) -> ObjectRef { if (inst->kind.same_as(inst_sample_compute_location)) { - // The decision made - int old_decision = Downcast(decision)->value; - // Extract the inputs + // Step 1. Extract the instruction input and the old decision. ICHECK_EQ(inputs.size(), 1); tir::StmtSRef block_sref = sch->GetSRef(Downcast(inputs[0])); - // Extract locations that can be computed at - + int old_decision = Downcast(decision)->value; + // Step 2. Collect all the compute-at locations. Array location_srefs; std::vector location_indices; std::tie(location_srefs, location_indices) = CollectComputeLocation(sch->state(), block_sref); - - // Remove `old_decision` + // Step 3. Remove the old decision. auto it = std::find(location_indices.begin(), location_indices.end(), old_decision); if (it != location_indices.end()) { location_srefs.erase(location_srefs.begin() + (it - location_indices.begin())); location_indices.erase(it); } ICHECK_EQ(location_srefs.size(), location_indices.size()); - - // Add the candidate + // Step 4. Add a new candidate if there are at least one remaining compute-at position. if (!location_srefs.empty()) { candidates.emplace_back(inst, std::move(location_indices)); } @@ -103,12 +105,7 @@ std::vector FindCandidates(const Trace& tr } Optional MutateComputeLocationNode::Apply(const Trace& trace, TRandState* rand_state) { - tir::Schedule sch = tir::Schedule::Traced( // - /*mod=*/Downcast(LoadJSON(this->json_mod_)), // - /*rand_state=*/ForkSeed(rand_state), // - /*debug_mode=*/0, - /*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone); - std::vector candidates = FindCandidates(trace, sch); + std::vector candidates = FindCandidates(trace, rand_state); if (candidates.empty()) { return NullOpt; } diff --git a/tests/python/unittest/test_meta_schedule_mutator_mutate_compute_location.py b/tests/python/unittest/test_meta_schedule_mutator_mutate_compute_location.py index 75ddeb473a1e..439a4c19dba1 100644 --- a/tests/python/unittest/test_meta_schedule_mutator_mutate_compute_location.py +++ b/tests/python/unittest/test_meta_schedule_mutator_mutate_compute_location.py @@ -57,9 +57,8 @@ def _sch(decision: int) -> Schedule: sch = Schedule(add, debug_mask="all") # pylint: disable=invalid-name b0 = sch.get_block(name="move", func_name="main") - (b1,) = sch.get_consumers(block=b0) - l2 = sch.sample_compute_location(block=b1, decision=decision) - sch.compute_at(block=b0, loop=l2, preserve_unit_loops=True) + l1 = sch.sample_compute_location(block=b0, decision=decision) + sch.compute_at(block=b0, loop=l1, preserve_unit_loops=True) # pylint: enable=invalid-name return sch From 7c9aae59f2c86686493e7e80fe5c895aa9eb9d34 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 31 Dec 2021 00:29:55 +0800 Subject: [PATCH 12/26] Complete RandomComputeLocation with test --- .../schedule_rule/random_compute_location.cc | 12 ++++-------- src/tir/schedule/analysis.h | 8 -------- src/tir/schedule/analysis/analysis.cc | 5 ----- ...schedule_schedule_rule_random_compute_location.py | 5 ++--- 4 files changed, 6 insertions(+), 24 deletions(-) diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index 5dc6f66e3e07..b05eb2186fa3 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -32,7 +32,10 @@ class RandomComputeLocationNode : public ScheduleRuleNode { return false; } // Cond 2. The block should be the direct child block of the root block. - if (!tir::IsSubrootBlock(sch->state(), block_sref)) { // Todo + if (GetScopeRoot(sch->state(), block_sref, // + /*require_stage_pipeline=*/false, // + /*require_subtree_compact_dataflow=*/false) + ->parent != nullptr) { return false; } // Cond 3 & 4. The block has at least one outer loop, and the outermost loop has only one child @@ -61,12 +64,7 @@ class RandomComputeLocationNode : public ScheduleRuleNode { return {sch}; } - int err_cnt = 0; for (;;) { - if (err_cnt == 100) { - LOG(INFO) << "err_cnt = 100, force quit"; - break; - } tir::LoopRV compute_at_loc = sch->SampleComputeLocation(block_rv); try { sch->ComputeAt(block_rv, compute_at_loc, true); @@ -74,8 +72,6 @@ class RandomComputeLocationNode : public ScheduleRuleNode { // ComputeAt fails, cleanup the following before re-try: // 1) trace: instruction & decisions // 2) sym_tab - LOG(INFO) << "fail"; - ++err_cnt; sch->trace().value()->Pop(); sch->RemoveRV(compute_at_loc); continue; diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 9083dd9e15a1..6b76ce233858 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -331,14 +331,6 @@ IterVarType GetLoopIterType(const StmtSRef& loop_sref); */ bool HasSingleChild(const StmtSRef& loop_or_block_sref); -/*! - * \brief Check if a block is the direct children of the root block - * \param self The schedule state - * \param block_sref The block to be analyzed - * \return A boolean flag indicating if the block is the subroot block - */ -bool IsSubrootBlock(const tir::ScheduleState& self, const tir::StmtSRef& block_sref); - /*! * \brief Get the lowest common ancestor of an array of blocks or loops on the sref tree * \param srefs The block srefs or loop srefs whose lowest common ancestor is to be queried diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 3f904d2f3af9..9c958a734c68 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -796,11 +796,6 @@ bool HasSingleChild(const StmtSRef& loop_or_block_sref) { return true; } -bool IsSubrootBlock(const tir::ScheduleState& self, const tir::StmtSRef& block_sref) { - tir::StmtSRef parent_block_sref = GetScopeRoot(self, block_sref, false, false); - return parent_block_sref->parent == nullptr; -} - StmtSRef GetSRefLowestCommonAncestor(const Array& srefs) { CHECK(!srefs.empty()) << "ValueError: The input array is required to have at least one sref"; diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py b/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py index 1dff38c1a826..b4d1964d3775 100644 --- a/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py +++ b/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py @@ -73,9 +73,8 @@ def test_random_compute_location(): expected = [ [ 'b0 = sch.get_block(name="move", func_name="main")', - "b1, = sch.get_consumers(block=b0)", - "l2 = sch.sample_compute_location(block=b1)", - "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=True)", + "l1 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l1, preserve_unit_loops=True)", ] ] mod = Add From cfb23635841f8e9badcd0e83cc984f6e8350c9f2 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 31 Dec 2021 00:42:53 +0800 Subject: [PATCH 13/26] Complete SampleComputeLocation in sampling.cc --- src/tir/schedule/concrete_schedule.h | 1 - src/tir/schedule/primitive/sampling.cc | 10 ++++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index d625c295079e..cacd8e389dff 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -284,7 +284,6 @@ inline Array GetSRefsHelper(const ConcreteScheduleNode* sch, const Arr return result; } -// Todo: expose it to schedule.h? inline Array ConcreteScheduleNode::GetSRefs(const Array& rvs) const { return GetSRefsHelper(this, rvs); } diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index 45f27dcaf0f7..dbfb50a65b86 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -348,15 +348,16 @@ std::vector SamplePerfectTile( tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, support::LinearCongruentialEngine::TRandState* rand_state, const StmtSRef& block_sref, Optional* decision) { - // Find all possible compute-at locations + // Step 1. Collect all possible compute-at locations. Array location_srefs; std::vector location_indices; std::tie(location_srefs, location_indices) = CollectComputeLocation(self, block_sref); ICHECK_EQ(location_srefs.size(), location_indices.size()); - // The decision made, by default it is -1 + // Step 2. If there was a previous decision, keep the decision unchanged if it exists in the + // location candidates. Otherwise, pick the location before the previous decision. + // Step 3. If there was not a previous decision, sample a decision from the collected locations. if (decision->defined()) { - // Handle existing decision int64_t old_decision = Downcast(*decision)->value; auto it = std::lower_bound(location_indices.begin(), location_indices.end(), old_decision); int idx = it - location_indices.begin(); @@ -376,11 +377,12 @@ tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, return StmtSRef::RootMark(); } } else { - // Sample possible combinations int sampled_idx = SampleInt(rand_state, 0, location_indices.size()); *decision = Integer(location_indices[sampled_idx]); return location_srefs[sampled_idx]; } + ICHECK(false) << "Cannot reach here"; + throw; } /******** InstructionKind Registration ********/ From 1a2883a975d75e42d610078a14b6c5fe0628d364 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 11:27:07 +0800 Subject: [PATCH 14/26] Do random-compute-location in AddRFactor --- .../schedule_rule/add_rfactor.cc | 35 ++++++++++++++++--- .../space_generator/post_order_apply.cc | 3 +- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/src/meta_schedule/schedule_rule/add_rfactor.cc b/src/meta_schedule/schedule_rule/add_rfactor.cc index c4c05b226b5c..a1c6416f13d9 100644 --- a/src/meta_schedule/schedule_rule/add_rfactor.cc +++ b/src/meta_schedule/schedule_rule/add_rfactor.cc @@ -89,11 +89,7 @@ Array AddRFactorNode::Apply(const tir::Schedule& sch, const tir:: ReorderAndFuseReductionLoops(sch, block_rv, &fused_reduce_loop, &num_spatial_loops); // Split the fused reduction loop. - Array factors; - do { - factors = sch->SamplePerfectTile(fused_reduce_loop, 2, max_innermost_factor); - } while (*tir::as_const_int(sch->Get(factors[0])) == 1 || - *tir::as_const_int(sch->Get(factors[1])) == 1); + Array factors = sch->SamplePerfectTile(fused_reduce_loop, 2, max_innermost_factor); const Array& split_loops = sch->Split(fused_reduce_loop, {factors.begin(), factors.end()}); @@ -104,6 +100,35 @@ Array AddRFactorNode::Apply(const tir::Schedule& sch, const tir:: const tir::BlockRV& block_rf = sch_tmp->RFactor(split_loop, num_spatial_loops); Array axes = sch_tmp->GetLoops(block_rf); ICHECK_GT(axes.size(), num_spatial_loops); + + for (;;) { + tir::LoopRV compute_at_loc = sch_tmp->SampleComputeLocation(block_rv); + try { + sch_tmp->ComputeAt(block_rv, compute_at_loc, true); + } catch (const dmlc::Error& e) { + // ComputeAt fails, cleanup the following before re-try: + // 1) trace: instruction & decisions + // 2) sym_tab + sch_tmp->trace().value()->Pop(); + sch_tmp->RemoveRV(compute_at_loc); + continue; + } + break; + } + for (;;) { + tir::LoopRV compute_at_loc = sch_tmp->SampleComputeLocation(block_rf); + try { + sch_tmp->ComputeAt(block_rf, compute_at_loc, true); + } catch (const dmlc::Error& e) { + // ComputeAt fails, cleanup the following before re-try: + // 1) trace: instruction & decisions + // 2) sym_tab + sch_tmp->trace().value()->Pop(); + sch_tmp->RemoveRV(compute_at_loc); + continue; + } + break; + } res.push_back(sch_tmp); } diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index 324422aa80b3..3f685407817b 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -105,9 +105,10 @@ class PostOrderApplyNode : public SpaceGeneratorNode { Array result{sch}; // Enumerate the schedule rules first because you can // always concat multiple schedule rules as one + Array all_blocks = BlockCollector::Collect(sch); for (ScheduleRule sch_rule : sch_rules_) { for (const tir::Schedule& sch : result) { - stack.emplace_back(sch, BlockCollector::Collect(sch)); + stack.emplace_back(sch, all_blocks); } result.clear(); From 8f1971798d1a1eb2090450f2e4d1cb4daaf06764 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 11:28:52 +0800 Subject: [PATCH 15/26] Comment out the warning, and disable n_leading_iter --- src/tir/schedule/analysis/analysis.cc | 3 ++- src/tir/schedule/primitive/sampling.cc | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 9c958a734c68..6912ffaa2e6f 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -845,7 +845,8 @@ std::pair, std::vector> CollectComputeLocation(const Schedu std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary_sref) - loop_srefs.begin(); ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); int n_leading_datapar_iter = GetNumOfLeadingDataParIter(block_sref); - int n_candidate = std::min(lca_pos + 1, n_leading_datapar_iter); + // int n_candidate = std::min(lca_pos + 1, n_leading_datapar_iter); + int n_candidate = lca_pos + 1; std::vector loop_iter_types; loop_iter_types.reserve(n_candidate); diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index dbfb50a65b86..048017fda77f 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -367,13 +367,13 @@ tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, return location_srefs[idx]; } else if (it != location_indices.begin()) { *decision = Integer(*--it); - LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " - << *decision; + // LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " + // << *decision; return location_srefs[idx - 1]; } else { *decision = Integer(-1); - LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " - << *decision; + // LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " + // << *decision; return StmtSRef::RootMark(); } } else { From e4f4abf599839ea50b0f78b577fa0a3aacfab79b Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 17:59:10 +0800 Subject: [PATCH 16/26] Use the annotation trick --- include/tvm/tir/stmt.h | 4 +++ .../schedule_rule/add_rfactor.cc | 29 +--------------- .../schedule_rule/random_compute_location.cc | 34 ++++++++++++++++++- src/tir/schedule/utils.h | 14 +++++++- ...meta_schedule_schedule_rule_add_rfactor.py | 2 ++ 5 files changed, 53 insertions(+), 30 deletions(-) diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index 5048ce8e5826..eb0476449fe9 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -1390,6 +1390,10 @@ constexpr const int meta_schedule_cache_type_read = 0; /*! \sa meta_schedule_cache_type */ constexpr const int meta_schedule_cache_type_write = 1; +/*! \brief Mark the block whose producer needs to be applied by rule Random-Compute-Location */ +constexpr const char* meta_schedule_random_compute_producer = + "meta_schedule.random_compute_producer"; + /*! \brief Mark auto-parallel setting on the block. */ constexpr const char* meta_schedule_parallel = "meta_schedule.parallel"; diff --git a/src/meta_schedule/schedule_rule/add_rfactor.cc b/src/meta_schedule/schedule_rule/add_rfactor.cc index a1c6416f13d9..057979412239 100644 --- a/src/meta_schedule/schedule_rule/add_rfactor.cc +++ b/src/meta_schedule/schedule_rule/add_rfactor.cc @@ -101,34 +101,7 @@ Array AddRFactorNode::Apply(const tir::Schedule& sch, const tir:: Array axes = sch_tmp->GetLoops(block_rf); ICHECK_GT(axes.size(), num_spatial_loops); - for (;;) { - tir::LoopRV compute_at_loc = sch_tmp->SampleComputeLocation(block_rv); - try { - sch_tmp->ComputeAt(block_rv, compute_at_loc, true); - } catch (const dmlc::Error& e) { - // ComputeAt fails, cleanup the following before re-try: - // 1) trace: instruction & decisions - // 2) sym_tab - sch_tmp->trace().value()->Pop(); - sch_tmp->RemoveRV(compute_at_loc); - continue; - } - break; - } - for (;;) { - tir::LoopRV compute_at_loc = sch_tmp->SampleComputeLocation(block_rf); - try { - sch_tmp->ComputeAt(block_rf, compute_at_loc, true); - } catch (const dmlc::Error& e) { - // ComputeAt fails, cleanup the following before re-try: - // 1) trace: instruction & decisions - // 2) sym_tab - sch_tmp->trace().value()->Pop(); - sch_tmp->RemoveRV(compute_at_loc); - continue; - } - break; - } + sch_tmp->Annotate(block_rv, tir::attr::meta_schedule_random_compute_producer, Bool(true)); res.push_back(sch_tmp); } diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index b05eb2186fa3..bd3ef7beb1e2 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -64,6 +64,38 @@ class RandomComputeLocationNode : public ScheduleRuleNode { return {sch}; } + // Step 1. If the producer of the input block needs a random compute-at location (specified by + // the annotation), we colect the producer first, and transform the producer block later. + // - The reason we collect the producer before transforming the input block is that, if the + // decision of Sample-Compute-Location is "compute-inline" for the input block, we can no longer + // access the input block. Hence we collect its producer ahead of time. + // - Note that only single producer is allowed in this case. + Array producers{nullptr}; + if (tir::HasAnn(sch->GetSRef(block_rv), tir::attr::meta_schedule_random_compute_producer, + Bool(true))) { + producers = sch->GetProducers(block_rv); + sch->Unannotate(block_rv, tir::attr::meta_schedule_random_compute_producer); + ICHECK_EQ(producers.size(), 1); + } + + // Step 2. Transform the input block. + tir::Schedule res = RandomlyComputeAt(sch, block_rv); + + // Step 3. Transform the producer block if compute-location sampling is needed. + if (producers.defined()) { + res = RandomlyComputeAt(res, producers[0]); + } + + return {res}; + } + + /*! + * \brief Keep sampling a compute-at location for the input block until success. + * \param sch The TIR schedule + * \param block_rv The block whose compute-at location is to be sampled + * \return The TIR schedule after transformation + */ + tir::Schedule RandomlyComputeAt(const tir::Schedule& sch, const tir::BlockRV& block_rv) { for (;;) { tir::LoopRV compute_at_loc = sch->SampleComputeLocation(block_rv); try { @@ -78,7 +110,7 @@ class RandomComputeLocationNode : public ScheduleRuleNode { } break; } - return {sch}; + return sch; } public: diff --git a/src/tir/schedule/utils.h b/src/tir/schedule/utils.h index b3dc655e18b1..f9787d29b250 100644 --- a/src/tir/schedule/utils.h +++ b/src/tir/schedule/utils.h @@ -356,7 +356,7 @@ inline Optional GetAnn(const StmtSRef& sref, const String& ann_key) * \brief Check if a Block/For has a specific pair of annotation key and values * \param sref The sref to the block or the for loop * \param ann_key The annotation key to be checked - * \param ann_val The annotation value to be checked + * \param ann_val The string annotation value to be checked * \return Whether a Block/For has a specific pair of annotation key and values */ inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const String& ann_val) { @@ -364,6 +364,18 @@ inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const String& an return result.defined() && result.value() == ann_val; } +/*! + * \brief Check if a Block/For has a specific pair of annotation key and values + * \param sref The sref to the block or the for loop + * \param ann_key The annotation key to be checked + * \param ann_val The boolean annotation value to be checked + * \return Whether a Block/For has a specific pair of annotation key and values + */ +inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const Bool& ann_val) { + Optional result = GetAnn(sref, ann_key); + return result.defined() && result.value()->value == ann_val->value; +} + /******** Tensorization ******/ /*! * \brief Rewrite the block's outer loops to match the tensor intrin diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_add_rfactor.py b/tests/python/unittest/test_meta_schedule_schedule_rule_add_rfactor.py index 300b5aeedf3b..5a8031220354 100644 --- a/tests/python/unittest/test_meta_schedule_schedule_rule_add_rfactor.py +++ b/tests/python/unittest/test_meta_schedule_schedule_rule_add_rfactor.py @@ -48,6 +48,7 @@ def test_cpu_matmul(): "v4, v5 = sch.sample_perfect_tile(loop=l3, n=2, max_innermost_factor=64)", "l6, l7 = sch.split(loop=l3, factors=[v4, v5])", "b8 = sch.rfactor(loop=l7, factor_axis=2)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', ], [ 'b0 = sch.get_block(name="C", func_name="main")', @@ -55,6 +56,7 @@ def test_cpu_matmul(): "v4, v5 = sch.sample_perfect_tile(loop=l3, n=2, max_innermost_factor=64)", "l6, l7 = sch.split(loop=l3, factors=[v4, v5])", "b8 = sch.rfactor(loop=l6, factor_axis=2)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', ], ] target = Target("llvm --num-cores=32") From 5d5aa5380830a4b419bb7c6c0ea70e36cdf14fc5 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 19:27:37 +0800 Subject: [PATCH 17/26] Refactor and add docstring for CollectComputeLocation --- src/tir/schedule/analysis.h | 7 ---- src/tir/schedule/analysis/analysis.cc | 46 ++++++++++++--------------- 2 files changed, 20 insertions(+), 33 deletions(-) diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 6b76ce233858..ed2507c76b8c 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -659,13 +659,6 @@ bool HasOp(const Stmt& stmt, const Array& ops); */ bool HasIfThenElse(const Stmt& stmt); -/*! - * \brief Get the number of continuously leading data-parallel block iterator of the given block - * \param block_sref The sref of the block to be queried - * \return The number of continuously leading data-parallel block iterator of the input block - */ -int GetNumOfLeadingDataParIter(const StmtSRef& block_sref); - /******** Storage Scope ********/ /*! diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index 6912ffaa2e6f..faf577cb4ecd 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -821,6 +821,9 @@ std::pair, std::vector> CollectComputeLocation(const Schedu const StmtSRef& block_sref) { Array location_srefs; std::vector location_indices; + + // Step 1. Add the "compute-root" candidate. Add the "compute-inline" candidate if the block can + // be inlined. if (CanComputeInline(self, block_sref)) { location_srefs.push_back(StmtSRef::InlineMark()); location_indices.push_back(-2); @@ -828,26 +831,28 @@ std::pair, std::vector> CollectComputeLocation(const Schedu location_srefs.push_back(StmtSRef::RootMark()); location_indices.push_back(-1); + // Step 2. If the block has no consumer, there is no more candidate. Array consumers = GetConsumers(self, block_sref); if (consumers.empty()) { return std::make_pair(location_srefs, location_indices); } - - StmtSRef loop_boundary_sref = consumers.size() > 1 ? GetSRefLowestCommonAncestor(consumers) - : GetRef(consumers[0]->parent); - if (loop_boundary_sref->StmtAs() != nullptr) { + // Step 3. Get the deepest loop that the input block can be computed at (namely "boundary"). If + // such a loop cannot be found, there is no more candidate and we just return. + StmtSRef loop_boundary = consumers.size() > 1 ? GetSRefLowestCommonAncestor(consumers) + : GetRef(consumers[0]->parent); + if (loop_boundary->StmtAs() != nullptr) { return std::make_pair(location_srefs, location_indices); } + // Step 4. Collect the loops outside the first consumer and locate the boundary loop. The position + // of the boundary loop reveals the number of possible additional candidates. Array loop_srefs = GetLoops(consumers[0]); - - int lca_pos = - std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary_sref) - loop_srefs.begin(); + int lca_pos = std::find(loop_srefs.begin(), loop_srefs.end(), loop_boundary) - loop_srefs.begin(); ICHECK_LT(lca_pos, static_cast(loop_srefs.size())); - int n_leading_datapar_iter = GetNumOfLeadingDataParIter(block_sref); - // int n_candidate = std::min(lca_pos + 1, n_leading_datapar_iter); int n_candidate = lca_pos + 1; + // Step 5. Find the position of the deepest data-parallel loop among the candidate loops. This + // position is used for removing the unwanted candidates from the perspective of performance. std::vector loop_iter_types; loop_iter_types.reserve(n_candidate); int i_last_datapar = -1; @@ -858,11 +863,13 @@ std::pair, std::vector> CollectComputeLocation(const Schedu i_last_datapar = i; } } - + // Step 6. Check and add the candidates in turn according to the following rules: + // - skip the unit loops (loops with extent 1); + // - do not consider the data-parallel loops after a not-data-parallel loop; + // - do not consider the trailing not-data-parallel loops. location_srefs.reserve(n_candidate + 2); location_indices.reserve(n_candidate + 2); bool visited_reduce = false; - for (int i = 0; i < n_candidate; ++i) { const int64_t* loop_extent = GetLoopIntExtent(loop_srefs[i]); if (loop_extent != nullptr && *loop_extent == 1) { @@ -873,14 +880,13 @@ std::pair, std::vector> CollectComputeLocation(const Schedu if (visited_reduce) { break; } - } else if (loop_iter_types[i] == IterVarType::kCommReduce) { + } else { visited_reduce = true; if (i > i_last_datapar) { break; } - } else { - break; } + location_srefs.push_back(loop_srefs[i]); location_indices.push_back(i); } @@ -2082,18 +2088,6 @@ bool HasIfThenElse(const Stmt& stmt) { return has_branch; } -int GetNumOfLeadingDataParIter(const StmtSRef& block_sref) { - const BlockNode* block = TVM_SREF_TO_BLOCK(block, block_sref); - int n_iter = static_cast(block->iter_vars.size()); - - for (int i = 0; i < n_iter; ++i) { - if (block->iter_vars[i]->iter_type != kDataPar) { - return i; - } - } - return n_iter; -} - /******** Storage Scope ********/ void CheckStorageScope(const ScheduleState& self, String storage_scope) { From 945dc08ab801ef31409f6fc36d5fe624420edf68 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 21:55:11 +0800 Subject: [PATCH 18/26] Minor --- src/meta_schedule/mutator/mutate_compute_location.cc | 2 +- src/meta_schedule/schedule_rule/add_rfactor.cc | 2 ++ src/tir/schedule/primitive/sampling.cc | 4 ---- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc index 9b3f2265171e..9c495e1c50cd 100644 --- a/src/meta_schedule/mutator/mutate_compute_location.cc +++ b/src/meta_schedule/mutator/mutate_compute_location.cc @@ -39,7 +39,7 @@ class MutateComputeLocationNode : public MutatorNode { struct Candidate { /*! \brief The SampleComputeLocation instruction */ Instruction inst; - /*! \brief The candidate compute locations */ + /*! \brief The candidate compute-at locations */ std::vector locs; explicit Candidate(Instruction inst, std::vector locs) diff --git a/src/meta_schedule/schedule_rule/add_rfactor.cc b/src/meta_schedule/schedule_rule/add_rfactor.cc index 057979412239..75bb47c23dc3 100644 --- a/src/meta_schedule/schedule_rule/add_rfactor.cc +++ b/src/meta_schedule/schedule_rule/add_rfactor.cc @@ -101,6 +101,8 @@ Array AddRFactorNode::Apply(const tir::Schedule& sch, const tir:: Array axes = sch_tmp->GetLoops(block_rf); ICHECK_GT(axes.size(), num_spatial_loops); + // Annotate that the rfactor block, which is now the producer of the original block, needs to be + // considered by the rule Random-Compute-Location. sch_tmp->Annotate(block_rv, tir::attr::meta_schedule_random_compute_producer, Bool(true)); res.push_back(sch_tmp); } diff --git a/src/tir/schedule/primitive/sampling.cc b/src/tir/schedule/primitive/sampling.cc index 048017fda77f..cbb4e66918e9 100644 --- a/src/tir/schedule/primitive/sampling.cc +++ b/src/tir/schedule/primitive/sampling.cc @@ -367,13 +367,9 @@ tir::StmtSRef SampleComputeLocation(tir::ScheduleState self, return location_srefs[idx]; } else if (it != location_indices.begin()) { *decision = Integer(*--it); - // LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " - // << *decision; return location_srefs[idx - 1]; } else { *decision = Integer(-1); - // LOG(WARNING) << "old decision " << old_decision << " is outdated. Change the decision to " - // << *decision; return StmtSRef::RootMark(); } } else { From d6367ebc76120698b940276a4ff268312dd1c6fd Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Sun, 2 Jan 2022 21:55:29 +0800 Subject: [PATCH 19/26] Test for SampleComputeLocation --- .../unittest/test_tir_schedule_sampling.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/python/unittest/test_tir_schedule_sampling.py b/tests/python/unittest/test_tir_schedule_sampling.py index 5d2676e41d1c..e46ef2899af2 100644 --- a/tests/python/unittest/test_tir_schedule_sampling.py +++ b/tests/python/unittest/test_tir_schedule_sampling.py @@ -37,6 +37,28 @@ def elementwise(a: T.handle, b: T.handle) -> None: B[vi, vj, vk] = A[vi, vj, vk] * 2.0 +@T.prim_func +def tiled_conv2d_with_padding(inputs: T.Buffer[(1, 224, 224, 3), "float32"], weight: T.Buffer[(7, 7, 3, 64), "float32"], conv2d_nhwc: T.Buffer[(1, 112, 112, 64), "float32"]) -> None: + PadInput = T.alloc_buffer([1, 230, 230, 3], dtype="float32") + for i0, i1, i2, i3 in T.grid(1, 230, 230, 3): + with T.block("PadInput"): + i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3]) + T.reads(inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1]) + T.writes(PadInput[i0_1, i1_1, i2_1, i3_1]) + PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(3 <= i1_1 and i1_1 < 227 and 3 <= i2_1 and i2_1 < 227, inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1], T.float32(0), dtype="float32") + for i0_0, i1_0, i2_0, i3_0, i0_1_1, i1_1_1, i2_1_1, i3_1_1, i4_0, i5_0, i6_0, i0_2, i1_2, i2_2, i3_2, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3, i3_3 in T.grid(1, 1, 4, 1, 1, 2, 4, 1, 7, 7, 1, 1, 1, 1, 1, 1, 1, 3, 1, 56, 7, 64): + with T.block("conv2d_nhwc"): + n = T.axis.spatial(1, 0) + h = T.axis.spatial(112, i1_1_1 * 56 + i1_3) + w = T.axis.spatial(112, i2_0 * 28 + i2_1_1 * 7 + i2_3) + co, rh, rw, rc = T.axis.remap("SRRR", [i3_3, i4_0, i5_0, i6_1]) + T.reads(conv2d_nhwc[n, h, w, co], PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc], weight[rh, rw, rc, co]) + T.writes(conv2d_nhwc[n, h, w, co]) + with T.init(): + conv2d_nhwc[n, h, w, co] = T.float32(0) + conv2d_nhwc[n, h, w, co] = conv2d_nhwc[n, h, w, co] + PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc] * weight[rh, rw, rc, co] + + # pylint: enable=no-member,invalid-name,unused-variable @@ -116,5 +138,22 @@ def test_sample_perfect_tile_composite(): verify_trace_roundtrip(sch, mod=elementwise) +def test_sample_compute_location(): + n = 100 + sch = tir.Schedule(tiled_conv2d_with_padding, seed=42, debug_mask="all") + pad_input = sch.get_block("PadInput") + decision_dict = dict() + for _ in range(n): + _ = sch.sample_compute_location(pad_input) # pylint: disable=invalid-name + decision = sch.trace.decisions[sch.trace.insts[-1]] + decision_dict[decision] = decision_dict[decision] + 1 if decision in decision_dict else 1 + + n_candidates = 8 + expected_rate = 1.0 / n_candidates + for _, cnt in decision_dict.items(): + assert (expected_rate - 0.03) * n <= cnt <= (expected_rate + 0.03) * n + + + if __name__ == "__main__": sys.exit(pytest.main([__file__] + sys.argv[1:])) From 93cac6d574e2181c100d513f8fd31feb20780d1a Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 7 Jan 2022 15:58:17 +0800 Subject: [PATCH 20/26] Annotate the tiling structure --- include/tvm/tir/stmt.h | 3 +++ src/meta_schedule/schedule_rule/multi_level_tiling.cc | 2 ++ ...t_meta_schedule_schedule_rule_multi_level_tiling.py | 10 ++++++++++ tests/python/unittest/test_meta_schedule_sketch_cpu.py | 6 ++++++ .../python/unittest/test_meta_schedule_sketch_cuda.py | 4 ++++ tests/python/unittest/test_tir_schedule_sampling.py | 4 ++-- 6 files changed, 27 insertions(+), 2 deletions(-) diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index eb0476449fe9..20ad447a9b3c 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -1390,6 +1390,9 @@ constexpr const int meta_schedule_cache_type_read = 0; /*! \sa meta_schedule_cache_type */ constexpr const int meta_schedule_cache_type_write = 1; +/*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */ +constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure"; + /*! \brief Mark the block whose producer needs to be applied by rule Random-Compute-Location */ constexpr const char* meta_schedule_random_compute_producer = "meta_schedule.random_compute_producer"; diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/meta_schedule/schedule_rule/multi_level_tiling.cc index eac907771c3b..a0ffe7e00426 100644 --- a/src/meta_schedule/schedule_rule/multi_level_tiling.cc +++ b/src/meta_schedule/schedule_rule/multi_level_tiling.cc @@ -293,6 +293,8 @@ class MultiLevelTilingNode : public ScheduleRuleNode { if (!NeedsMultiLevelTiling(sch->state(), sch->GetSRef(block_rv))) { return {sch}; } + sch->Annotate(block_rv, tir::attr::meta_schedule_tiling_structure, structure); + std::vector states{State(sch, block_rv)}; states = SubRule(std::move(states), [&](State state) { return DetectTensorCore(state); }); states = SubRule(std::move(states), [&](State state) { return AddWriteReuse(state); }); 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 03488ed5fe99..dd703e49ff0e 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 @@ -47,6 +47,7 @@ def test_cpu_matmul(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', 'b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", @@ -60,6 +61,7 @@ def test_cpu_matmul(): ], [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', 'b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", @@ -73,6 +75,7 @@ def test_cpu_matmul(): ], [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "l1, l2, l3 = sch.get_loops(block=b0)", "v4, v5, v6, v7 = sch.sample_perfect_tile(loop=l1, n=4, max_innermost_factor=64)", "l8, l9, l10, l11 = sch.split(loop=l1, factors=[v4, v5, v6, v7])", @@ -105,6 +108,7 @@ def test_cpu_matmul_relu(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "b1, = sch.get_consumers(block=b0)", "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", @@ -118,6 +122,7 @@ def test_cpu_matmul_relu(): ], [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "b1, = sch.get_consumers(block=b0)", "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", @@ -131,6 +136,7 @@ def test_cpu_matmul_relu(): ], [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "l1, l2, l3 = sch.get_loops(block=b0)", "v4, v5, v6, v7 = sch.sample_perfect_tile(loop=l1, n=4, max_innermost_factor=64)", "l8, l9, l10, l11 = sch.split(loop=l1, factors=[v4, v5, v6, v7])", @@ -164,6 +170,7 @@ def test_cuda_matmul(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l2, n=5, max_innermost_factor=64)", @@ -217,6 +224,7 @@ def test_cuda_matmul_relu(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l2, n=5, max_innermost_factor=64)", @@ -269,6 +277,7 @@ def test_cuda_tensor_core_matmul(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', "l1, l2, l3 = sch.get_loops(block=b0)", "l4, l5 = sch.split(loop=l1, factors=[32, 16])", "l6, l7 = sch.split(loop=l2, factors=[32, 16])", @@ -340,6 +349,7 @@ def test_cuda_tensor_core_matmul_relu(): expected = [ [ 'b0 = sch.get_block(name="C", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', "l1, l2, l3 = sch.get_loops(block=b0)", "l4, l5 = sch.split(loop=l1, factors=[32, 16])", "l6, l7 = sch.split(loop=l2, factors=[32, 16])", diff --git a/tests/python/unittest/test_meta_schedule_sketch_cpu.py b/tests/python/unittest/test_meta_schedule_sketch_cpu.py index 1c896addd408..dcf7264815ba 100644 --- a/tests/python/unittest/test_meta_schedule_sketch_cpu.py +++ b/tests/python/unittest/test_meta_schedule_sketch_cpu.py @@ -33,6 +33,7 @@ def test_meta_schedule_cpu_sketch_matmul(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", "l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8])", @@ -49,6 +50,7 @@ def test_meta_schedule_cpu_sketch_matmul(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', 'b2 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', "l3, l4, l5 = sch.get_loops(block=b0)", "v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", @@ -67,6 +69,7 @@ def test_meta_schedule_cpu_sketch_matmul(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', 'b2 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', "l3, l4, l5 = sch.get_loops(block=b0)", "v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", @@ -105,6 +108,7 @@ def test_meta_schedule_cpu_sketch_matmul_relu(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "l2, l3, l4 = sch.get_loops(block=b0)", "v5, v6, v7, v8 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", "l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8])", @@ -121,6 +125,7 @@ def test_meta_schedule_cpu_sketch_matmul_relu(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "b2, = sch.get_consumers(block=b0)", "l3, l4, l5 = sch.get_loops(block=b0)", "v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", @@ -139,6 +144,7 @@ def test_meta_schedule_cpu_sketch_matmul_relu(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', "b2, = sch.get_consumers(block=b0)", "l3, l4, l5 = sch.get_loops(block=b0)", "v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", diff --git a/tests/python/unittest/test_meta_schedule_sketch_cuda.py b/tests/python/unittest/test_meta_schedule_sketch_cuda.py index f1a5626f3f0f..86bbfecd6980 100644 --- a/tests/python/unittest/test_meta_schedule_sketch_cuda.py +++ b/tests/python/unittest/test_meta_schedule_sketch_cuda.py @@ -32,6 +32,7 @@ def test_meta_schedule_cuda_sketch_matmul(): [ 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b2 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")', "l3, l4, l5 = sch.get_loops(block=b0)", "v6, v7, v8, v9, v10 = sch.sample_perfect_tile(loop=l3, n=5, max_innermost_factor=64)", @@ -87,6 +88,7 @@ def test_meta_schedule_cuda_sketch_matmul_relu(): 'b0 = sch.get_block(name="C", func_name="main")', 'b1 = sch.get_block(name="compute", func_name="main")', 'b2 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b3 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")', "l4, l5, l6 = sch.get_loops(block=b0)", "v7, v8, v9, v10, v11 = sch.sample_perfect_tile(loop=l4, n=5, max_innermost_factor=64)", @@ -143,6 +145,7 @@ def test_meta_schedule_cuda_sketch_conv2d_nchw(): 'b0 = sch.get_block(name="pad_temp", func_name="main")', 'b1 = sch.get_block(name="compute", func_name="main")', 'b2 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b3 = sch.cache_write(block=b1, write_buffer_index=0, storage_scope="local")', "l4, l5, l6, l7, l8, l9, l10 = sch.get_loops(block=b1)", "v11, v12, v13, v14, v15 = sch.sample_perfect_tile(loop=l4, n=5, max_innermost_factor=64)", @@ -218,6 +221,7 @@ def test_meta_schedule_cuda_sketch_conv2d_nchw_bias_bn_relu(): # pylint: disabl 'b4 = sch.get_block(name="bn_add", func_name="main")', 'b5 = sch.get_block(name="compute_1", func_name="main")', 'b6 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")', 'b7 = sch.cache_write(block=b1, write_buffer_index=0, storage_scope="local")', "l8, l9, l10, l11, l12, l13, l14 = sch.get_loops(block=b1)", "v15, v16, v17, v18, v19 = sch.sample_perfect_tile(loop=l8, n=5, max_innermost_factor=64)", diff --git a/tests/python/unittest/test_tir_schedule_sampling.py b/tests/python/unittest/test_tir_schedule_sampling.py index e46ef2899af2..cf1f17b8a133 100644 --- a/tests/python/unittest/test_tir_schedule_sampling.py +++ b/tests/python/unittest/test_tir_schedule_sampling.py @@ -24,7 +24,7 @@ from tvm.tir.schedule.testing import verify_trace_roundtrip -# pylint: disable=no-member,invalid-name,unused-variable +# pylint: disable=no-member,invalid-name,unused-variable,line-too-long @T.prim_func @@ -59,7 +59,7 @@ def tiled_conv2d_with_padding(inputs: T.Buffer[(1, 224, 224, 3), "float32"], wei conv2d_nhwc[n, h, w, co] = conv2d_nhwc[n, h, w, co] + PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc] * weight[rh, rw, rc, co] -# pylint: enable=no-member,invalid-name,unused-variable +# pylint: enable=no-member,invalid-name,unused-variable,line-too-long def test_sample_categorical(): From ee28c30fc3abab50c6d1446cfff3aa628719f919 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Mon, 3 Jan 2022 10:15:01 +0800 Subject: [PATCH 21/26] Skip tiled blocks in RandomComputeLocation --- .../meta_schedule/testing/schedule_rule.py | 1 + .../schedule_rule/random_compute_location.cc | 6 +- .../unittest/test_meta_schedule_sketch_cpu.py | 623 ++++++++++++++---- 3 files changed, 497 insertions(+), 133 deletions(-) diff --git a/python/tvm/meta_schedule/testing/schedule_rule.py b/python/tvm/meta_schedule/testing/schedule_rule.py index 6b4f84a4023f..dec5ab68e701 100644 --- a/python/tvm/meta_schedule/testing/schedule_rule.py +++ b/python/tvm/meta_schedule/testing/schedule_rule.py @@ -38,6 +38,7 @@ def get(target: Target) -> List[ScheduleRule]: add_rfactor(target), multi_level_tiling(target), parallel_vectorize_unroll(target), + random_compute_location(target), ] if target.kind.name == "cuda": return [ diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index bd3ef7beb1e2..8add558b59a8 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -47,7 +47,11 @@ class RandomComputeLocationNode : public ScheduleRuleNode { if (tir::GetChildBlockSRefOnSRefTree(sch->state(), loop_srefs[0]).size() > 1) { return false; } - // Cond 5. The block has at lease one consumer. + // Cond 5. The block is not tiled. We check this condition by examine the block's annotation. + if (tir::GetAnn(block_sref, tir::attr::meta_schedule_tiling_structure).defined()) { + return false; + } + // Cond 6. The block has at lease one consumer. if (tir::GetConsumers(sch->state(), sch->GetSRef(block_rv)).empty()) { return false; } diff --git a/tests/python/unittest/test_meta_schedule_sketch_cpu.py b/tests/python/unittest/test_meta_schedule_sketch_cpu.py index dcf7264815ba..d0b20a3dd104 100644 --- a/tests/python/unittest/test_meta_schedule_sketch_cpu.py +++ b/tests/python/unittest/test_meta_schedule_sketch_cpu.py @@ -181,34 +181,11 @@ def test_meta_schedule_cpu_sketch_conv2d_nchw(): # pylint: disable=line-too-long expected = [ [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="root", func_name="main")', - "l2, l3, l4, l5, l6, l7, l8 = sch.get_loops(block=b0)", - "v9, v10, v11, v12 = sch.sample_perfect_tile(loop=l2, n=4, max_innermost_factor=64)", - "l13, l14, l15, l16 = sch.split(loop=l2, factors=[v9, v10, v11, v12])", - "v17, v18, v19, v20 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", - "l21, l22, l23, l24 = sch.split(loop=l3, factors=[v17, v18, v19, v20])", - "v25, v26, v27, v28 = sch.sample_perfect_tile(loop=l4, n=4, max_innermost_factor=64)", - "l29, l30, l31, l32 = sch.split(loop=l4, factors=[v25, v26, v27, v28])", - "v33, v34, v35, v36 = sch.sample_perfect_tile(loop=l5, n=4, max_innermost_factor=64)", - "l37, l38, l39, l40 = sch.split(loop=l5, factors=[v33, v34, v35, v36])", - "v41, v42 = sch.sample_perfect_tile(loop=l6, n=2, max_innermost_factor=64)", - "l43, l44 = sch.split(loop=l6, factors=[v41, v42])", - "v45, v46 = sch.sample_perfect_tile(loop=l7, n=2, max_innermost_factor=64)", - "l47, l48 = sch.split(loop=l7, factors=[v45, v46])", - "v49, v50 = sch.sample_perfect_tile(loop=l8, n=2, max_innermost_factor=64)", - "l51, l52 = sch.split(loop=l8, factors=[v49, v50])", - "sch.reorder(l13, l21, l29, l37, l14, l22, l30, l38, l43, l47, l51, l15, l23, l31, l39, l44, l48, l52, l16, l24, l32, l40)", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', - "v53 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v53)', - ], - [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="root", func_name="main")', - 'b2 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', - "l3, l4, l5, l6, l7, l8, l9 = sch.get_loops(block=b0)", + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + "l3, l4, l5, l6, l7, l8, l9 = sch.get_loops(block=b1)", "v10, v11, v12, v13 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", "l14, l15, l16, l17 = sch.split(loop=l3, factors=[v10, v11, v12, v13])", "v18, v19, v20, v21 = sch.sample_perfect_tile(loop=l4, n=4, max_innermost_factor=64)", @@ -224,37 +201,72 @@ def test_meta_schedule_cpu_sketch_conv2d_nchw(): "v50, v51 = sch.sample_perfect_tile(loop=l9, n=2, max_innermost_factor=64)", "l52, l53 = sch.split(loop=l9, factors=[v50, v51])", "sch.reorder(l14, l22, l30, l38, l15, l23, l31, l39, l44, l48, l52, l16, l24, l32, l40, l45, l49, l53, l17, l25, l33, l41)", - "sch.reverse_compute_at(block=b2, loop=l38, preserve_unit_loops=True)", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.vectorize", ann_val=32)', "v54 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v54)', + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.unroll_explicit", ann_val=v54)', + "l55 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l55, preserve_unit_loops=True)", ], [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="root", func_name="main")', - 'b2 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")', - "l3, l4, l5, l6, l7, l8, l9 = sch.get_loops(block=b0)", - "v10, v11, v12, v13 = sch.sample_perfect_tile(loop=l3, n=4, max_innermost_factor=64)", - "l14, l15, l16, l17 = sch.split(loop=l3, factors=[v10, v11, v12, v13])", - "v18, v19, v20, v21 = sch.sample_perfect_tile(loop=l4, n=4, max_innermost_factor=64)", - "l22, l23, l24, l25 = sch.split(loop=l4, factors=[v18, v19, v20, v21])", - "v26, v27, v28, v29 = sch.sample_perfect_tile(loop=l5, n=4, max_innermost_factor=64)", - "l30, l31, l32, l33 = sch.split(loop=l5, factors=[v26, v27, v28, v29])", - "v34, v35, v36, v37 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", - "l38, l39, l40, l41 = sch.split(loop=l6, factors=[v34, v35, v36, v37])", - "v42, v43 = sch.sample_perfect_tile(loop=l7, n=2, max_innermost_factor=64)", - "l44, l45 = sch.split(loop=l7, factors=[v42, v43])", - "v46, v47 = sch.sample_perfect_tile(loop=l8, n=2, max_innermost_factor=64)", - "l48, l49 = sch.split(loop=l8, factors=[v46, v47])", - "v50, v51 = sch.sample_perfect_tile(loop=l9, n=2, max_innermost_factor=64)", - "l52, l53 = sch.split(loop=l9, factors=[v50, v51])", - "sch.reorder(l14, l22, l30, l38, l15, l23, l31, l39, l44, l48, l52, l16, l24, l32, l40, l45, l49, l53, l17, l25, l33, l41)", - "sch.reverse_compute_at(block=b2, loop=l39, preserve_unit_loops=True)", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', - "v54 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v54)', + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + 'b3 = sch.cache_write(block=b1, write_buffer_index=0, storage_scope="global")', + "l4, l5, l6, l7, l8, l9, l10 = sch.get_loops(block=b1)", + "v11, v12, v13, v14 = sch.sample_perfect_tile(loop=l4, n=4, max_innermost_factor=64)", + "l15, l16, l17, l18 = sch.split(loop=l4, factors=[v11, v12, v13, v14])", + "v19, v20, v21, v22 = sch.sample_perfect_tile(loop=l5, n=4, max_innermost_factor=64)", + "l23, l24, l25, l26 = sch.split(loop=l5, factors=[v19, v20, v21, v22])", + "v27, v28, v29, v30 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", + "l31, l32, l33, l34 = sch.split(loop=l6, factors=[v27, v28, v29, v30])", + "v35, v36, v37, v38 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", + "l39, l40, l41, l42 = sch.split(loop=l7, factors=[v35, v36, v37, v38])", + "v43, v44 = sch.sample_perfect_tile(loop=l8, n=2, max_innermost_factor=64)", + "l45, l46 = sch.split(loop=l8, factors=[v43, v44])", + "v47, v48 = sch.sample_perfect_tile(loop=l9, n=2, max_innermost_factor=64)", + "l49, l50 = sch.split(loop=l9, factors=[v47, v48])", + "v51, v52 = sch.sample_perfect_tile(loop=l10, n=2, max_innermost_factor=64)", + "l53, l54 = sch.split(loop=l10, factors=[v51, v52])", + "sch.reorder(l15, l23, l31, l39, l16, l24, l32, l40, l45, l49, l53, l17, l25, l33, l41, l46, l50, l54, l18, l26, l34, l42)", + "sch.reverse_compute_at(block=b3, loop=l39, preserve_unit_loops=True)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.vectorize", ann_val=32)', + "v55 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.unroll_explicit", ann_val=v55)', + "l56 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l56, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + 'b3 = sch.cache_write(block=b1, write_buffer_index=0, storage_scope="global")', + "l4, l5, l6, l7, l8, l9, l10 = sch.get_loops(block=b1)", + "v11, v12, v13, v14 = sch.sample_perfect_tile(loop=l4, n=4, max_innermost_factor=64)", + "l15, l16, l17, l18 = sch.split(loop=l4, factors=[v11, v12, v13, v14])", + "v19, v20, v21, v22 = sch.sample_perfect_tile(loop=l5, n=4, max_innermost_factor=64)", + "l23, l24, l25, l26 = sch.split(loop=l5, factors=[v19, v20, v21, v22])", + "v27, v28, v29, v30 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", + "l31, l32, l33, l34 = sch.split(loop=l6, factors=[v27, v28, v29, v30])", + "v35, v36, v37, v38 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", + "l39, l40, l41, l42 = sch.split(loop=l7, factors=[v35, v36, v37, v38])", + "v43, v44 = sch.sample_perfect_tile(loop=l8, n=2, max_innermost_factor=64)", + "l45, l46 = sch.split(loop=l8, factors=[v43, v44])", + "v47, v48 = sch.sample_perfect_tile(loop=l9, n=2, max_innermost_factor=64)", + "l49, l50 = sch.split(loop=l9, factors=[v47, v48])", + "v51, v52 = sch.sample_perfect_tile(loop=l10, n=2, max_innermost_factor=64)", + "l53, l54 = sch.split(loop=l10, factors=[v51, v52])", + "sch.reorder(l15, l23, l31, l39, l16, l24, l32, l40, l45, l49, l53, l17, l25, l33, l41, l46, l50, l54, l18, l26, l34, l42)", + "sch.reverse_compute_at(block=b3, loop=l40, preserve_unit_loops=True)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.vectorize", ann_val=32)', + "v55 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.unroll_explicit", ann_val=v55)', + "l56 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l56, preserve_unit_loops=True)", ], ] # pylint: enable=line-too-long @@ -283,46 +295,17 @@ def test_meta_schedule_cpu_sketch_conv2d_nchw_bias_bn_relu(): # pylint: disable # pylint: disable=line-too-long expected = [ [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="bias_add", func_name="main")', - 'b2 = sch.get_block(name="bn_mul", func_name="main")', - 'b3 = sch.get_block(name="bn_add", func_name="main")', - 'b4 = sch.get_block(name="root", func_name="main")', - "sch.compute_inline(block=b3)", - "sch.compute_inline(block=b2)", - "sch.compute_inline(block=b1)", - "l5, l6, l7, l8, l9, l10, l11 = sch.get_loops(block=b0)", - "v12, v13, v14, v15 = sch.sample_perfect_tile(loop=l5, n=4, max_innermost_factor=64)", - "l16, l17, l18, l19 = sch.split(loop=l5, factors=[v12, v13, v14, v15])", - "v20, v21, v22, v23 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", - "l24, l25, l26, l27 = sch.split(loop=l6, factors=[v20, v21, v22, v23])", - "v28, v29, v30, v31 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", - "l32, l33, l34, l35 = sch.split(loop=l7, factors=[v28, v29, v30, v31])", - "v36, v37, v38, v39 = sch.sample_perfect_tile(loop=l8, n=4, max_innermost_factor=64)", - "l40, l41, l42, l43 = sch.split(loop=l8, factors=[v36, v37, v38, v39])", - "v44, v45 = sch.sample_perfect_tile(loop=l9, n=2, max_innermost_factor=64)", - "l46, l47 = sch.split(loop=l9, factors=[v44, v45])", - "v48, v49 = sch.sample_perfect_tile(loop=l10, n=2, max_innermost_factor=64)", - "l50, l51 = sch.split(loop=l10, factors=[v48, v49])", - "v52, v53 = sch.sample_perfect_tile(loop=l11, n=2, max_innermost_factor=64)", - "l54, l55 = sch.split(loop=l11, factors=[v52, v53])", - "sch.reorder(l16, l24, l32, l40, l17, l25, l33, l41, l46, l50, l54, l18, l26, l34, l42, l47, l51, l55, l19, l27, l35, l43)", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.vectorize", ann_val=32)', - "v56 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.unroll_explicit", ann_val=v56)', - ], - [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="bias_add", func_name="main")', - 'b2 = sch.get_block(name="bn_mul", func_name="main")', - 'b3 = sch.get_block(name="bn_add", func_name="main")', - 'b4 = sch.get_block(name="root", func_name="main")', + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="bias_add", func_name="main")', + 'b3 = sch.get_block(name="bn_mul", func_name="main")', + 'b4 = sch.get_block(name="bn_add", func_name="main")', + 'b5 = sch.get_block(name="root", func_name="main")', + "sch.compute_inline(block=b4)", "sch.compute_inline(block=b3)", "sch.compute_inline(block=b2)", - "sch.compute_inline(block=b1)", - "b5, = sch.get_consumers(block=b0)", - "l6, l7, l8, l9, l10, l11, l12 = sch.get_loops(block=b0)", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + "l6, l7, l8, l9, l10, l11, l12 = sch.get_loops(block=b1)", "v13, v14, v15, v16 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", "l17, l18, l19, l20 = sch.split(loop=l6, factors=[v13, v14, v15, v16])", "v21, v22, v23, v24 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", @@ -338,43 +321,84 @@ def test_meta_schedule_cpu_sketch_conv2d_nchw_bias_bn_relu(): # pylint: disable "v53, v54 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", "l55, l56 = sch.split(loop=l12, factors=[v53, v54])", "sch.reorder(l17, l25, l33, l41, l18, l26, l34, l42, l47, l51, l55, l19, l27, l35, l43, l48, l52, l56, l20, l28, l36, l44)", - "sch.reverse_compute_at(block=b5, loop=l41, preserve_unit_loops=True)", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.vectorize", ann_val=32)', + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.vectorize", ann_val=32)', "v57 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.unroll_explicit", ann_val=v57)', + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.unroll_explicit", ann_val=v57)', + "l58 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l58, preserve_unit_loops=True)", ], [ - 'b0 = sch.get_block(name="compute", func_name="main")', - 'b1 = sch.get_block(name="bias_add", func_name="main")', - 'b2 = sch.get_block(name="bn_mul", func_name="main")', - 'b3 = sch.get_block(name="bn_add", func_name="main")', - 'b4 = sch.get_block(name="root", func_name="main")', + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="bias_add", func_name="main")', + 'b3 = sch.get_block(name="bn_mul", func_name="main")', + 'b4 = sch.get_block(name="bn_add", func_name="main")', + 'b5 = sch.get_block(name="root", func_name="main")', + "sch.compute_inline(block=b4)", "sch.compute_inline(block=b3)", "sch.compute_inline(block=b2)", - "sch.compute_inline(block=b1)", - "b5, = sch.get_consumers(block=b0)", - "l6, l7, l8, l9, l10, l11, l12 = sch.get_loops(block=b0)", - "v13, v14, v15, v16 = sch.sample_perfect_tile(loop=l6, n=4, max_innermost_factor=64)", - "l17, l18, l19, l20 = sch.split(loop=l6, factors=[v13, v14, v15, v16])", - "v21, v22, v23, v24 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", - "l25, l26, l27, l28 = sch.split(loop=l7, factors=[v21, v22, v23, v24])", - "v29, v30, v31, v32 = sch.sample_perfect_tile(loop=l8, n=4, max_innermost_factor=64)", - "l33, l34, l35, l36 = sch.split(loop=l8, factors=[v29, v30, v31, v32])", - "v37, v38, v39, v40 = sch.sample_perfect_tile(loop=l9, n=4, max_innermost_factor=64)", - "l41, l42, l43, l44 = sch.split(loop=l9, factors=[v37, v38, v39, v40])", - "v45, v46 = sch.sample_perfect_tile(loop=l10, n=2, max_innermost_factor=64)", - "l47, l48 = sch.split(loop=l10, factors=[v45, v46])", - "v49, v50 = sch.sample_perfect_tile(loop=l11, n=2, max_innermost_factor=64)", - "l51, l52 = sch.split(loop=l11, factors=[v49, v50])", - "v53, v54 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", - "l55, l56 = sch.split(loop=l12, factors=[v53, v54])", - "sch.reorder(l17, l25, l33, l41, l18, l26, l34, l42, l47, l51, l55, l19, l27, l35, l43, l48, l52, l56, l20, l28, l36, l44)", - "sch.reverse_compute_at(block=b5, loop=l42, preserve_unit_loops=True)", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.vectorize", ann_val=32)', - "v57 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b4, ann_key="meta_schedule.unroll_explicit", ann_val=v57)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + "b6, = sch.get_consumers(block=b1)", + "l7, l8, l9, l10, l11, l12, l13 = sch.get_loops(block=b1)", + "v14, v15, v16, v17 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", + "l18, l19, l20, l21 = sch.split(loop=l7, factors=[v14, v15, v16, v17])", + "v22, v23, v24, v25 = sch.sample_perfect_tile(loop=l8, n=4, max_innermost_factor=64)", + "l26, l27, l28, l29 = sch.split(loop=l8, factors=[v22, v23, v24, v25])", + "v30, v31, v32, v33 = sch.sample_perfect_tile(loop=l9, n=4, max_innermost_factor=64)", + "l34, l35, l36, l37 = sch.split(loop=l9, factors=[v30, v31, v32, v33])", + "v38, v39, v40, v41 = sch.sample_perfect_tile(loop=l10, n=4, max_innermost_factor=64)", + "l42, l43, l44, l45 = sch.split(loop=l10, factors=[v38, v39, v40, v41])", + "v46, v47 = sch.sample_perfect_tile(loop=l11, n=2, max_innermost_factor=64)", + "l48, l49 = sch.split(loop=l11, factors=[v46, v47])", + "v50, v51 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l52, l53 = sch.split(loop=l12, factors=[v50, v51])", + "v54, v55 = sch.sample_perfect_tile(loop=l13, n=2, max_innermost_factor=64)", + "l56, l57 = sch.split(loop=l13, factors=[v54, v55])", + "sch.reorder(l18, l26, l34, l42, l19, l27, l35, l43, l48, l52, l56, l20, l28, l36, l44, l49, l53, l57, l21, l29, l37, l45)", + "sch.reverse_compute_at(block=b6, loop=l42, preserve_unit_loops=True)", + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.vectorize", ann_val=32)', + "v58 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.unroll_explicit", ann_val=v58)', + "l59 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l59, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="compute", func_name="main")', + 'b2 = sch.get_block(name="bias_add", func_name="main")', + 'b3 = sch.get_block(name="bn_mul", func_name="main")', + 'b4 = sch.get_block(name="bn_add", func_name="main")', + 'b5 = sch.get_block(name="root", func_name="main")', + "sch.compute_inline(block=b4)", + "sch.compute_inline(block=b3)", + "sch.compute_inline(block=b2)", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.tiling_structure", ann_val="SSRSRS")', + "b6, = sch.get_consumers(block=b1)", + "l7, l8, l9, l10, l11, l12, l13 = sch.get_loops(block=b1)", + "v14, v15, v16, v17 = sch.sample_perfect_tile(loop=l7, n=4, max_innermost_factor=64)", + "l18, l19, l20, l21 = sch.split(loop=l7, factors=[v14, v15, v16, v17])", + "v22, v23, v24, v25 = sch.sample_perfect_tile(loop=l8, n=4, max_innermost_factor=64)", + "l26, l27, l28, l29 = sch.split(loop=l8, factors=[v22, v23, v24, v25])", + "v30, v31, v32, v33 = sch.sample_perfect_tile(loop=l9, n=4, max_innermost_factor=64)", + "l34, l35, l36, l37 = sch.split(loop=l9, factors=[v30, v31, v32, v33])", + "v38, v39, v40, v41 = sch.sample_perfect_tile(loop=l10, n=4, max_innermost_factor=64)", + "l42, l43, l44, l45 = sch.split(loop=l10, factors=[v38, v39, v40, v41])", + "v46, v47 = sch.sample_perfect_tile(loop=l11, n=2, max_innermost_factor=64)", + "l48, l49 = sch.split(loop=l11, factors=[v46, v47])", + "v50, v51 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l52, l53 = sch.split(loop=l12, factors=[v50, v51])", + "v54, v55 = sch.sample_perfect_tile(loop=l13, n=2, max_innermost_factor=64)", + "l56, l57 = sch.split(loop=l13, factors=[v54, v55])", + "sch.reorder(l18, l26, l34, l42, l19, l27, l35, l43, l48, l52, l56, l20, l28, l36, l44, l49, l53, l57, l21, l29, l37, l45)", + "sch.reverse_compute_at(block=b6, loop=l43, preserve_unit_loops=True)", + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.vectorize", ann_val=32)', + "v58 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b5, ann_key="meta_schedule.unroll_explicit", ann_val=v58)', + "l59 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l59, preserve_unit_loops=True)", ], ] # pylint: enable=line-too-long @@ -399,15 +423,18 @@ def test_meta_schedule_cpu_sketch_conv2d_nchw_bias_bn_relu(): # pylint: disable check_trace(spaces, expected) -def test_meta_schedule_sketch_cpu_max_pool2d_nchw(): +def test_meta_schedule_sketch_cpu_max_pool2d_nchw(): # pylint: disable=invalid-name # pylint: disable=line-too-long expected: List[List[str]] = [ [ - 'b0 = sch.get_block(name="root", func_name="main")', - 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.parallel", ann_val=256)', - 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.vectorize", ann_val=32)', - "v1 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", - 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.unroll_explicit", ann_val=v1)', + 'b0 = sch.get_block(name="pad_temp", func_name="main")', + 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', + "v2 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v2)', + "l3 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l3, preserve_unit_loops=True)", ], ] # pylint: enable=line-too-long @@ -428,9 +455,341 @@ def test_meta_schedule_sketch_cpu_max_pool2d_nchw(): check_trace(spaces, expected) +def test_meta_schedule_cpu_sketch_batchnorm(): # pylint: disable=invalid-name + # pylint: disable=line-too-long + expected = [ + [ + 'b0 = sch.get_block(name="C", func_name="main")', + 'b1 = sch.get_block(name="root", func_name="main")', + "l2, l3, l4 = sch.get_loops(block=b0)", + "l5 = sch.fuse(l3, l4)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l8, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "b12, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l13 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l13, preserve_unit_loops=True)", + "l14 = sch.sample_compute_location(block=b12)", + "sch.compute_at(block=b12, loop=l14, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="C", func_name="main")', + 'b1 = sch.get_block(name="root", func_name="main")', + "l2, l3, l4 = sch.get_loops(block=b0)", + "l5 = sch.fuse(l3, l4)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l9, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "b12, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l13 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l13, preserve_unit_loops=True)", + "l14 = sch.sample_compute_location(block=b12)", + "sch.compute_at(block=b12, loop=l14, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="C", func_name="main")', + 'b1 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.vectorize", ann_val=32)', + "v2 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v2)', + "l3 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l3, preserve_unit_loops=True)", + ], + ] + # pylint: enable=line-too-long + ctx = create_context( + create_prim_func(te_workload.norm_bmn(B=1, M=256, N=256)), + target=_target(), + ) + spaces = ctx.space_generator.generate_design_space(mod=ctx.mod) + assert len(spaces) == 3 + check_trace(spaces, expected) + + +def test_meta_schedule_cpu_sketch_softmax(): # pylint: disable=invalid-name + # pylint: disable=line-too-long + expected = [ + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l8, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + "l11, l12 = sch.get_loops(block=b0)", + "v13, v14 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l15, l16 = sch.split(loop=l12, factors=[v13, v14])", + "b17 = sch.rfactor(loop=l15, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v18 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v18)', + "b19, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l20 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l20, preserve_unit_loops=True)", + "l21 = sch.sample_compute_location(block=b19)", + "sch.compute_at(block=b19, loop=l21, preserve_unit_loops=True)", + "l22 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l22, preserve_unit_loops=True)", + "b23, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l24 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l24, preserve_unit_loops=True)", + "l25 = sch.sample_compute_location(block=b23)", + "sch.compute_at(block=b23, loop=l25, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l8, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + "l11, l12 = sch.get_loops(block=b0)", + "v13, v14 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l15, l16 = sch.split(loop=l12, factors=[v13, v14])", + "b17 = sch.rfactor(loop=l16, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v18 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v18)', + "b19, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l20 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l20, preserve_unit_loops=True)", + "l21 = sch.sample_compute_location(block=b19)", + "sch.compute_at(block=b19, loop=l21, preserve_unit_loops=True)", + "l22 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l22, preserve_unit_loops=True)", + "b23, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l24 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l24, preserve_unit_loops=True)", + "l25 = sch.sample_compute_location(block=b23)", + "sch.compute_at(block=b23, loop=l25, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l8, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "b12, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l13 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l13, preserve_unit_loops=True)", + "l14 = sch.sample_compute_location(block=b12)", + "sch.compute_at(block=b12, loop=l14, preserve_unit_loops=True)", + "l15 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l15, preserve_unit_loops=True)", + "l16 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l16, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l9, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + "l11, l12 = sch.get_loops(block=b0)", + "v13, v14 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l15, l16 = sch.split(loop=l12, factors=[v13, v14])", + "b17 = sch.rfactor(loop=l15, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v18 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v18)', + "b19, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l20 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l20, preserve_unit_loops=True)", + "l21 = sch.sample_compute_location(block=b19)", + "sch.compute_at(block=b19, loop=l21, preserve_unit_loops=True)", + "l22 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l22, preserve_unit_loops=True)", + "b23, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l24 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l24, preserve_unit_loops=True)", + "l25 = sch.sample_compute_location(block=b23)", + "sch.compute_at(block=b23, loop=l25, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l9, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + "l11, l12 = sch.get_loops(block=b0)", + "v13, v14 = sch.sample_perfect_tile(loop=l12, n=2, max_innermost_factor=64)", + "l15, l16 = sch.split(loop=l12, factors=[v13, v14])", + "b17 = sch.rfactor(loop=l16, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v18 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v18)', + "b19, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l20 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l20, preserve_unit_loops=True)", + "l21 = sch.sample_compute_location(block=b19)", + "sch.compute_at(block=b19, loop=l21, preserve_unit_loops=True)", + "l22 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l22, preserve_unit_loops=True)", + "b23, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l24 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l24, preserve_unit_loops=True)", + "l25 = sch.sample_compute_location(block=b23)", + "sch.compute_at(block=b23, loop=l25, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b2)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l9, factor_axis=1)", + 'sch.annotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "b12, = sch.get_producers(block=b2)", + 'sch.unannotate(block_or_loop=b2, ann_key="meta_schedule.random_compute_producer")', + "l13 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l13, preserve_unit_loops=True)", + "l14 = sch.sample_compute_location(block=b12)", + "sch.compute_at(block=b12, loop=l14, preserve_unit_loops=True)", + "l15 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l15, preserve_unit_loops=True)", + "l16 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l16, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b0)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l8, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "l12 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l12, preserve_unit_loops=True)", + "l13 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l13, preserve_unit_loops=True)", + "b14, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l15 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l15, preserve_unit_loops=True)", + "l16 = sch.sample_compute_location(block=b14)", + "sch.compute_at(block=b14, loop=l16, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + "l4, l5 = sch.get_loops(block=b0)", + "v6, v7 = sch.sample_perfect_tile(loop=l5, n=2, max_innermost_factor=64)", + "l8, l9 = sch.split(loop=l5, factors=[v6, v7])", + "b10 = sch.rfactor(loop=l9, factor_axis=1)", + 'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer", ann_val=1)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v11 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v11)', + "l12 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l12, preserve_unit_loops=True)", + "l13 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l13, preserve_unit_loops=True)", + "b14, = sch.get_producers(block=b0)", + 'sch.unannotate(block_or_loop=b0, ann_key="meta_schedule.random_compute_producer")', + "l15 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l15, preserve_unit_loops=True)", + "l16 = sch.sample_compute_location(block=b14)", + "sch.compute_at(block=b14, loop=l16, preserve_unit_loops=True)", + ], + [ + 'b0 = sch.get_block(name="T_softmax_maxelem", func_name="main")', + 'b1 = sch.get_block(name="T_softmax_exp", func_name="main")', + 'b2 = sch.get_block(name="T_softmax_expsum", func_name="main")', + 'b3 = sch.get_block(name="root", func_name="main")', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.parallel", ann_val=256)', + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.vectorize", ann_val=32)', + "v4 = sch.sample_categorical(candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25])", + 'sch.annotate(block_or_loop=b3, ann_key="meta_schedule.unroll_explicit", ann_val=v4)', + "l5 = sch.sample_compute_location(block=b2)", + "sch.compute_at(block=b2, loop=l5, preserve_unit_loops=True)", + "l6 = sch.sample_compute_location(block=b1)", + "sch.compute_at(block=b1, loop=l6, preserve_unit_loops=True)", + "l7 = sch.sample_compute_location(block=b0)", + "sch.compute_at(block=b0, loop=l7, preserve_unit_loops=True)", + ], + ] + # pylint: enable=line-too-long + ctx = create_context( + create_prim_func(te_workload.softmax_mn(m=256, n=256)), + target=_target(), + ) + spaces = ctx.space_generator.generate_design_space(mod=ctx.mod) + assert len(spaces) == 9 + check_trace(spaces, expected) + + if __name__ == "__main__": test_meta_schedule_cpu_sketch_matmul() test_meta_schedule_cpu_sketch_matmul_relu() test_meta_schedule_cpu_sketch_conv2d_nchw() test_meta_schedule_cpu_sketch_conv2d_nchw_bias_bn_relu() test_meta_schedule_sketch_cpu_max_pool2d_nchw() + test_meta_schedule_cpu_sketch_batchnorm() + test_meta_schedule_cpu_sketch_softmax() From b5fda317f1c06cbf36fd4f4829e6e0f5064edd0f Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Mon, 3 Jan 2022 10:56:26 +0800 Subject: [PATCH 22/26] Minor updates --- python/tvm/meta_schedule/tune.py | 2 +- python/tvm/tir/schedule/schedule.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/meta_schedule/tune.py b/python/tvm/meta_schedule/tune.py index 245de28744e0..da08bf899078 100644 --- a/python/tvm/meta_schedule/tune.py +++ b/python/tvm/meta_schedule/tune.py @@ -106,7 +106,7 @@ def _sch_rules() -> List[ScheduleRule]: ), M.ParallelizeVectorizeUnroll( max_jobs_per_core=16, - max_vectorize_extent=32, + max_vectorize_extent=64, unroll_max_steps=[0, 16, 64, 512], unroll_explicit=True, ), diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index cb207fea6ab7..2e70c3e22802 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -370,7 +370,7 @@ def sample_perfect_tile( ) ) - def sample_compute_location( # Todo: add some unittests + def sample_compute_location( self, block: BlockRV, decision: Optional[int] = None, From 38d7712acf6639f4b7cef4f9fdc3a4aec34b9add Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Mon, 3 Jan 2022 11:23:42 +0800 Subject: [PATCH 23/26] Use bool instead of Bool object --- src/meta_schedule/schedule_rule/random_compute_location.cc | 2 +- src/tir/schedule/utils.h | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc index 8add558b59a8..20f6ac51595f 100644 --- a/src/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/meta_schedule/schedule_rule/random_compute_location.cc @@ -76,7 +76,7 @@ class RandomComputeLocationNode : public ScheduleRuleNode { // - Note that only single producer is allowed in this case. Array producers{nullptr}; if (tir::HasAnn(sch->GetSRef(block_rv), tir::attr::meta_schedule_random_compute_producer, - Bool(true))) { + true)) { producers = sch->GetProducers(block_rv); sch->Unannotate(block_rv, tir::attr::meta_schedule_random_compute_producer); ICHECK_EQ(producers.size(), 1); diff --git a/src/tir/schedule/utils.h b/src/tir/schedule/utils.h index f9787d29b250..d7fdedbbec99 100644 --- a/src/tir/schedule/utils.h +++ b/src/tir/schedule/utils.h @@ -371,9 +371,9 @@ inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const String& an * \param ann_val The boolean annotation value to be checked * \return Whether a Block/For has a specific pair of annotation key and values */ -inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const Bool& ann_val) { +inline bool HasAnn(const StmtSRef& sref, const String& ann_key, bool ann_val) { Optional result = GetAnn(sref, ann_key); - return result.defined() && result.value()->value == ann_val->value; + return result.defined() && result.value()->value == ann_val; } /******** Tensorization ******/ From 9ae015ceec527f10505a717381772f5ac53d1ced Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 7 Jan 2022 10:59:43 +0800 Subject: [PATCH 24/26] Overload another HasAnn --- src/tir/schedule/utils.h | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/tir/schedule/utils.h b/src/tir/schedule/utils.h index d7fdedbbec99..144ba5f77e99 100644 --- a/src/tir/schedule/utils.h +++ b/src/tir/schedule/utils.h @@ -364,6 +364,18 @@ inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const String& an return result.defined() && result.value() == ann_val; } +/*! + * \brief Check if a Block/For has a specific pair of annotation key and values + * \param sref The sref to the block or the for loop + * \param ann_key The annotation key to be checked + * \param ann_val The string annotation value to be checked + * \return Whether a Block/For has a specific pair of annotation key and values + */ +inline bool HasAnn(const StmtSRef& sref, const String& ann_key, const char* ann_val) { + Optional result = GetAnn(sref, ann_key); + return result.defined() && result.value() == ann_val; +} + /*! * \brief Check if a Block/For has a specific pair of annotation key and values * \param sref The sref to the block or the for loop From db10385f1a75bb0417969871da1b57449e5768d2 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 7 Jan 2022 11:02:41 +0800 Subject: [PATCH 25/26] Add prefix --- src/tir/transforms/memhammer_lower_auto_copy.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tir/transforms/memhammer_lower_auto_copy.cc b/src/tir/transforms/memhammer_lower_auto_copy.cc index 79984762f66e..9fad4c4c10ee 100644 --- a/src/tir/transforms/memhammer_lower_auto_copy.cc +++ b/src/tir/transforms/memhammer_lower_auto_copy.cc @@ -30,7 +30,7 @@ #include "tvm/tir/transform.h" /*! * \brief Automatically do memory optimizations for auto copy blocks - * \file lower_auto_copy.cc + * \file memhammer_lower_auto_copy.cc */ namespace tvm { From c03d8943ca904df83a43ee57df3328d8ac0eb9ce Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Mon, 10 Jan 2022 04:19:12 -0800 Subject: [PATCH 26/26] Update shell scripts --- tests/python/meta_schedule/run_ansor_cpu.sh | 1 + tests/python/meta_schedule/run_meta_schedule_cpu.sh | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/meta_schedule/run_ansor_cpu.sh b/tests/python/meta_schedule/run_ansor_cpu.sh index 954c88c255b2..a080ded8fdd9 100644 --- a/tests/python/meta_schedule/run_ansor_cpu.sh +++ b/tests/python/meta_schedule/run_ansor_cpu.sh @@ -33,6 +33,7 @@ run DIL run GMM run GRP run NRM +run SFM run T2D # Subgraph run C2d-BN-RELU diff --git a/tests/python/meta_schedule/run_meta_schedule_cpu.sh b/tests/python/meta_schedule/run_meta_schedule_cpu.sh index 5c01875907f6..87bc17f9e8b6 100644 --- a/tests/python/meta_schedule/run_meta_schedule_cpu.sh +++ b/tests/python/meta_schedule/run_meta_schedule_cpu.sh @@ -31,7 +31,8 @@ run DEP run DIL run GMM run GRP -# run NRM +run NRM +run SFM run T2D # Subgraph run C2d-BN-RELU