[MXNET-1315] Add checks for dynamic-shaped operators in CachedOp (apa…

…che#14018)

* Initial commit

* Try this

* Boy next door!

* Add comments per discussion with Da

* Try this

* boy try this

* change the boss of the gym

src/imperative/cached_op.cc

@@ -285,12 +285,20 @@ bool CachedOp::SetForwardGraph(
   }
 
   bool match = true;
-  match &= CheckAndInferShape(&g, std::move(shape_inputs), true);
+  bool contain_dynamic_shape = false;
+  match &= CheckAndInferShape(&g, std::move(shape_inputs), true,
+                              {0, 0}, {0, 0}, &contain_dynamic_shape);
   match &= CheckAndInferType(&g, std::move(dtype_inputs), true);
   exec::DevMaskVector dev_mask(g.indexed_graph().num_nodes(), inputs[0]->ctx().dev_mask());
   match &= CheckAndInferStorageType(&g, std::move(dev_mask),
                                     std::move(storage_type_inputs), true);
 
+  // When dynmaic shape exists, it is not feasible to plan memory ahead of time
+  if (contain_dynamic_shape) {
+    g.attrs.erase("forward_mem_plan");
+    g.attrs.erase("full_mem_plan");
+    return false;
+  }
   if (!match) {
     g.attrs.erase("forward_mem_plan");
     g.attrs.erase("full_mem_plan");

src/imperative/cached_op.h

@@ -35,6 +35,7 @@ struct CachedOpConfig : public dmlc::Parameter<CachedOpConfig> {
   uint32_t backward_bulk_size;
   bool static_alloc;
   bool static_shape;
+  bool is_dynamic;
   nnvm::Tuple<uint32_t> data_indices;
   nnvm::Tuple<uint32_t> param_indices;
   std::string subgraph;
@@ -66,6 +67,9 @@ struct CachedOpConfig : public dmlc::Parameter<CachedOpConfig> {
     DMLC_DECLARE_FIELD(subgraph)
     .set_default(std::string(""))
     .describe("JSON string of a subgraph.");
+    DMLC_DECLARE_FIELD(is_dynamic)
+    .set_default(false)
+    .describe("Whether the graph contains dynamic shape operators.");
   }
 };
 

src/imperative/imperative_utils.h

@@ -179,7 +179,7 @@ inline void SetShapeType(const Context& ctx,
 
   for (size_t i = 0; i < outputs.size(); ++i) {
     NDArrayStorageType storage_type = static_cast<NDArrayStorageType>(out_storage_types[i]);
-    if (outputs[i]->is_none()) {
+    if (outputs[i]->is_none() || outputs[i]->shape().ndim() == 0) {
       if (is_dynamic_shape_existing) {
         // once there is dynamic shape somewhere, we could not pre-determine the shape.
         *outputs[i] = NDArray(ctx, out_types[i]);
@@ -566,8 +566,12 @@ inline void PushOperator(const OpStatePtr& state,
 inline bool CheckAndInferShape(nnvm::Graph* p_g, nnvm::ShapeVector&& shapes,
                                bool use_inputs,
                                std::pair<uint32_t, uint32_t> node_range = {0, 0},
-                               std::pair<uint32_t, uint32_t> entry_range = {0, 0}) {
+                               std::pair<uint32_t, uint32_t> entry_range = {0, 0},
+                               bool *contain_unknown = nullptr) {
   using namespace nnvm;
+  if (contain_unknown != nullptr) {
+    *contain_unknown = false;
+  }
   nnvm::Graph& g = *p_g;
   if (use_inputs) {
     if (g.attrs.count("shape_inputs") &&
@@ -601,8 +605,11 @@ inline bool CheckAndInferShape(nnvm::Graph* p_g, nnvm::ShapeVector&& shapes,
     g.attrs["shape"] = std::make_shared<dmlc::any>(std::move(shapes));
     g = exec::InferShape(std::move(g));
   }
-  CHECK_EQ(g.GetAttr<size_t>("shape_num_unknown_nodes"), 0U);
-
+  if (contain_unknown == nullptr) {
+    CHECK_EQ(g.GetAttr<size_t>("shape_num_unknown_nodes"), 0U);
+  } else {
+    *contain_unknown = g.GetAttr<size_t>("shape_num_unknown_nodes") != 0U;
+  }
   return false;
 }
 

src/operator/contrib/boolean_mask.cc

@@ -98,6 +98,8 @@ which stands for the rows in x where the corresonding element in index is non-ze
 .set_attr<FComputeEx>("FComputeEx<cpu>", BooleanMaskForward<cpu>)
 .set_attr<FInferStorageType>("FInferStorageType", BooleanMaskStorageType)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_contrib_boolean_mask"})
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"data", "index"};})
 .add_argument("data", "NDArray-or-Symbol", "Data")
 .add_argument("index", "NDArray-or-Symbol", "Mask")
 .add_arguments(BooleanMaskParam::__FIELDS__());

src/operator/subgraph_op_common.cc

@@ -222,8 +222,14 @@ void LoopState::Forward(int iter_no,
   // If an input and an output share the array, the output array will be changed
   // by CachedOp. We need to copy data to the real output.
   for (size_t i = 0; i < out_bufs.size(); i++)
-    if (!out_bufs[i].IsSame(coutputs[i]))
+    if (!out_bufs[i].IsSame(coutputs[i])) {
+      // The line below checks whether dynamic shape exists.
+      // If so, re-initialize the shape.
+      if (coutputs[i].shape().ndim() == 0) {
+        const_cast<NDArray &>(coutputs[i]).Init(out_bufs[i].shape());
+      }
       CopyFromTo(out_bufs[i], coutputs[i]);
+    }
   if (is_recording) {
     all_inputs.push_back(cinputs);
     all_outputs.push_back(coutputs);

src/operator/subgraph_op_common.h

@@ -161,7 +161,8 @@ class LoopState {
     // only static_alloc supports nested call of CachedOp.
     std::vector<std::pair<std::string, std::string> > kwargs = {
       {"inline_limit", "0"},
-      {"static_alloc", "1"}
+      {"static_alloc", "1"},
+      {"is_dynamic", "1"}
     };
     return std::make_shared<CachedOp>(sym, kwargs);
   }