Task infer blob desc support choosing method

oneflow/core/graph/boxing_identity_task_node.cpp

@@ -47,7 +47,7 @@ void BoxingIdentityTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
-  node->InferBlobDescs(nullptr);
+  (node->*GetInferBlobDescsMethod())(nullptr);
 }
 
 void BoxingIdentityTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/boxing_zeros_task_node.cpp

@@ -51,7 +51,7 @@ void BoxingZerosTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
-  node->InferBlobDescs(nullptr);
+  (node->*GetInferBlobDescsMethod())(nullptr);
 }
 
 void BoxingZerosTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/collective_boxing_pack_task_node.cpp

@@ -60,7 +60,7 @@ void CollectiveBoxingPackTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
-  node->InferBlobDescs(nullptr);
+  (node->*GetInferBlobDescsMethod())(nullptr);
 }
 
 void CollectiveBoxingPackTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/collective_boxing_task_node.cpp

@@ -55,7 +55,7 @@ void CollectiveBoxingGenericTaskNode::BuildExecGphAndRegst() {
     node->BindBnWithRegst(obn, out_regst);
     out_regst->AddLbi(boxing_op->BnInOp2Lbi(obn));
   }
-  node->InferBlobDescs(nullptr);
+  (node->*GetInferBlobDescsMethod())(nullptr);
 }
 
 void CollectiveBoxingGenericTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/collective_boxing_unpack_task_node.cpp

@@ -60,7 +60,7 @@ void CollectiveBoxingUnpackTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
-  node->InferBlobDescs(nullptr);
+  (node->*GetInferBlobDescsMethod())(nullptr);
 }
 
 void CollectiveBoxingUnpackTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/compute_task_node.h

@@ -43,6 +43,10 @@ class CompTaskNode : public TaskNode {
   // op
   std::shared_ptr<const Operator> op() const { return op_node_->shared_op(); }
 
+  ExecNode::InferBlobDescsMethod GetInferBlobDescsMethod() const override {
+    return &ExecNode::InferBlobDescsByInputs;
+  }
+
  protected:
   const OpNode* GetOneSuccOpNodeOnEdge(TaskEdge* edge);
   const OpNode* GetOnePredOpNodeOnEdge(TaskEdge* edge);

oneflow/core/graph/exec_graph.cpp

@@ -103,9 +103,31 @@ Maybe<void> CheckPhysicalBlobDesc(
   return Maybe<void>::Ok();
 }
 
+// A helper function to infer blob's physical shape with ND SBP.
+Maybe<void> InferPhysicalBlobDesc(
+    const Operator& op, const PbRpf<std::string>& bns,
+    const std::function<Maybe<const BlobDesc>(const std::string&)>& GetLogicalBlobDesc,
+    const NdSbpSignature* nd_sbp_signature, const ParallelContext* parallel_ctx,
+    const std::function<BlobDesc*(const std::string&)>& GetPhysicalBlobDesc) {
+  const std::shared_ptr<const ParallelDesc> op_parallel_desc = JUST(op.GetOpParallelDesc());
+  for (const auto& bn : bns) {
+    BlobDesc* physical_blob_desc = GetPhysicalBlobDesc(bn);
+    const auto& logical_blob_desc = *JUST(GetLogicalBlobDesc(bn));
+    CHECK_NOTNULL_OR_RETURN(physical_blob_desc)
+        << "physical_blob_desc should not be nullptr. op location: " << op.op_loc();
+    *physical_blob_desc = logical_blob_desc;
+    const auto& physical_shape = JUST_MSG(
+        GetPhysicalShape(logical_blob_desc.shape(), nd_sbp_signature->bn_in_op2nd_sbp().at(bn),
+                         *op_parallel_desc, *parallel_ctx),
+        std::stringstream() << " check physical shape failed, op name " << op.op_loc());
+    physical_blob_desc->set_shape(*physical_shape);
+  }
+  return Maybe<void>::Ok();
+}
+
 }  // namespace
 
-void ExecNode::InferBlobDescs(const ParallelContext* parallel_ctx) {
+void ExecNode::InferBlobDescsByInputs(const ParallelContext* parallel_ctx) {
   auto GetBlobDesc4BnInOp = GetBlobDesc4BnInOpFunc();
   const OpNode* op_node = Singleton<OpGraph>::Get()->OpNode4OpName(op()->op_name());
   const NdSbpSignature* nd_sbp_signature = nullptr;
@@ -128,7 +150,50 @@ void ExecNode::InferBlobDescs(const ParallelContext* parallel_ctx) {
   CHECK_JUST_MSG(op_->InferInplaceObn2IbnIf(&mut_inplace_obn2ibn_, &con_inplace_obn2ibn_,
                                             GetBlobDesc4BnInOp, parallel_ctx),
                  std::stringstream()
-                     << " infer inplace obn to ibn if failed, op name " << op_->op_loc());
+                     << " infer inplace obn to ibn is failed, op name " << op_->op_loc());
+}
+
+void ExecNode::InferBlobDescsByNdSbp(const ParallelContext* parallel_ctx) {
+  const HashSet<std::string> ibns{op()->input_bns().begin(), op()->input_bns().end()};
+  HashMap<std::string, BlobDesc> ibn2blob_desc{};
+  const auto& GetBlobDesc4BnInOp = [&](const std::string& bn_in_op) -> BlobDesc* {
+    // Generate temp regst to store input blob desc, and will be released after infer output blob
+    // desc.
+    if (ibns.count(bn_in_op) > 0) {
+      auto iter = ibn2blob_desc.find(bn_in_op);
+      if (iter == ibn2blob_desc.end()) {
+        iter = ibn2blob_desc.emplace(bn_in_op, kInvalidDataType).first;
+      }
+      return &iter->second;
+    }
+    auto it = bn_in_op2regst_.find(bn_in_op);
+    if (it == bn_in_op2regst_.end()) { return nullptr; }
+    std::shared_ptr<RegstDesc> regst = it->second;
+    CHECK(regst);
+    return regst->MutBlobDesc(op()->BnInOp2Lbi(bn_in_op));
+  };
+  const OpNode* op_node = Singleton<OpGraph>::Get()->OpNode4OpName(op()->op_name());
+  const NdSbpSignature* nd_sbp_signature = &CHECK_NOTNULL(op_node)->nd_sbp_signature();
+
+  // TODO(strint): user op can infer output with SBP, so there is no need to infer the input.
+  // Reference: https://github.com/Oneflow-Inc/oneflow/pull/8971
+  // Infer input blob desc with SBP, the infer results are set into the temp input blob desc.
+  CHECK_JUST(InferPhysicalBlobDesc(
+      *op(), op()->input_bns(),
+      std::bind(&Operator::GetLogicalBlobDesc4Ibn, op().get(), std::placeholders::_1),
+      nd_sbp_signature, parallel_ctx, GetBlobDesc4BnInOp));
+
+  // Infer output blob desc with input.
+  CHECK_JUST_MSG(op_->InferBlobDescsIf(GetBlobDesc4BnInOp, parallel_ctx, &GlobalJobDesc()),
+                 std::stringstream() << " infer blob descs is failed, op name " << op_->op_loc());
+  CHECK_JUST(CheckPhysicalBlobDesc(
+      *op(), op()->output_bns(),
+      std::bind(&Operator::GetLogicalBlobDesc4Obn, op().get(), std::placeholders::_1),
+      nd_sbp_signature, parallel_ctx, GetBlobDesc4BnInOp));
+  CHECK_JUST_MSG(op_->InferInplaceObn2IbnIf(&mut_inplace_obn2ibn_, &con_inplace_obn2ibn_,
+                                            GetBlobDesc4BnInOp, parallel_ctx),
+                 std::stringstream()
+                     << " infer inplace obn to ibn is failed, op name " << op_->op_loc());
 }
 
 std::function<BlobDesc*(const std::string&)> ExecNode::GetBlobDesc4BnInOpFunc() const {

oneflow/core/graph/exec_graph.h

@@ -72,7 +72,9 @@ class ExecNode final : public Node<ExecNode, ExecEdge> {
   std::string VisualStr() const override { return op_->op_name(); }
   void ToProto(const ParallelContext*, ExecNodeProto*) const;
 
-  void InferBlobDescs(const ParallelContext* parallel_ctx);
+  typedef void (ExecNode::*InferBlobDescsMethod)(const ParallelContext*);
+  void InferBlobDescsByInputs(const ParallelContext* parallel_ctx);
+  void InferBlobDescsByNdSbp(const ParallelContext* parallel_ctx);
 
   const HashMap<std::string, std::string>& mut_inplace_obn2ibn() const {
     return mut_inplace_obn2ibn_;

oneflow/core/graph/nccl_send_recv_boxing_task_node.cpp

@@ -85,7 +85,7 @@ void NcclSendRecvBoxingTaskNode::BuildExecGphAndRegst() {
     node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
   }
   node->AddBnToRegstAndBindIt(&Operator::tmp_bns, GetProducedRegst("tmp"));
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void NcclSendRecvBoxingTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/slice_boxing_task_node.cpp

@@ -64,7 +64,7 @@ void SliceBoxingTaskNode::BuildExecGphAndRegst() {
   out_regst->AddLbi(lbi());
   node->BindBnWithRegst(op->SoleObn(), out_regst);
   node->AddBnToRegstAndBindIt(&Operator::tmp_bns, GetProducedRegst("tmp"));
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void SliceBoxingTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph/task_node.h

@@ -71,6 +71,9 @@ class TaskNode : public Node<TaskNode, TaskEdge> {
   DeviceType device_type() const;
   virtual const ParallelContext* parallel_ctx() const { return nullptr; }
 
+  // Different types of TaskNode/Compile Mode choose different output BlobDesc inference methods
+  virtual ExecNode::InferBlobDescsMethod GetInferBlobDescsMethod() const = 0;
+
   // Setters
   void set_machine_id(int64_t val);
   void set_thrd_id(int64_t val);

oneflow/core/graph/transport_task_node.h

@@ -37,9 +37,16 @@ class TransportTaskNode : public TaskNode {
                                            const TaskGraphRebuildCtx& ctx);
   void ToTransportTaskProtoIf(TransportTaskProto*) const;
 
+  ExecNode::InferBlobDescsMethod GetInferBlobDescsMethod() const override {
+    // TransportTaskNode infers output BlobDesc based on input BlobDesc, because it can't infers
+    // output BlobDesc with SBP.
+    return &ExecNode::InferBlobDescsByInputs;
+  }
+
  private:
   virtual Maybe<void> InitTransportTaskFromProto(const TransportTaskProto&,
                                                  const TaskGraphRebuildCtx& ctx) = 0;
+
   virtual void ToTransportTaskProto(TransportTaskProto*) const = 0;
   LogicalBlobId lbi_;
 };

oneflow/core/graph_impl/acc_compute_task_node.cpp

@@ -44,7 +44,7 @@ void AccCompTaskNode::BuildExecGphAndRegst() {
   exec_node->BindBnWithRegst(op()->SoleIbn(), in_regst);
   out_regst->AddLbi(op()->BnInOp2Lbi(op()->SoleObn()));
   exec_node->BindBnWithRegst(op()->SoleObn(), out_regst);
-  exec_node->InferBlobDescs(parallel_ctx());
+  (exec_node->*GetInferBlobDescsMethod())(parallel_ctx());
   out_regst->ForEachLbi([out_regst](const LogicalBlobId& lbi) {
     const BlobDesc* blob_desc = out_regst->GetBlobDesc(lbi);
     CHECK_EQ(blob_desc->is_dynamic(), false);

oneflow/core/graph_impl/acc_ctrl_tick_compute_task_node.cpp

@@ -47,7 +47,7 @@ void AccCtrlTickCompTaskNode::BuildExecGphAndRegst() {
   exec_node->BindBnWithRegst(op->SoleIbn(), in_regst);
   out_regst->AddLbi(op->BnInOp2Lbi(op->SoleObn()));
   exec_node->BindBnWithRegst(op->SoleObn(), out_regst);
-  exec_node->InferBlobDescs(parallel_ctx());
+  (exec_node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_COMP_TASK_STREAM_INDEX_GETTER(TaskType::kAccCtrlTick);

oneflow/core/graph_impl/acc_tick_compute_task_node.cpp

@@ -47,7 +47,7 @@ void AccTickCompTaskNode::BuildExecGphAndRegst() {
   exec_node->BindBnWithRegst(op->SoleIbn(), in_regst);
   out_regst->AddLbi(op->BnInOp2Lbi(op->SoleObn()));
   exec_node->BindBnWithRegst(op->SoleObn(), out_regst);
-  exec_node->InferBlobDescs(parallel_ctx());
+  (exec_node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_COMP_TASK_STREAM_INDEX_GETTER(TaskType::kAccTick);

oneflow/core/graph_impl/case_compute_task_node.cpp

@@ -69,7 +69,7 @@ void CaseCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(sole_op->BnInOp2Lbi(name));
     node->BindBnWithRegst(name, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void CaseCompTaskNode::InferProducedDataRegstTimeShape() { NaiveInferProducedDataRegstTimeShape(); }

oneflow/core/graph_impl/critical_section_wait_compute_task_node.cpp

@@ -56,7 +56,7 @@ void CriticalSectionWaitTickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_INDEPENDENT_TASK_STREAM_INDEX_GETTER(TaskType::kCriticalSectionWaitTick);

oneflow/core/graph_impl/decode_h2d_compute_task_node.cpp

@@ -54,7 +54,7 @@ void DecodeH2DCompTaskNode::BuildExecGphAndRegst() {
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
   node->AddBnToRegstAndBindIt(&Operator::tmp_bns, GetProducedRegst("tmp"));
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_NAMED_TASK_STREAM_INDEX_GETTER(DeviceType::kCUDA, TaskType::kDecodeH2D, "DECODE_H2D")

oneflow/core/graph_impl/device_tick_compute_task_node.cpp

@@ -56,7 +56,7 @@ void DeviceTickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_COMP_TASK_STREAM_INDEX_GETTER(TaskType::kDeviceTick);

oneflow/core/graph_impl/distribute_concat_compute_task_node.cpp

@@ -52,7 +52,8 @@ void DistributeConcatCompTaskNode::ConsumeAllRegsts() {
 void DistributeConcatCompTaskNode::BuildExecGphAndRegst() {
   BuildExecGphStructAndBindInRegst();
   BuildOutRegst();
-  mut_exec_gph().TopoForEachNode([this](ExecNode* node) { node->InferBlobDescs(parallel_ctx()); });
+  mut_exec_gph().TopoForEachNode(
+      [this](ExecNode* node) { (node->*GetInferBlobDescsMethod())(parallel_ctx()); });
 }
 
 void DistributeConcatCompTaskNode::BuildExecGphStructAndBindInRegst() {

oneflow/core/graph_impl/distribute_split_compute_task_node.cpp

@@ -47,7 +47,8 @@ void DistributeSplitCompTaskNode::ConsumeAllRegsts() {
 void DistributeSplitCompTaskNode::BuildExecGphAndRegst() {
   BuildExecGphStructAndBindInRegst();
   BuildOutRegst();
-  mut_exec_gph().TopoForEachNode([this](ExecNode* node) { node->InferBlobDescs(parallel_ctx()); });
+  mut_exec_gph().TopoForEachNode(
+      [this](ExecNode* node) { (node->*GetInferBlobDescsMethod())(parallel_ctx()); });
 }
 
 void DistributeSplitCompTaskNode::BuildExecGphStructAndBindInRegst() {

oneflow/core/graph_impl/dst_subset_tick_compute_task_node.cpp

@@ -56,7 +56,7 @@ void DstSubsetTickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_TICK_TASK_STREAM_INDEX_GETTER(TaskType::kDstSubsetTick);

oneflow/core/graph_impl/esac_compute_task_node.cpp

@@ -70,7 +70,7 @@ void EsacCompTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi("out"));
   node->BindBnWithRegst("out", out_regst);
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void EsacCompTaskNode::InferProducedDataRegstTimeShape() { NaiveInferProducedDataRegstTimeShape(); }

oneflow/core/graph_impl/normal_forward_compute_task_node.cpp

@@ -99,7 +99,8 @@ void NormalForwardCompTaskNode::BuildExecGphAndRegst() {
   BuildExecGphStructAndBindInRegst();
   BuildOutRegst();
   BuildTmp7BufRegsts();
-  mut_exec_gph().TopoForEachNode([this](ExecNode* node) { node->InferBlobDescs(parallel_ctx()); });
+  mut_exec_gph().TopoForEachNode(
+      [this](ExecNode* node) { (node->*GetInferBlobDescsMethod())(parallel_ctx()); });
 }
 
 void NormalForwardCompTaskNode::BuildExecGphStructAndBindInRegst() {

oneflow/core/graph_impl/pack_compute_task_node.cpp

@@ -50,7 +50,7 @@ void PackCompTaskNode::BuildExecGphAndRegst() {
   out_regst->AddLbi(op()->BnInOp2Lbi(op()->SoleObn()));
   exec_node->BindBnWithRegst(op()->SoleObn(), out_regst);
 
-  exec_node->InferBlobDescs(parallel_ctx());
+  (exec_node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_COMP_TASK_STREAM_INDEX_GETTER(TaskType::kPack);

oneflow/core/graph_impl/reentrant_lock_compute_task_node.cpp

@@ -56,7 +56,7 @@ void ReentrantLockCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void ReentrantLockCompTaskNode::InferProducedDataRegstTimeShape() {

oneflow/core/graph_impl/repeat_compute_task_node.cpp

@@ -51,7 +51,7 @@ void RepeatCompTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(sole_op->BnInOp2Lbi(sole_op->SoleObn()));
   node->BindBnWithRegst(sole_op->SoleObn(), out_regst);
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 
   // NOTE(chengcheng): force inplace
   CHECK_EQ(in_regst->NumOfLbi(), 1);

oneflow/core/graph_impl/source_tick_compute_task_node.cpp

@@ -46,7 +46,7 @@ void SourceTickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_TICK_TASK_STREAM_INDEX_GETTER(TaskType::kSourceTick);

oneflow/core/graph_impl/src_subset_tick_compute_task_node.cpp

@@ -56,7 +56,7 @@ void SrcSubsetTickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_TICK_TASK_STREAM_INDEX_GETTER(TaskType::kSrcSubsetTick);

oneflow/core/graph_impl/ssp_variable_proxy_task_node.cpp

@@ -74,7 +74,7 @@ class SspVariableProxyCompTaskNode final : public CompTaskNode {
     BuildExecGphStructAndBindInRegst();
     BuildOutRegst();
     mut_exec_gph().TopoForEachNode(
-        [this](ExecNode* node) { node->InferBlobDescs(parallel_ctx()); });
+        [this](ExecNode* node) { (node->*GetInferBlobDescsMethod())(parallel_ctx()); });
   }
 
   void BuildExecGphStructAndBindInRegst() {

oneflow/core/graph_impl/tick_compute_task_node.cpp

@@ -56,7 +56,7 @@ void TickCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_TICK_TASK_STREAM_INDEX_GETTER(TaskType::kTick);

oneflow/core/graph_impl/unpack_compute_task_node.cpp

@@ -50,7 +50,7 @@ void UnpackCompTaskNode::BuildExecGphAndRegst() {
   std::shared_ptr<RegstDesc> out_regst = GetProducedRegst("out");
   out_regst->AddLbi(op()->BnInOp2Lbi(op()->SoleObn()));
   exec_node->BindBnWithRegst(op()->SoleObn(), out_regst);
-  exec_node->InferBlobDescs(parallel_ctx());
+  (exec_node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 REGISTER_COMP_TASK_STREAM_INDEX_GETTER(TaskType::kUnpack);

oneflow/core/graph_impl/wait_and_send_ids_compute_task_node.cpp

@@ -49,7 +49,7 @@ void WaitAndSendIdsCompTaskNode::BuildExecGphAndRegst() {
     out_regst->AddLbi(lbi);
     node->BindBnWithRegst(obn, out_regst);
   }
-  node->InferBlobDescs(parallel_ctx());
+  (node->*GetInferBlobDescsMethod())(parallel_ctx());
 }
 
 void WaitAndSendIdsCompTaskNode::InferProducedDataRegstTimeShape() {