[MXNET-1417][Performance] Caching Dynamic Shape Checking Result (#15262)

* fix

* address comments

Author: junrushao

src/imperative/cached_op.cc

@@ -100,6 +100,7 @@ CachedOp::CachedOp(
   static const std::vector<const Op*> zero_ops{Op::Get("zeros_like"), Op::Get("_zeros")};
   static const auto _copy_op = Op::Get("_copy");
   config_.Init(flags);
+  this->dynamic_shape_checked_ = false;
 
   if (config_.static_shape) {
     CHECK(config_.static_alloc) << "static_alloc must be True when static_shape is True";
@@ -272,6 +273,11 @@ bool CachedOp::CheckDynamicShapeExists(const Context& default_ctx,
                                        bool erase_result) {
   using namespace nnvm;
   using namespace imperative;
+  if (this->dynamic_shape_checked_) {
+    return config_.is_dynamic;
+  } else {
+    this->dynamic_shape_checked_ = true;
+  }
   CHECK_EQ(inputs.size(), num_inputs());
 
   auto state_ptr = GetCachedOpState(default_ctx);

src/imperative/cached_op.h

@@ -196,6 +196,7 @@ class CachedOp {
   nnvm::Graph grad_graph_;
   nnvm::Graph full_graph_;
   bool inlining_;
+  bool dynamic_shape_checked_;
   std::vector<nnvm::NodeEntry> ograd_entries_;
   std::vector<uint32_t> bwd_in_dep_, bwd_out_dep_, bwd_ograd_dep_;
   std::unordered_map<uint32_t, uint32_t> fwd_input_to_grad_output_;