Fix FillOpParallelDesc order (#4453)

* Fix FillOpParallelDesc order

* CHECK_JUST
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/core/kernel/runtime_blob_shape_infer_helper.cpp

@@ -26,15 +26,15 @@ RuntimeBlobShapeInferHelper::RuntimeBlobShapeInferHelper(const OperatorConf& op_
                                                          const JobDesc* job_desc) {
   op_ = ConstructOp(op_conf);
   const OpAttribute& op_attribute = kernel_conf.op_attribute();
+  if (op_attribute.has_parallel_conf_signature()
+      && op_attribute.parallel_conf_signature().has_op_parallel_conf()) {
+    CHECK_JUST(op_->FillOpParallelDesc(
+        ParallelDesc(op_attribute.parallel_conf_signature().op_parallel_conf())));
+  }
   if (op_attribute.has_sbp_signature()) {
     sbp_signature_.reset(new SbpSignature(op_attribute.sbp_signature()));
     CHECK_JUST(op_->FillSbpSignature(*sbp_signature_));
   }
-  if (op_attribute.has_parallel_conf_signature()
-      && op_attribute.parallel_conf_signature().has_op_parallel_conf()) {
-    op_->FillOpParallelDesc(
-        ParallelDesc(op_attribute.parallel_conf_signature().op_parallel_conf()));
-  }
   op_->ForEachBnInOp([&](const std::string& bn_in_op) { bn_in_op2blob_desc_[bn_in_op].reset(); });
   if (op_attribute.has_logical_blob_desc_signature()) {
     HashMap<std::string, std::unique_ptr<BlobDesc>> bn_in_op2logical_blob_desc;

oneflow/core/operator/operator.cpp

@@ -478,7 +478,7 @@ void Operator::ForEachBnInOp(std::function<void(const std::string&)> Handler) co
 Maybe<void> Operator::FillSbpSignature(const SbpSignature& sbp_signature) {
   ParallelDistributionSignature parallel_distribution_signature;
   SbpSignatureToParallelDistributionSignature(sbp_signature, &parallel_distribution_signature);
-  FillParallelDistributionSignature(parallel_distribution_signature);
+  JUST(FillParallelDistributionSignature(parallel_distribution_signature));
   return Maybe<void>::Ok();
 }
 
@@ -698,7 +698,7 @@ Maybe<void> Operator::InferParallelDistributionSignature(
     SbpSignatureToParallelDistributionSignature(sbp_signature, signature);
     return Maybe<void>::Ok();
   } else {
-    UNIMPLEMENTED();
+    UNIMPLEMENTED_THEN_RETURN();
   }
 }