engine.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef MXNET_ENGINE_H_
#define MXNET_ENGINE_H_

#include <dmlc/base.h>
#if DMLC_USE_CXX11
#include <algorithm>
#include <memory>
#include <functional>
#endif
#include <vector>
#include "./base.h"

namespace mxnet {

// forward declare engine
class Engine;

namespace engine {
struct Var;
struct Opr;
typedef Var* VarHandle;
typedef Opr* OprHandle;
class CallbackOnComplete {
 public:
  // use implicit copy and assign
 inline void operator()() const {
    (*callback_)(engine_, param_);
  }

 private:
  friend class ::mxnet::Engine;
  void (*callback_)(Engine *, void *);
  Engine* engine_;
  void* param_;
};
}  // namespace engine

#if DMLC_USE_CXX11

enum class FnProperty {
  kNormal,
  kCopyFromGPU,
  kCopyToGPU,
  kCPUPrioritized,
  kAsync,
  kDeleteVar
};  // enum class FnProperty

class MXNET_API Engine {
 public:
  typedef engine::CallbackOnComplete CallbackOnComplete;
  typedef std::function<void(RunContext)> SyncFn;
  typedef std::function<void(RunContext, CallbackOnComplete)> AsyncFn;
  typedef engine::VarHandle VarHandle;
  typedef engine::OprHandle OprHandle;
  virtual void NotifyShutdown() = 0;
  virtual VarHandle NewVariable() = 0;
  virtual OprHandle NewOperator(AsyncFn fn,
                                std::vector<VarHandle> const& const_vars,
                                std::vector<VarHandle> const& mutable_vars,
                                FnProperty prop = FnProperty::kNormal,
                                const char* opr_name = nullptr) = 0;
  virtual void DeleteOperator(OprHandle op) = 0;
  virtual void Push(OprHandle op, Context exec_ctx, int priority = 0, bool profiling = false) = 0;
  virtual void PushAsync(AsyncFn exec_fun, Context exec_ctx,
                         std::vector<VarHandle> const& const_vars,
                         std::vector<VarHandle> const& mutable_vars,
                         FnProperty prop = FnProperty::kNormal,
                         int priority = 0,
                         const char* opr_name = nullptr) = 0;
  virtual void DeleteVariable(SyncFn delete_fn,
                              Context exec_ctx,
                              VarHandle var) = 0;
  virtual void WaitForVar(VarHandle var) = 0;
  virtual void WaitForAll() = 0;
  virtual ~Engine() noexcept(false) {}
  static Engine* Get();
  static std::shared_ptr<Engine> _GetSharedRef();
  inline void PushSync(SyncFn exec_fn, Context exec_ctx,
                       std::vector<VarHandle> const& const_vars,
                       std::vector<VarHandle> const& mutable_vars,
                       FnProperty prop = FnProperty::kNormal,
                       int priority = 0,
                       const char* opr_name = nullptr) {
    this->PushAsync([exec_fn](RunContext ctx, CallbackOnComplete on_complete) {
        exec_fn(ctx);
        on_complete();
      }, exec_ctx, const_vars, mutable_vars, prop, priority, opr_name);
  }

  inline CallbackOnComplete CreateCallback(
      void (*callback)(Engine *, void *), void *param) {
    CallbackOnComplete ret;
    ret.callback_ = callback;
    ret.engine_ = this;
    ret.param_ = param;
    return ret;
  }
  // For each var vector, sort it and remove the duplicated vars.
  // Also remove vars from read_vars if it also appears in write_vars
  inline void DeduplicateVarHandle(std::vector<engine::VarHandle> *read_vars,
                                   std::vector<engine::VarHandle> *write_vars) {
    std::sort(write_vars->begin(), write_vars->end());
    write_vars->resize(std::unique(write_vars->begin(), write_vars->end()) -
                      write_vars->begin());
    std::sort(read_vars->begin(), read_vars->end());
    read_vars->resize(std::unique(read_vars->begin(), read_vars->end()) -
                      read_vars->begin());
    auto wit = write_vars->begin();
    auto rtop = read_vars->begin();
    for (auto rit = read_vars->begin(); rit != read_vars->end(); ++rit) {
      while (wit != write_vars->end() && *wit < *rit) ++wit;
      if (wit == write_vars->end() || *wit != *rit) {
        *rtop = *rit;
        ++rtop;
      }
    }
    read_vars->resize(rtop - read_vars->begin());
  }

  virtual int num_omp_threads_per_worker() const = 0;

  virtual void set_num_omp_threads_per_worker(int num_omp_threads_per_worker) = 0;
};  // class Engine
#endif  // DMLC_USE_CXX11
}  // namespace mxnet
#endif  // MXNET_ENGINE_H_