Sarkars/flib ngenc compute

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,30 @@ 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");
+      if (fdef == nullptr) {
+        OP_REQUIRES_OK(
+            ctx, 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:

test/python/common.py

@@ -32,7 +32,7 @@
 
 class NgraphTest(object):
 
-    def with_ngraph(self, l, config=tf.ConfigProto()):
+    def with_ngraph(self, l, config=tf.ConfigProto(), graph=None):
         if ngraph_bridge.is_grappler_enabled():
             rewrite_options = rewriter_config_pb2.RewriterConfig(
                 meta_optimizer_iterations=rewriter_config_pb2.RewriterConfig.
@@ -50,7 +50,7 @@ def with_ngraph(self, l, config=tf.ConfigProto()):
 
         os.environ['NGRAPH_TF_DISABLE_DEASSIGN_CLUSTERS'] = '1'
         ngraph_bridge.enable()
-        with tf.Session(config=config) as sess:
+        with tf.Session(graph=graph, config=config) as sess:
             retval = l(sess)
 
         os.environ.pop('NGRAPH_TF_DISABLE_DEASSIGN_CLUSTERS', None)
@@ -61,12 +61,12 @@ def with_ngraph(self, l, config=tf.ConfigProto()):
 
         return retval
 
-    def without_ngraph(self, l, config=tf.ConfigProto()):
+    def without_ngraph(self, l, config=tf.ConfigProto(), graph=None):
         ngraph_tf_disable_deassign_clusters = os.environ.pop(
             'NGRAPH_TF_DISABLE_DEASSIGN_CLUSTERS', None)
 
         ngraph_bridge.disable()
-        with tf.Session(config=config) as sess:
+        with tf.Session(graph=graph, config=config) as sess:
             retval = l(sess)
 
         if ngraph_tf_disable_deassign_clusters is not None:

test/python/flib_graph_1.pbtxt

@@ -0,0 +1,284 @@
+node {
+  name: "Placeholder_2"
+  op: "Placeholder"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "shape"
+    value {
+      shape {
+        dim {
+          size: 2
+        }
+        dim {
+          size: 3
+        }
+      }
+    }
+  }
+}
+node {
+  name: "Placeholder_1"
+  op: "Placeholder"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "shape"
+    value {
+      shape {
+        dim {
+          size: 2
+        }
+        dim {
+          size: 3
+        }
+      }
+    }
+  }
+}
+node {
+  name: "Placeholder"
+  op: "Placeholder"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "shape"
+    value {
+      shape {
+        dim {
+          size: 2
+        }
+        dim {
+          size: 3
+        }
+      }
+    }
+  }
+}
+node {
+  name: "ngraph_cluster_0"
+  op: "NGraphEncapsulate"
+  input: "Placeholder"
+  input: "Placeholder_1"
+  input: "Placeholder_2"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "Targuments"
+    value {
+      list {
+        type: DT_FLOAT
+        type: DT_FLOAT
+        type: DT_FLOAT
+      }
+    }
+  }
+  attr {
+    key: "Tresults"
+    value {
+      list {
+        type: DT_FLOAT
+        type: DT_FLOAT
+      }
+    }
+  }
+  attr {
+    key: "_ngraph_backend"
+    value {
+      s: "CPU"
+    }
+  }
+  attr {
+    key: "ngraph_cluster"
+    value {
+      i: 0
+    }
+  }
+  attr {
+    key: "ngraph_graph_id"
+    value {
+      i: 0
+    }
+  }
+}
+node {
+  name: "add_1"
+  op: "IdentityN"
+  input: "ngraph_cluster_0"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "T"
+    value {
+      list {
+        type: DT_FLOAT
+      }
+    }
+  }
+}
+node {
+  name: "Sigmoid"
+  op: "IdentityN"
+  input: "ngraph_cluster_0:1"
+  device: "/job:localhost/replica:0/task:0/device:CPU:0"
+  attr {
+    key: "T"
+    value {
+      list {
+        type: DT_FLOAT
+      }
+    }
+  }
+}
+library {
+  function {
+    signature {
+      name: "Enc_0_native_segment"
+      input_arg {
+        name: "ngraph_input_0"
+        type: DT_FLOAT
+      }
+      input_arg {
+        name: "ngraph_input_1"
+        type: DT_FLOAT
+      }
+      input_arg {
+        name: "ngraph_input_2"
+        type: DT_FLOAT
+      }
+      output_arg {
+        name: "ngraph_output_0"
+        type: DT_FLOAT
+      }
+      output_arg {
+        name: "ngraph_output_1"
+        type: DT_FLOAT
+      }
+    }
+    node_def {
+      name: "add"
+      op: "Add"
+      input: "ngraph_input_0"
+      input: "ngraph_input_1"
+      device: "/job:localhost/replica:0/task:0/device:CPU:0"
+      attr {
+        key: "T"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "_ngraph_backend"
+        value {
+          s: "CPU"
+        }
+      }
+      attr {
+        key: "_ngraph_cluster"
+        value {
+          i: 0
+        }
+      }
+      attr {
+        key: "_ngraph_marked_for_clustering"
+        value {
+          b: true
+        }
+      }
+      experimental_debug_info {
+        original_node_names: "add"
+      }
+    }
+    node_def {
+      name: "add_1_ngraph/_0"
+      op: "Add"
+      input: "add:z:0"
+      input: "ngraph_input_2"
+      device: "/job:localhost/replica:0/task:0/device:CPU:0"
+      attr {
+        key: "T"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "_ngraph_backend"
+        value {
+          s: "CPU"
+        }
+      }
+      attr {
+        key: "_ngraph_cluster"
+        value {
+          i: 0
+        }
+      }
+      attr {
+        key: "_ngraph_marked_for_clustering"
+        value {
+          b: true
+        }
+      }
+      experimental_debug_info {
+        original_node_names: "add_1_ngraph/_0"
+      }
+    }
+    node_def {
+      name: "Sigmoid_ngraph/_1"
+      op: "Sigmoid"
+      input: "add_1_ngraph/_0:z:0"
+      device: "/job:localhost/replica:0/task:0/device:CPU:0"
+      attr {
+        key: "T"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "_ngraph_backend"
+        value {
+          s: "CPU"
+        }
+      }
+      attr {
+        key: "_ngraph_cluster"
+        value {
+          i: 0
+        }
+      }
+      attr {
+        key: "_ngraph_marked_for_clustering"
+        value {
+          b: true
+        }
+      }
+      experimental_debug_info {
+        original_node_names: "Sigmoid_ngraph/_1"
+      }
+    }
+    ret {
+      key: "ngraph_output_0"
+      value: "add_1_ngraph/_0:z:0"
+    }
+    ret {
+      key: "ngraph_output_1"
+      value: "Sigmoid_ngraph/_1:y:0"
+    }
+  }
+}
+versions {
+  producer: 27
+}