Refactor AGInfo and Imperative

CMakeLists.txt

@@ -143,7 +143,7 @@ else(MSVC)
     add_definitions(-DMSHADOW_USE_F16C=0)
   endif()
   set(CMAKE_POSITION_INDEPENDENT_CODE ON)
-  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wno-unknown-pragmas -Wno-sign-compare")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wno-unknown-pragmas -Wno-sign-compare -Werror=return-type")
   if ("${CMAKE_CXX_COMPILER_ID}" MATCHES ".*Clang$")
     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wno-braced-scalar-init")
   endif()

include/mxnet/c_api.h

@@ -1224,7 +1224,7 @@ MXNET_DLL int MXAutogradBackward(mx_uint num_output,
  * \param output_handles output NDArrays
  * \param ograd_handles head gradient for NDArrays
  * \param num_variables number of variables
- * \param
+ * \param var_handles variables to compute gradient with respect to (d / d var)
  * \param retain_graph whether to keep the graph after backward
  * \param is_train whether to do backward for training or inference
  * \return 0 when success, -1 when failure happens

include/mxnet/imperative.h

@@ -35,48 +35,53 @@
 #include "./ndarray.h"
 
 namespace mxnet {
-/*! \brief runtime functions for NDArray */
-class Imperative {
+/*!
+ * Autograd Info used in class: nnvm::Node::info
+ */
+class AGInfo {
  public:
-  /*! \brief */
-  class AGInfo {
-   public:
-    Context ctx;
-    OpReqType grad_req;
-    OpStatePtr state;
-    std::vector<NDArray> outputs;
-    std::vector<NDArray> out_grads;
-    bool fresh_out_grad;
+  Context ctx;
+  OpReqType grad_req;
+  OpStatePtr state;
+  std::vector<NDArray> outputs;
+  std::vector<NDArray> out_grads;
+  bool fresh_out_grad;
 
-    AGInfo() :
+  AGInfo() :
       grad_req(kNullOp), fresh_out_grad(false) {}
 
-    static void Clear(const nnvm::NodePtr& node) {
-      if (node == nullptr || node->info.empty()) return;
-      AGInfo& info = Get(node);
-      if (info.grad_req != kNullOp) return;
-      node->info.clear();
-    }
+  static void Clear(const nnvm::NodePtr& node) {
+    if (node == nullptr || node->info.empty()) return;
+    AGInfo& info = Get(node);
+    if (info.grad_req != kNullOp) return;
+    node->info.clear();
+  }
 
-    static AGInfo& Get(const nnvm::NodePtr& node) {
-      return dmlc::get<AGInfo>(node->info);
-    }
+  static AGInfo& Get(const nnvm::NodePtr& node) {
+    return dmlc::get<AGInfo>(node->info);
+  }
 
-    static AGInfo& Create(const nnvm::NodePtr& node) {
-      node->info.construct<AGInfo>();
-      return Get(node);
-    }
+  static AGInfo& Create(const nnvm::NodePtr& node) {
+    node->info.construct<AGInfo>();
+    return Get(node);
+  }
 
-    static bool IsNone(const NDArray& arr) {
-      return arr.entry_.node == nullptr || arr.entry_.node->info.empty();
-    }
+  static bool IsNone(const NDArray& arr) {
+    return arr.autograd_.node == nullptr || arr.autograd_.node->info.empty();
+  }
+
+  static bool IsVariable(const nnvm::NodePtr& node) {
+    AGInfo& info = Get(node);
+    return info.grad_req != kNullOp && info.outputs.size() == 1
+        && info.out_grads.size() == 1;
+  }
+};
+
+/*! \brief runtime functions for NDArray */
+class Imperative {
+ public:
+  /*! \brief */
 
-    static bool IsVariable(const nnvm::NodePtr& node) {
-      AGInfo& info = Get(node);
-      return info.grad_req != kNullOp && info.outputs.size() == 1
-             && info.out_grads.size() == 1;
-    }
-  };
   /*! \brief whether operator recording is on. */
   bool is_training() const {
     return is_train_;
@@ -97,11 +102,11 @@ class Imperative {
       is_recording_ = is_recording;
       return old;
   }
-  /*! brief whether numpy compatibility is on. */
+  /*! \brief whether numpy compatibility is on. */
   bool is_np_shape() const {
     return is_np_shape_;
   }
-  /*! brief turn on or turn off numpy compatibility switch. */
+  /*! \brief turn on or turn off numpy compatibility switch. */
   bool set_is_np_shape(bool is_np_shape) {
     bool old = is_np_shape_;
     is_np_shape_ = is_np_shape;
@@ -160,17 +165,39 @@ class Imperative {
 
  private:
   friend class NDArray;
-  /*! \brief make constructor protected. */
+  /*! Create a forward graph
+   * @param output_nodes graph node vector to add nodes to
+   * @param outputs source ndarrays
+   * @return vector of nodes
+   */
+  static nnvm::Graph CreateGraph(const std::vector<NDArray *> &outputs);
+  /*! Create gradient nodes using output shapes and ctx.
+   * Gradient heads are initialized to 1 if they are not present (nullptr)
+   * @return vector of nodes
+   */
+  static std::vector<nnvm::NodeEntry> CreateHeadGradientNodes(const std::vector<NDArray*>& outputs,
+                                                              const std::vector<NDArray*>& ograds);
+
+  struct GradientVariableNodes;
+  /*! Create variable nodes.
+   * If variables is provided, gradient nodes are crated for them. Otherwise it uses read only
+   * inputs reachable from the outputs.
+   * @param variables
+   * @param outputs
+   * @return aux data structure with nodes and arrays for gradients
+   */
+  GradientVariableNodes CreateGradientVariableNodes(const std::vector<NDArray*>& variables,
+                                            const std::vector<nnvm::NodeEntry>& outputs);
   Imperative() {
     if (PreferBulkExecTrain())
       backward_bulk_size_ = BulkExecMaxNodeTrainBwd();
   }
   /*! \brief find the input/output ndarrays that are needed for backward */
   void GetBackwardDependency(
       const nnvm::NodePtr& node,
-      uint32_t num_inputs, uint32_t num_outputs,
-      std::vector<bool> *p_save_inputs,
-      std::vector<bool> *p_save_outputs);
+      size_t num_inputs, size_t num_outputs,
+      std::vector<bool> *save_inputs,
+      std::vector<bool> *save_outputs);
   /*! \brief indicate whether is training. */
 #if DMLC_CXX11_THREAD_LOCAL
   static thread_local bool is_train_;

include/mxnet/ndarray.h

@@ -80,10 +80,12 @@ class MKLDNNMemory;
  * \brief ndarray interface
  */
 class NDArray {
+  friend class AGInfo;
+  friend class Imperative;
  public:
   /*! \brief default constructor */
   NDArray()
-    : entry_(nullptr) {
+    : autograd_(nullptr) {
   }
   /*!
    * \brief constructs a new dynamic NDArray
@@ -98,7 +100,7 @@ class NDArray {
         shape_(shape),
         dtype_(dtype),
         storage_type_(kDefaultStorage),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
   /*! \brief constructor for NDArray with storage type
    */
@@ -117,7 +119,7 @@ class NDArray {
         shape_(),
         dtype_(dtype),
         storage_type_(kDefaultStorage),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
   /*!
    * \brief constructing a static NDArray that shares data with TBlob
@@ -131,7 +133,7 @@ class NDArray {
         shape_(data.shape_),
         dtype_(data.type_flag_),
         storage_type_(kDefaultStorage),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
 
   /*!
@@ -149,7 +151,7 @@ class NDArray {
         }),
         shape_(data.shape_),
         dtype_(data.type_flag_), storage_type_(kDefaultStorage),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
 
   /*! \brief create ndarray from shared memory */
@@ -158,7 +160,7 @@ class NDArray {
         shape_(shape),
         dtype_(dtype),
         storage_type_(kDefaultStorage),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
 
   /*!
@@ -177,7 +179,7 @@ class NDArray {
         shape_(shape),
         dtype_(data.type_flag_),
         storage_type_(stype),
-        entry_(nullptr) {
+        autograd_(nullptr) {
   }
   /*!
    * \brief initialize the NDArray, assuming it is not assigned a meaningful shape before
@@ -387,6 +389,7 @@ class NDArray {
   }
   /*! \return the associated variable of the ndarray.*/
   inline Engine::VarHandle var() const {
+    CHECK(ptr_);
     return ptr_->var;
   }
   /*! \return byte offset in chunk of the ndarray*/
@@ -395,6 +398,7 @@ class NDArray {
   }
   /*! \brief return var version of the NDArray*/
   inline size_t version() const {
+    CHECK(var());
     return var()->version();
   }
   /*!
@@ -649,7 +653,7 @@ class NDArray {
    */
   NDArray Detach() const {
     NDArray ret(*this);
-    ret.entry_ = nnvm::NodeEntry(nullptr);
+    ret.autograd_ = nnvm::NodeEntry(nullptr);
     return ret;
   }
 
@@ -812,7 +816,6 @@ class NDArray {
                    std::vector<std::string>* keys);
 
  private:
-  friend class Imperative;
   /*! \brief the real data chunk that backs NDArray */
   // shandle is used to store the actual values in the NDArray
   // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
@@ -1102,7 +1105,7 @@ class NDArray {
   /*! \brief storage type of data */
   NDArrayStorageType storage_type_ = kUndefinedStorage;
   /*! \brief node entry for autograd */
-  nnvm::NodeEntry entry_;
+  nnvm::NodeEntry autograd_;
   /*!
    * \brief internal TBlob
    * \note When user access tblob_ by some const methods like

include/mxnet/op_attr_types.h

@@ -54,7 +54,8 @@ enum OpReqType {
    */
   kWriteInplace,
   /*! \brief add to the provided space */
-  kAddTo
+  kAddTo,
+  kOpReqTypeMax
 };
 
 /*!

src/c_api/c_api_ndarray.cc

@@ -355,7 +355,7 @@ int MXAutogradBackwardEx(mx_uint num_output,
   }
 
   auto grads = Imperative::Get()->Backward(outputs, ograds, variables, is_train,
-                                                  retain_graph, create_graph);
+                                           retain_graph, create_graph);
   if (num_variables != 0) {
     ret->ret_handles.clear();
     ret->out_types.clear();

src/engine/stream_manager.h

@@ -102,6 +102,7 @@ RunContext StreamManager<kNumGpus, kStreams>::GetRunContext(
 #endif  // MXNET_USE_CUDA
     default:
       LOG(FATAL) << "Not Reached";
+      break;
     }
   }
   return ret;

src/executor/exec_pass.h

@@ -258,6 +258,7 @@ inline Graph MXGradient(
   if (copy_op_str != std::string()) {
       graph.attrs["copy_op"] = std::make_shared<any>(std::move(copy_op_str));
   }
+  /// @sa nnvm::pass::Gradient in gradient.cc
   return ApplyPass(std::move(graph), "MXGradient");
 }
 }  // namespace pass