Sarkars/flib ngenc compute

src/grappler/ngraph_optimizer.cc

@@ -64,7 +64,8 @@ Status NgraphOptimizer::Optimize(tensorflow::grappler::Cluster* cluster,
   // we will not do anything; all subsequent
   // passes become a no-op.
   if (config::IsEnabled() == false ||
-      std::getenv("NGRAPH_TF_DISABLE") != nullptr) {
+      std::getenv("NGRAPH_TF_DISABLE") != nullptr ||
+      IsProcessedByNgraphPass(&graph)) {
     NGRAPH_VLOG(0) << "NGTF_OPTIMIZER: Ngraph is disabled ";
     graph.ToGraphDef(output);
     return Status::OK();

src/ngraph_cluster_manager.cc

@@ -34,7 +34,7 @@ int NGraphClusterManager::NewCluster() {
 
 GraphDef* NGraphClusterManager::GetClusterGraph(int idx) {
   std::lock_guard<std::mutex> guard(s_cluster_graphs_mutex);
-  return s_cluster_graphs[idx];
+  return idx < s_cluster_graphs.size() ? s_cluster_graphs[idx] : nullptr;
 }
 
 }  // namespace ngraph_bridge

src/ngraph_encapsulate_op.cc

@@ -18,6 +18,7 @@
 #include <utility>
 
 #include "tensorflow/core/common_runtime/dma_helper.h"
+#include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/common_runtime/optimization_registry.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/node_def_util.h"
@@ -93,9 +94,29 @@ class NGraphEncapsulateOp : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr<int>("ngraph_cluster", &m_ngraph_cluster));
     graph_def = NGraphClusterManager::GetClusterGraph(m_ngraph_cluster);
 
-    GraphConstructorOptions opts;
-    opts.allow_internal_ops = true;
-    OP_REQUIRES_OK(ctx, ConvertGraphDefToGraph(opts, *graph_def, &m_graph));
+    if (graph_def == nullptr) {
+      // Read graphdef from function library
+      const FunctionLibraryDefinition flib =
+          *ctx->function_library()->GetFunctionLibraryDefinition();
+      const FunctionDef* fdef =
+          flib.Find("Enc_" + to_string(m_ngraph_cluster) + "_native_segment");
+      OP_REQUIRES(
+          ctx, fdef != nullptr,
+          errors::Internal("Did not find graphdef for encapsulate ",
+                           m_ngraph_cluster,
+                           " in NGraphClusterManager or function library"));
+      // TODO: how to convert from functiondef to graphdef. Anything easier?
+      FunctionBody* fnbody;
+      const auto get_func_sig = [&flib](const string& op, const OpDef** sig) {
+        return flib.LookUpOpDef(op, sig);
+      };
+      FunctionDefToBodyHelper(*fdef, {}, &flib, get_func_sig, &fnbody);
+      CopyGraph(*fnbody->graph, &m_graph);
+    } else {
+      GraphConstructorOptions opts;
+      opts.allow_internal_ops = true;
+      OP_REQUIRES_OK(ctx, ConvertGraphDefToGraph(opts, *graph_def, &m_graph));
+    }
     OP_REQUIRES_OK(ctx, ctx->GetAttr("ngraph_graph_id", &m_graph_id));
     //
     // Initialize the "m_input_is_static" vector as follows:

src/ngraph_mark_for_clustering.cc

@@ -291,6 +291,7 @@ Status MarkForClustering(Graph* graph,
       confirmation_function_map["Minimum"] = SimpleConfirmationFunction();
       confirmation_function_map["Mul"] = SimpleConfirmationFunction();
       confirmation_function_map["Neg"] = SimpleConfirmationFunction();
+      confirmation_function_map["NoOp"] = SimpleConfirmationFunction();
       confirmation_function_map["OneHot"] = SimpleConfirmationFunction();
       confirmation_function_map["Pad"] = SimpleConfirmationFunction();
       confirmation_function_map["Pow"] = SimpleConfirmationFunction();

src/ngraph_utils.cc

@@ -378,7 +378,20 @@ void AllreduceOpControlOrder(
       }
     }
   }
-};
+}
+
+bool IsProcessedByNgraphPass(Graph* g) {
+  // TODO: place a dummy node as a marker
+  // Current method may fail when graph had no encapsulates or ngvars after
+  // first pass
+  // Also variable/optimizer change introduces other types of ng nodes
+  for (Node* node : g->nodes()) {
+    if (node->type_string() == "NGraphEncapsulate" ||
+        node->type_string() == "NGraphVariable")
+      return true;
+  }
+  return false;
+}
 
 }  // namespace ngraph_bridge
 

src/ngraph_utils.h

@@ -63,6 +63,8 @@ bool IsNGVariableType(string node_type);
 // Node-types that are executed on nGraph
 bool IsNGSupportedType(string node_type);
 
+bool IsProcessedByNgraphPass(Graph* g);
+
 // Taken from: tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
 // Extract values from a Const op to `values`. Returns true if succeeds.
 //

test/CMakeLists.txt

@@ -108,3 +108,8 @@ install(
     DESTINATION ${CMAKE_INSTALL_PREFIX}/test
     FILES_MATCHING PATTERN "*.py"
 )
+install(
+    DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/python
+    DESTINATION ${CMAKE_INSTALL_PREFIX}/test
+    FILES_MATCHING PATTERN "*.pbtxt"
+)

test/python/CMakeLists.txt

@@ -23,4 +23,15 @@ foreach(file ${files})
 )
 endforeach()
 
+execute_process(
+    COMMAND ${CMAKE_COMMAND} -E create_symlink
+        ${CMAKE_CURRENT_SOURCE_DIR}/flib_graph_1.pbtxt
+        ${CMAKE_CURRENT_BINARY_DIR}/flib_graph_1.pbtxt
+)
+
+execute_process(
+    COMMAND ${CMAKE_COMMAND} -E create_symlink
+        ${CMAKE_CURRENT_SOURCE_DIR}/flib_graph_2.pbtxt
+        ${CMAKE_CURRENT_BINARY_DIR}/flib_graph_2.pbtxt
+)