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

Method wait_tensor

imperative/src/impl/interpreter/interpreter_impl.cpp:1049–1137  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1047}
1048
1049TensorPtr ChannelImpl::wait_tensor(TensorInfo* info, TensorProp prop) {
1050 std::unique_lock<decltype(m_mutex)> lock(m_mutex);
1051 mgb_assert(!m_waitee, "duplicate waitee");
1052 m_waitee = info;
1053 m_waitee_id = Profiler::next_id();
1054 auto backtrace_getter = [bt = get_backtrace()]() {
1055 std::unordered_map<std::string, std::string> infos;
1056 if (bt != nullptr) {
1057 if (bt->py_stack_info != nullptr)
1058 infos["python_backtrace"] = bt->py_traceback();
1059 if (bt->trans_stack_info.size() > 0)
1060 infos["transformation_backtrace"] = bt->transformation_traceback();
1061 }
1062 return infos;
1063 };
1064 MGB_RECORD_EVENT(
1065 TensorWaitPropEvent, info->id, m_waitee_id, prop, backtrace_getter);
1066 bool require_host = prop == TensorProp::HostValue;
1067 bool require_dev = prop == TensorProp::DevValue;
1068 auto host_available = [&] { return info->ptr && info->ptr->value_fetched(); };
1069 auto dev_available = [&] { return info->ptr; };
1070 bool wait_host = false;
1071 bool wait_regen = false;
1072 if (require_host && !host_available()) {
1073 // avoid dead lock
1074 lock.unlock();
1075 if (Profiler::is_profiling()) {
1076 m_worker.add_task(
1077 {Profiler::next_id(), GetValue{info},
1078 get_channel_state().stack_manager.dump()});
1079 } else {
1080 m_worker.add_task({
1081 Profiler::next_id(),
1082 GetValue{info},
1083 });
1084 }
1085 lock.lock();
1086 wait_host = true;
1087 }
1088 if (require_dev && !dev_available()) {
1089 lock.unlock();
1090 if (Profiler::is_profiling()) {
1091 m_worker.add_task(
1092 {Profiler::next_id(), StartRegen{info},
1093 get_channel_state().stack_manager.dump()});
1094 } else {
1095 m_worker.add_task({
1096 Profiler::next_id(),
1097 StartRegen{info},
1098 });
1099 }
1100 lock.lock();
1101 wait_regen = true;
1102 }
1103 if (require_dev) {
1104 m_cv.wait(lock, [&]() {
1105 check_worker_exc_unsafe();
1106 return dev_available();

Callers

nothing calls this directly

Calls 13

is_profilingFunction · 0.85
py_tracebackMethod · 0.80
value_fetchedMethod · 0.80
check_async_errorMethod · 0.80
sizeMethod · 0.45
unlockMethod · 0.45
add_taskMethod · 0.45
dumpMethod · 0.45
lockMethod · 0.45
waitMethod · 0.45
comp_nodeMethod · 0.45

Tested by

no test coverage detected