MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / search_action

Method search_action

src/core/impl/graph/seq_sublinear_memory.cpp:923–973  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

921}
922
923SeqModifierForSublinearMemory::SeqModifyAction SeqModifierForSublinearMemory::
924 search_action(const CompNode::UnorderedMap<OprNodeArray>* cn2oprseq) {
925 m_thread2planner.clear();
926
927 size_t planner_concur;
928 if (auto env = MGB_GETENV("MGB_SUBLINEAR_MEMORY_WORKERS")) {
929 auto set = static_cast<size_t>(std::stoi(env));
930 mgb_assert(
931 set && set <= static_cast<size_t>(sys::get_cpu_count()) * 4,
932 "invalid planner concurrency: %zu", set);
933 planner_concur = set;
934 } else {
935 planner_concur = m_config->num_worker;
936 }
937
938 std::string msg = ssprintf(
939 "use %zu threads to search for sublinear memory plan; this can be changed "
940 "via %c%cB_SUBLINEAR_MEMORY_WORKERS env var",
941 planner_concur, 'M', 'G');
942 mgb_log_debug("%s", msg.c_str());
943 for (auto&& i : m_planner_thread_pool.start(planner_concur))
944 m_thread2planner[i].reset(new ModifyActionPlanner{this});
945
946 std::vector<std::unique_ptr<ActionSearcherSingleCN>> searchers;
947 searchers.reserve(cn2oprseq->size());
948
949 using WorkerPool = FutureThreadPool<const SeqModifyAction&>;
950 WorkerPool workers;
951 workers.start(cn2oprseq->size());
952
953 m_prev_min_bottleneck.clear();
954 for (auto&& i : *cn2oprseq) {
955 m_prev_min_bottleneck[i.first] = 0;
956 }
957
958 std::vector<WorkerPool::Future> futures;
959 for (auto&& i : *cn2oprseq) {
960 searchers.emplace_back(std::make_unique<ActionSearcherSingleCN>(this));
961 futures.emplace_back(workers.launch(
962 &ActionSearcherSingleCN::search, searchers.back().get(), i.first,
963 &i.second));
964 }
965
966 SeqModifyAction action;
967 for (auto&& i : futures) {
968 auto&& cur = i.get();
969 action.insert(cur.begin(), cur.end());
970 }
971 m_thread2planner.clear();
972 return action;
973}
974
975void SeqModifierForSublinearMemory::apply_action(
976 SeqModifyAction& action, const OprNodeArray& oprseq) {

Callers

nothing calls this directly

Calls 12

emplace_backMethod · 0.80
backMethod · 0.80
clearMethod · 0.45
startMethod · 0.45
resetMethod · 0.45
reserveMethod · 0.45
sizeMethod · 0.45
launchMethod · 0.45
getMethod · 0.45
insertMethod · 0.45
beginMethod · 0.45
endMethod · 0.45

Tested by

no test coverage detected