anirudh2290 commented on a change in pull request #9283: Fix custom op
multi-gpu scaling
URL: https://github.com/apache/incubator-mxnet/pull/9283#discussion_r159801209
##########
File path: src/operator/custom/custom-inl.h
##########
@@ -63,11 +64,80 @@ class Registry {
return nullptr;
}
- static Registry* Get();
+ template<typename Func>
+ void Push(const Func& func,
+ const OpContext& ctx,
+ bool recording,
+ bool training,
+ const std::vector<NDArray>& arrs) {
+ if (naive_engine_) {
+ func();
+ ctx.async_on_complete();
+ return;
+ }
+ std::unique_lock<std::mutex> lock(mutex_);
+ q_.push(
+ [=]() mutable {
+ bool prev_recording = Imperative::Get()->set_is_recording(recording);
+ bool prev_training = Imperative::Get()->set_is_training(training);
+
+ func();
+
+ Imperative::Get()->set_is_training(prev_training);
+ Imperative::Get()->set_is_recording(prev_recording);
+
+ std::vector<Engine::VarHandle> vars;
+ for (const auto& i : arrs) vars.push_back(i.var());
+ Engine::Get()->PushSync([=](RunContext rctx) {
+ ctx.async_on_complete();
+ }, ctx.run_ctx.ctx, vars, {},
+ FnProperty::kNormal, 0, PROFILER_MESSAGE("CustomOperator"));
+ });
+ cv_.notify_all();
+ }
+
+ ~CustomOperator() {
+ if (naive_engine_) return;
+ {
+ std::unique_lock<std::mutex> lock(mutex_);
+ destructing_ = true;
+ cv_.notify_all();
+ }
+ worker_.join();
+ }
+
+ static CustomOperator* Get();
+
private:
- Registry() {}
+ CustomOperator() {
+ destructing_ = false;
+ naive_engine_ = true;
+ if (std::string("NaiveEngine") != dmlc::GetEnv("MXNET_ENGINE_TYPE",
std::string())) {
+ naive_engine_ = false;
+ worker_ = std::thread(
+ [&]() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ while (!q_.empty() || !destructing_) {
+ cv_.wait(lock, [&] {return !q_.empty() || destructing_;});
+ while (!q_.empty()) {
+ auto fn = q_.front();
+ lock.unlock();
Review comment:
What happens when after we unlock, something is pushed to the queue ? Will
we pop the correct item from the queue ?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
