[runtime] dynamic shape with AOT memory planing

python/setup.py

@@ -129,6 +129,7 @@ def is_pure(self):
       install_requires=[
         'numpy',
         'decorator',
+        'attrs'
         ],
       packages=find_packages(),
       distclass=BinaryDistribution,

python/tvm/contrib/graph_runtime.py

@@ -113,6 +113,8 @@ def __init__(self, module):
         self._get_input = module["get_input"]
         self._get_num_outputs = module["get_num_outputs"]
         self._load_params = module["load_params"]
+        self._set_shape_variable = module["set_shape_variable"]
+        self._update_view = module["update_view"]
 
     def set_input(self, key=None, value=None, **params):
         """Set inputs to the module via kwargs
@@ -138,6 +140,18 @@ def set_input(self, key=None, value=None, **params):
             for k in keys:
                 self._get_input(k).copyfrom(params[k])
 
+    def set_shape_variable(self, var_bounds):
+        """Set variables value in symbolic shape
+
+        Parameters
+        ----------
+        var_bounds : dict pf str to int
+            shape variable bounds : value
+        """
+        for k, v in var_bounds.items():
+            self._set_shape_variable(k, v)
+        self._update_view()
+
     def run(self, **input_dict):
         """Run forward execution of the graph
 

python/tvm/relay/backend/graph_runtime_codegen.py

@@ -22,6 +22,7 @@
 import json
 from collections import defaultdict
 import attr
+import tvm
 from . import _backend
 from . import compile_engine
 from ..op import Op
@@ -90,27 +91,37 @@ def to_json(self):
         }
 
 
-def shape_to_json(shape):
-    """Convert symbolic shape to json compatible forma."""
-    return [sh.value for sh in shape]
-
+def shape_to_json(shape, var_bound_dict, bound_hint):
+    """Convert symbolic shape to json compatible format."""
+    ret = []
+    for dim in shape:
+        if var_bound_dict is None:
+            ret.append(dim.value)
+        else:
+            ret.append(str(dim))
+        if isinstance(dim, tvm.expr.Var):
+            assert bound_hint is not None
+            var_bound_dict[str(dim)] = bound_hint[dim]
+    return ret
 
 class GraphRuntimeCodegen(ExprFunctor):
     """The compiler from Relay to the TVM runtime system."""
     nodes = attr.ib()
     var_map = attr.ib()
 
-    def __init__(self, mod, target):
+    def __init__(self, mod, target, shape_var_bounds=None):
         ExprFunctor.__init__(self)
         self.mod = mod
         self.target = target
         self.nodes = []
         self.var_map = {}
         self.params = {}
+        self.sym_shape = False
         self.storage_device_map = None
         self.compile_engine = compile_engine.get()
         self.lowered_funcs = defaultdict(set)
         self._name_map = {}
+        self.shape_var_bounds = shape_var_bounds
 
     def add_node(self, node, expr):
         """
@@ -133,8 +144,9 @@ def add_node(self, node, expr):
         # setup storage ids
         assert expr in self.storage_device_map
         storage_device_info = self.storage_device_map[expr]
-        assert len(storage_device_info) == 2
+        assert len(storage_device_info) == 3
         node.attrs["storage_id"] = [x.value for x in storage_device_info[0]]
+        node.attrs["max_bytes"] = [x.value for x in storage_device_info[2]]
         device_types = [x.value for x in storage_device_info[1]]
         num_unknown_devices = device_types.count(0)
         if num_unknown_devices != 0 and num_unknown_devices != len(device_types):
@@ -146,6 +158,11 @@ def add_node(self, node, expr):
         if num_unknown_devices == 0:
             node.attrs["device_index"] = device_types
 
+        if self.sym_shape:
+            var_bound_dict = {}
+        else:
+            var_bound_dict = None
+
         node_id = len(self.nodes)
         self.nodes.append(node)
         # Tuple return value, flatten as tuple
@@ -157,7 +174,9 @@ def add_node(self, node, expr):
                 if not isinstance(typ, TensorType):
                     raise RuntimeError("type %s not supported" % typ)
                 ret.append(NodeRef(node_id, i))
-                shape.append(shape_to_json(typ.shape))
+                shape.append(shape_to_json(typ.shape,
+                                           var_bound_dict,
+                                           self.shape_var_bounds))
                 dtype.append(typ.dtype)
             node.attrs["shape"] = shape
             node.attrs["dtype"] = dtype
@@ -167,8 +186,12 @@ def add_node(self, node, expr):
         # Normal tensor return type
         if not isinstance(checked_type, TensorType):
             raise RuntimeError("type %s not supported" % checked_type)
-        node.attrs["shape"] = [shape_to_json(checked_type.shape)]
+        node.attrs["shape"] = [shape_to_json(checked_type.shape,
+                                             var_bound_dict,
+                                             self.shape_var_bounds)]
         node.attrs["dtype"] = [checked_type.dtype]
+        if self.sym_shape:
+            node.attrs["var_upper_bound"] = var_bound_dict.items()
         node.num_outputs = 1
         return NodeRef(node_id, 0)
 
@@ -312,23 +335,32 @@ def _get_json(self):
         num_entry = 0
         shapes = []
         storage_ids = []
+        max_bytes = []
         device_types = []
         dltypes = []
+        var_upper_bound = []
         node_row_ptr = [0]
         for node in self.nodes:
             assert node.num_outputs == len(node.attrs["shape"])
             shapes += node.attrs["shape"]
             dltypes += node.attrs["dtype"]
             storage_ids += node.attrs["storage_id"]
+            max_bytes += node.attrs["max_bytes"]
             if "device_index" in node.attrs:
                 device_types += node.attrs["device_index"]
             num_entry += node.num_outputs
             node_row_ptr.append(num_entry)
+            if self.sym_shape:
+                var_upper_bound += node.attrs["var_upper_bound"]
 
         # Compute "attrs" field.
         attrs = {}
+        attrs["symbolic_shape"] = 1 if self.sym_shape else 0
         attrs["shape"] = ["list_shape", shapes]
         attrs["storage_id"] = ["list_int", storage_ids]
+        if self.sym_shape:
+            attrs["max_bytes"] = ["list_int", max_bytes]
+            attrs["var_upper_bound"] = dict(var_upper_bound)
         if device_types:
             attrs["device_index"] = ["list_int", device_types]
         attrs["dltype"] = ["list_str", dltypes]
@@ -383,8 +415,17 @@ def codegen(self, func):
         params : Dict[str, tvm.nd.NDArray]
             Additional constant parameters.
         """
-        self.storage_device_map = _backend.GraphPlanMemory(func)
-        # First we convert all the parameters into input nodes.
+        self.storage_device_map = _backend.GraphPlanMemoryEx(func,
+                                                             self.shape_var_bounds)
+        # First we check whether there is symbolic shape
+        for param in func.params:
+            if self.sym_shape:
+                break
+            for dim in param._checked_type_.shape:
+                if isinstance(dim, tvm.expr.Var):
+                    self.sym_shape = True
+                    break
+        # Second we convert all the parameters into input nodes.
         for param in func.params:
             node = InputNode(param.name_hint, {})
             self.var_map[param] = self.add_node(node, param)

python/tvm/relay/build_module.py

@@ -197,7 +197,11 @@ def optimize(func, target, params=None):
     return func
 
 
-def build(func, target=None, target_host=None, params=None):
+def build(func,
+          target=None,
+          target_host=None,
+          params=None,
+          shape_var_bounds=None):
     """Build a function to run on TVM graph runtime.
 
     Parameters
@@ -269,7 +273,9 @@ def build(func, target=None, target_host=None, params=None):
         func = ir_pass.fuse_ops(func, cfg.opt_level)
         # Graph code generation
         func = ir_pass.infer_type(func)
-        graph_gen = _graph_gen.GraphRuntimeCodegen(mod=None, target=target)
+        graph_gen = _graph_gen.GraphRuntimeCodegen(mod=None,
+                                                   target=target,
+                                                   shape_var_bounds=shape_var_bounds)
         graph_json, lowered_funcs, params = graph_gen.codegen(func)
         mod = _tvm_build_module(
             lowered_funcs, target=target, target_host=target_host)

src/relay/backend/graph_plan_memory.cc

@@ -4,6 +4,8 @@
  * \brief Memory index assignment pass for executing
  *   the program in the graph runtime.
  */
+#include <tvm/ir.h>
+#include <tvm/arithmetic.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/pass.h>
@@ -182,6 +184,11 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     return total;
   }
 
+  // init variable upper bound hint
+  void InitSymbolicShapeBound(const std::unordered_map<size_t, int64_t> &var_upper_bound) {
+    ssb_parser_.input_upper_bounds = var_upper_bound;
+  }
+
   // Run storage allocation for a function.
   Map<Expr, Array<IntegerArray> > Plan(const Function& func) {
     prototype_ = StorageAllocaInit(&arena_).GetInitTokenMap(func);
@@ -196,15 +203,19 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     for (const auto& kv : token_map_) {
       std::vector<Integer> storage_ids;
       std::vector<Integer> device_types;
+      std::vector<Integer> max_bytes;
       for (StorageToken* tok : kv.second) {
         if (tok->device_type) {
           num_annotated_nodes++;
         }
         num_nodes++;
         storage_ids.push_back(tok->storage_id);
         device_types.push_back(tok->device_type);
+        max_bytes.push_back(tok->max_bytes);
       }
-      smap.Set(GetRef<Expr>(kv.first), Array<IntegerArray>({storage_ids, device_types}));
+      smap.Set(GetRef<Expr>(kv.first), Array<IntegerArray>({storage_ids,
+                                                            device_types,
+                                                            max_bytes}));
     }
     // Either all or none of the nodes should be annotated.
     if (num_annotated_nodes != 0 && num_annotated_nodes != num_nodes) {
@@ -276,14 +287,12 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     CHECK(ttype != nullptr);
     size_t size = 1;
     for (IndexExpr dim : ttype->shape) {
-      const int64_t* pval = as_const_int(dim);
-      CHECK(pval != nullptr)
-          << "Cannot allocate memory symbolic tensor shape "
-          << ttype->shape;
-      CHECK_GE(*pval, 0)
+      // symbolic expr
+      const int64_t pval = ssb_parser_.EvalDim(dim);
+      CHECK_GE(pval, 0)
           << "Cannot allocate memory for tensor with negative shape"
-          << *pval;
-      size *= static_cast<size_t>(pval[0]);
+          << pval;
+      size *= static_cast<size_t>(pval);
     }
     size *= DivRoundUp(ttype->dtype.bits() * ttype->dtype.lanes(), 8);
     return size;
@@ -354,6 +363,35 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     }
   }
 
+ private:
+  //  Parse symbolic shape to its bound
+  struct SymShapeBoundParser {
+    int32_t EvalDim(const IndexExpr& expr) {
+      const Variable* vptr = expr.as<Variable>();
+      int32_t ret = -1;
+      if (vptr != nullptr) {
+        auto hash_key = expr.hash();
+        CHECK_GT(input_upper_bounds.count(hash_key), 0)
+          << "Input var " << vptr->name_hint << " upper bound missing";
+        ret = input_upper_bounds.at(hash_key);
+        CHECK_GT(ret, 0) << "Cannot allocate memory for pure symbolic tensor shape";
+        if (var_upper_bound.count(vptr) == 0) {
+          var_upper_bound[vptr] = tvm::arith::IntSet::single_point(ret);
+        }
+      } else {
+        const tvm::IntImm* ret_ptr = tvm::arith::EvalSet(expr, var_upper_bound)\
+          .point_value().as<IntImm>();
+        if (ret_ptr == nullptr) {
+          LOG(FATAL) << "Symbolic shape evaluation fail";
+        }
+        ret = ret_ptr->value;
+      }
+      return ret;
+    }
+    std::unordered_map<size_t, int64_t> input_upper_bounds;
+    std::unordered_map<const Variable*, tvm::arith::IntSet> var_upper_bound;
+  };
+
  private:
   // allocator
   common::Arena arena_;
@@ -365,14 +403,40 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
   std::vector<StorageToken*> data_;
   /*! \brief internal prototype token map */
   std::unordered_map<const ExprNode*, std::vector<StorageToken*> > prototype_;
+  /*! \brief symbolic shape bound parser */
+  SymShapeBoundParser ssb_parser_;
 };
 
+Map<Expr, Array<IntegerArray> > GraphPlanMemoryEx(const Function& func,
+                                                  const Map<tvm::Var, Integer> shape_var_bound) {
+  std::unordered_map<size_t, int64_t> var_bounds;
+  if (shape_var_bound.size()) {
+    for (const auto &kv : shape_var_bound) {
+      auto& k = kv.first;
+      auto& v = kv.second;
+      var_bounds[k.hash()] = v->value;
+    }
+  }
+  auto allocator = StorageAllocator();
+  allocator.InitSymbolicShapeBound(var_bounds);
+  return allocator.Plan(func);
+}
+
 Map<Expr, Array<IntegerArray> > GraphPlanMemory(const Function& func) {
-  return StorageAllocator().Plan(func);
+  std::unordered_map<size_t, int64_t> var_bounds;
+  auto allocator = StorageAllocator();
+  allocator.InitSymbolicShapeBound(var_bounds);
+  return allocator.Plan(func);
 }
 
+
+
+TVM_REGISTER_GLOBAL("relay.backend.GraphPlanMemoryEx")
+.set_body_typed<Map<Expr, Array<IntegerArray> > \
+  (const Function&, const Map<tvm::Var, Integer>)>(GraphPlanMemoryEx);
+
 TVM_REGISTER_GLOBAL("relay.backend.GraphPlanMemory")
-.set_body_typed<Map<Expr, Array<IntegerArray> >(const Function&)>(GraphPlanMemory);
+.set_body_typed<Map<Expr, Array<IntegerArray> > (const Function&)>(GraphPlanMemory);
 
 }  // namespace relay
 }  // namespace tvm