[Graph] Support IteratorGetNext for SmartStage as a starting node for (…

…DeepRec-AI#901) searching. Support IteratorGetNext node as the starting node for searching the boundary of stage subgraph. Signed-off-by: chenbangduo.cbd <chenbangduo.cbd@alibaba-inc.com>
BrunoScaglione · Jun 30, 2023 · 7a4cab4 · 7a4cab4
1 parent 26c2755
commit 7a4cab4
Show file tree

Hide file tree

Showing 6 changed files with 528 additions and 134 deletions.
diff --git a/tensorflow/cc/BUILD b/tensorflow/cc/BUILD
@@ -742,7 +742,7 @@ cc_library(
     hdrs = ["training/prefetch_runner.h"],
     deps = [
         ":coordinator",
-        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_base_no_ops",
     ],
 )
 

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
@@ -3504,6 +3504,7 @@ tf_cuda_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "//third_party/eigen3",
+        "//tensorflow/cc:prefetch_runner",
         "//tensorflow/core/grappler/utils:functions",
         "//tensorflow/core/profiler:nvtx_utils",
         "//tensorflow/core/profiler/lib:traceme",