Improve cached_op performance for static mode

src/executor/attach_op_execs_pass.cc

@@ -261,7 +261,7 @@ class FComputeExExecutor : public OpExecutor {
   ExecType exec_type_;
 };
 
-void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i) {
+void CreateOpExecs(const Graph& g, OpExecVector* p_ret, OpStateVector* p_state, size_t i) {
   using nnvm::DTypeVector;
   using mxnet::ShapeVector;
   using nnvm::FMutateInputs;
@@ -302,6 +302,7 @@ void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i) {
 
     OpStatePtr state = fcreate_op_state[op](
         inode.source->attrs, vctx[i], ishape, itype);
+    if (p_state) p_state->at(i) = state;
     FStatefulComputeEx fcompute_ex = common::GetFCompute<FStatefulComputeEx>(
         op, "FStatefulComputeEx", vctx[i]);
     // FStatefulComputeEx is dispatched only when dispatch_mode is DispatchMode::kFComputeEx
@@ -359,7 +360,7 @@ Graph AttachOpExecs(Graph g) {
   const auto& idx = g.indexed_graph();
   OpExecVector ret(idx.num_nodes());
   for (size_t i = 0; i < idx.num_nodes(); ++i) {
-    CreateOpExecs(g, &ret, i);
+    CreateOpExecs(g, &ret, nullptr, i);
   }
   g.attrs["op_execs"] = std::make_shared<nnvm::any>(ret);
   return g;

src/executor/exec_pass.h

@@ -98,6 +98,12 @@ class OpExecutor {
  */
 using OpExecVector = std::vector<std::shared_ptr<OpExecutor> >;
 
+/*!
+ * \brief per node vector of operator states.
+ * \note stored under attribute "op_states"
+ */
+using OpStateVector = std::vector<OpStatePtr>;
+
 /*!
  * \brief per node context vector
  * \node stored under "context"
@@ -115,9 +121,10 @@ using DevMaskVector = std::vector<int>;
  *
  * \param g input graph
  * \param p_ret OpExecVector for input and output
+ * \param p_state OpStateVector if it has.
  * \param i the id of the node
  */
-void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i);
+void CreateOpExecs(const Graph& g, OpExecVector* p_ret, OpStateVector* p_state, size_t i);
 /*!
  * \brief Attach OpExecutor to the graph attributes.
  *

src/imperative/cached_op.cc

@@ -603,7 +603,7 @@ void CachedOp::StaticInitExec(
     }
   } else {
     for (size_t i = start_nid; i < end_nid; ++i) {
-      exec::CreateOpExecs(g, &state.execs, i);
+      exec::CreateOpExecs(g, &state.execs, &state.op_states, i);
     }
     exec::AttachOpResources(g, state.execs, start_nid, end_nid);
 
@@ -705,8 +705,6 @@ void CachedOp::StaticRunOps(
           arg_shapes.emplace_back(ndinput->shape());
           arg_dtypes.emplace_back(ndinput->dtype());
         }
-        state.op_states[i] = createop[node.source->op()](
-            node.source->attrs, default_ctx, arg_shapes, arg_dtypes);
         Imperative::Get()->InvokeOp(
             default_ctx, node.source->attrs, ndinputs, ndoutputs, req,
             dispatch_mode, state.op_states[i]);
@@ -910,7 +908,7 @@ OpStatePtr CachedOp::Forward(
 
   OpStatePtr op_state;
   try {
-    if (config_.is_dynamic || CheckDynamicShapeExists(default_ctx, inputs, true)) {
+    if (config_.is_dynamic && CheckDynamicShapeExists(default_ctx, inputs, true)) {
       config_.is_dynamic = true;
       config_.static_alloc = false;
       op_state = DynamicForward(default_ctx, inputs, outputs, true);