[COMPILER] GraphHash based cache system, allow dump and query duplica…

…ted functions. (apache#30)

nnvm/include/nnvm/graph.h

@@ -63,11 +63,11 @@ class Graph {
  * \return The indexed graph.
  * \sa IndexedGraph
  */
- const IndexedGraph& indexed_graph();
+ const IndexedGraph& indexed_graph() const;
 
  private:
  // internal structure of indexed graph
- std::shared_ptr<const IndexedGraph> indexed_graph_;
+ mutable std::shared_ptr<const IndexedGraph> indexed_graph_;
 };
 
 /*!

nnvm/include/nnvm/pass_functions.h

@@ -41,6 +41,17 @@ inline std::string SaveJSON(Graph graph) {
  return ret.GetAttr<std::string>("json");
 }
 
+
+/*!
+ * \brief Print graph ir
+ * \param graph The graph to be printed
+ * \return The graph ir string.
+ */
+inline std::string PrintGraphIR(Graph graph) {
+ Graph ret = ApplyPass(std::move(graph), "PrintGraphIR");
+ return ret.GetAttr<std::string>("graphir");
+}
+
 /*!
  * \brief Add control flow dependencies between nodes.
  *

nnvm/python/nnvm/compiler/__init__.py

@@ -5,6 +5,7 @@
 
 from . import build_module
 from . build_module import build, optimize, build_config
+from . compile_engine import engine, graph_key
 
 from .. import symbol as _symbol
 from .. import graph as _graph
@@ -14,5 +15,6 @@
 
 from .. import top as _top
 
+
 tvm.register_extension(_symbol.Symbol, _symbol.Symbol)
 tvm.register_extension(_graph.Graph, _graph.Graph)

nnvm/python/nnvm/compiler/build_module.py

@@ -184,7 +184,7 @@ def build(graph, target, shape, dtype="float32", params=None):
  graph._set_json_attr("target", target, "str")
  graph._set_json_attr("opt_level", cfg.opt_level, "int")
  graph = graph.apply("InferShape").apply("InferType")
- graph = graph.apply("GraphFusePartition").apply("GraphFuse")
+ graph = graph.apply("GraphFusePartition").apply("GraphFuseCompile")
  libmod = graph_attr._move_out_module(graph, "module")
  return graph, libmod, params
 

nnvm/python/nnvm/compiler/compile_engine.py

@@ -0,0 +1,99 @@
+# pylint: disable=invalid-name
+"""Compiler engine interface to internal engine"""
+import tvm
+
+_list_cache_items = tvm.get_global_func("nnvm.compiler.ListCacheItems")
+_clear_cache = tvm.get_global_func("nnvm.compiler.ClearCache")
+_get_cache_item = tvm.get_global_func("nnvm.compiler.GetCacheItem")
+_set_cache_item = tvm.get_global_func("nnvm.compiler.SetCacheItem")
+_graph_key_get_graph = tvm.get_global_func("nnvm.compiler.GraphKeyGetGraph")
+_make_graph_key = tvm.get_global_func("nnvm.compiler.MakeGraphKey")
+
+@tvm.register_node
+class GraphKey(tvm.node.NodeBase):
+ """Key of a graph compilation context"""
+ @property
+ def graph(self):
+ return _graph_key_get_graph(self)
+
+
+@tvm.register_node
+class GraphCacheEntry(tvm.node.NodeBase):
+ """CacheEntry of compilation into a TVM Function"""
+ pass
+
+
+@tvm.register_node
+class GraphFunc(tvm.node.NodeBase):
+ """Compiled result of a graph into a TVM Function"""
+ pass
+
+
+class Engine(object):
+ """Global singleton compilation engine."""
+ def items(self):
+ """List the available cache key value pairs.
+
+ Returns
+ -------
+ item_list : list of (GraphKey, GraphCacheEntry)
+ The existing cache items
+ """
+ res = _list_cache_items()
+ assert len(res) % 2 == 0
+ return [(res[2*i], res[2*i+1]) for i in range(len(res)/2)]
+
+ def clear_cache(self):
+ """Clear the existing cached functions."""
+ _clear_cache()
+
+ def __setitem__(self, key, value):
+ """Clear the existing cached functions."""
+ if isinstance(value, GraphCacheEntry):
+ _set_cache_item(key, value.graph_func)
+ else:
+ _set_cache_item(key, value)
+
+ def __getitem__(self, key):
+ """Clear the existing cached functions."""
+ return _get_cache_item(key)
+
+ def dump(self):
+ """Return a string representation of engine dump
+
+ Returns
+ -------
+ dump : str
+ The dumped string representation
+ """
+ items = self.items()
+ res = "====================================\n"
+ res += "CompilerEngine dump, %d items cached\n" % len(items)
+ for key, value in items:
+ res += "------------------------------------\n"
+ res += "target={}\n".format(key.target)
+ res += "inputs={}\n".format(key.inputs)
+ res += "use_count={}\n".format(value.use_count)
+ res += "func_name={}\n".format(value.graph_func.func_name)
+ res += key.graph.ir() + "\n"
+ res += "===================================\n"
+ return res
+
+engine = Engine()
+
+
+def graph_key(graph, inputs, target):
+ """Construct a new graph key.
+
+ Parameters
+ ----------
+ graph : Graph
+ The computation graph structure
+
+ inputs : list of Tensor(placeholder)
+ The input requirement to the graph.
+
+ target : str
+ The target of compilation.
+ """
+ return _make_graph_key(graph, inputs, target)

nnvm/python/nnvm/testing/__init__.py

@@ -1 +1,3 @@
 """Utilities for testcase"""
+
+from .config import ctx_list

nnvm/python/nnvm/testing/config.py

@@ -2,7 +2,7 @@
 import os
 import tvm
 
-def test_ctx_list():
+def ctx_list():
  """Get context list for testcases"""
  device_list = os.environ.get("NNVM_TEST_TARGETS", "")
  device_list = (device_list.split(",") if device_list