Replace src tick with with wait and send ids

oneflow/core/job_rewriter/autotick.cpp

@@ -427,6 +427,8 @@ Maybe<void> MultiClientAddWaitAndSendIds(JobBuilder* job_builder, int64_t machin
     parallel_conf.set_device_tag("cpu");
     parallel_conf.add_device_name(std::string("@") + std::to_string(machine_id) + ":0");
   }
+
+  // add wait_and_send_ids op conf
   OperatorConf wait_and_send_ids_op_conf;
   {
     wait_and_send_ids_op_conf.set_name(std::string("System-Src-WaitAndSendIds_") + NewUniqueId());
@@ -438,14 +440,33 @@ Maybe<void> MultiClientAddWaitAndSendIds(JobBuilder* job_builder, int64_t machin
     // wait_and_send_ids_conf->id_list() is unused in multi-client mode.
   }
   JUST(job_builder->AddOp(parallel_conf, wait_and_send_ids_op_conf));
-  OperatorConf source_tick_op = JUST(job_builder->OpConf4OpName(src_op_name));
-  {
-    CHECK_OR_RETURN(source_tick_op.has_source_tick_conf());
-    auto* source_tick_op_conf = source_tick_op.mutable_source_tick_conf();
-    CHECK_OR_RETURN(!source_tick_op_conf->has_wait_in());
-    source_tick_op_conf->set_wait_in(GenLogicalBlobName(wait_and_send_ids_op_conf.name(), "out"));
-  }
-  JUST(job_builder->MutOpOnlyOnce(source_tick_op));
+
+  // connect wait_and_send_ids to tick op which was connected to the src tick op
+  OperatorConf tick_op_conf;
+  bool find_src_tick_consumer_tick = false;
+  JUST(job_builder->ForEachOperator([&](const Operator& op) -> Maybe<void> {
+    // skip if the op is not a tick op
+    if (!op.op_conf().has_tick_conf()) { return Maybe<void>::Ok(); }
+    for (const auto& ibn : op.input_bns()) {
+      const auto& input_lbi = op.BnInOp2Lbi(ibn);
+      if (input_lbi.op_name() == src_op_name) {
+        CHECK_OR_RETURN(!find_src_tick_consumer_tick);
+        tick_op_conf.CopyFrom(op.op_conf());
+        find_src_tick_consumer_tick = true;
+      }
+    }
+    return Maybe<void>::Ok();
+  }));
+  CHECK_OR_RETURN(find_src_tick_consumer_tick);
+  CHECK_OR_RETURN(tick_op_conf.has_tick_conf());
+  CHECK_EQ_OR_RETURN(tick_op_conf.tick_conf().tick_size(), 1);
+  tick_op_conf.mutable_tick_conf()->clear_tick();
+  tick_op_conf.mutable_tick_conf()->add_tick(
+      GenLogicalBlobName(wait_and_send_ids_op_conf.name(), "out"));
+  JUST(job_builder->MutOpOnlyOnce(tick_op_conf));
+
+  // erase the src tick op
+  job_builder->DelOps({src_op_name});
   return Maybe<void>::Ok();
 }
 

oneflow/core/kernel/input_kernel.cpp

@@ -31,8 +31,6 @@ class InputKernel final : public KernelIf<device_type> {
   ~InputKernel() = default;
 
  private:
-  void Forward(const KernelCtx& ctx,
-               std::function<Blob*(const std::string&)> BnInOp2Blob) const override {}
   void ForwardDataContent(const KernelCtx& ctx,
                           std::function<Blob*(const std::string&)> BnInOp2Blob) const override {
     if (CHECK_JUST(*Global<Maybe<bool>, MultiClient>::Get())) {

oneflow/core/operator/op_conf.proto

@@ -220,8 +220,7 @@ message DstSubsetTickOpConf {
 }
 
 message SourceTickOpConf {
-  optional string wait_in = 1;
-  required string out = 2;
+  required string out = 1;
 }
 
 message SinkTickOpConf {

oneflow/core/operator/source_tick_op.cpp

@@ -21,7 +21,6 @@ namespace oneflow {
 Maybe<void> SourceTickOp::InitFromOpConf() {
   CHECK(op_conf().has_source_tick_conf());
   CHECK(op_conf().ctrl_in_op_name().empty());
-  if (op_conf().source_tick_conf().has_wait_in()) { EnrollInputBn("wait_in", false); }
   EnrollOutputBn("out", false);
   return Maybe<void>::Ok();
 }
@@ -46,8 +45,7 @@ Maybe<void> SourceTickOp::InferOutBlobDescs(
 }
 
 Maybe<void> SourceTickOp::GetSbpSignatures(cfg::SbpSignatureList* sbp_sig_list) const {
-  auto* sbp_signature = sbp_sig_list->mutable_sbp_signature()->Add();
-  SbpSignatureBuilder().Broadcast(input_bns()).Broadcast(output_bns()).Build(sbp_signature);
+  SbpSignatureBuilder().Broadcast(output_bns()).Build(sbp_sig_list->mutable_sbp_signature()->Add());
   return Maybe<void>::Ok();
 }
 

python/oneflow/test/graph/test_graph_relu.py

@@ -13,25 +13,21 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-
-import os
 import unittest
-
 import numpy as np
 
 import oneflow as flow
-import oneflow.framework.graph_build_util as graph_build_util
 import oneflow.unittest
 
 
-@unittest.skip(" nn.Graph cannnot run right now ")
-class TestReluGraph(flow.unittest.TestCase):
+class TestReluGraph(oneflow.unittest.TestCase):
     def test_relu_graph(test_case):
         data = np.array([2.0, 1.0, 0.0, -1.0, -2.0])
         x = flow.tensor(data, dtype=flow.float32)
+
         MyRelu = flow.nn.ReLU()
         y_eager = MyRelu(x)
-        print("eager out :", y_eager)
+        # print("eager out :", y_eager)
 
         class ReluGraph(flow.nn.Graph):
             def __init__(self):
@@ -42,8 +38,9 @@ def build(self, x):
                 return self.cc_relu(x)
 
         relu_g = ReluGraph()
-        y_lazy = relu_g(x)[0]
-        print("lazy out :", y_lazy)
+        y_lazy = relu_g(x)
+        # print(f"type of lazy y: {type(y_lazy)}")
+        # print(f"lazy y shape: {y_lazy.shape}, data: {y_lazy}")
         test_case.assertTrue(np.array_equal(y_eager.numpy(), y_lazy.numpy()))