Merge

Author: mbartling

utensor_cgen/backend/graph_lower/alloc_plan.py

@@ -53,7 +53,7 @@ def __contains__(self, offset):
 
 @attr.s
 class TimeSpaceAllocation(object):
-  entity_name = attr.ib(validator=instance_of(six.string_types))
+  entity = attr.ib()
   _time_alloc = attr.ib(validator=instance_of(TimeslotAllocation), repr=False)
   _space_alloc = attr.ib(validator=instance_of(SpaceAllocation), repr=False)
   time_slot_start = attr.ib(init=False)
@@ -70,11 +70,11 @@ def __attrs_post_init__(self):
     self.size = self._space_alloc.size
 
   @classmethod
-  def init(cls, entity_name, time_slot_start, time_slot_end, offset_start, size):
+  def init(cls, entity, time_slot_start, time_slot_end, offset_start, size):
     time_alloc = TimeslotAllocation(time_slot_start, time_slot_end)
     space_alloc = SpaceAllocation(offset_start, size)
     return cls(
-      entity_name=entity_name,
+      entity=entity,
       time_alloc=time_alloc,
       space_alloc=space_alloc
     )
@@ -94,30 +94,30 @@ def __init__(self, allocs, total_size):
         raise ValueError(
           'expecting value of {} of type {}, get {}'.format(k, TimeSpaceAllocation, type(v))
         )
-    self.plan = {alloc.entity_name: alloc for alloc in allocs}
+    self.plan = {alloc.entity: alloc for alloc in allocs}
     self.total_size = total_size
 
-  def __setitem__(self, entity_name, alloc):
+  def __setitem__(self, entity, alloc):
     if not isinstance(alloc, TimeSpaceAllocation):
       raise ValueError(
         'the value should be of type {}, get {}'.format(TimeSpaceAllocation, type(alloc))
       )
-    if entity_name in self._plan:
+    if entity in self._plan:
       logger.warning(
-        'duplicate entity_name detected: {}'.format(entity_name)
+        'duplicate entity detected: {}'.format(entity)
       )
-    self._plan[entity_name] = alloc
+    self._plan[entity] = alloc
 
-  def __getitem__(self, entity_name):
-    if entity_name not in self.plan:
-      raise KeyError('%s not found' % entity_name)
-    return self.plan[entity_name]
+  def __getitem__(self, entity):
+    if entity not in self.plan:
+      raise KeyError('%s not found' % entity)
+    return self.plan[entity]
 
-  def __contains__(self, entity_name):
-    return entity_name in self.plan
+  def __contains__(self, entity):
+    return entity in self.plan
 
-  def __delitem__(self, entity_name):
-    del self.plan[entity_name]
+  def __delitem__(self, entity):
+    del self.plan[entity]
 
   def __getattr__(self, attr_name):
     return getattr(self.plan, attr_name)

utensor_cgen/backend/graph_lower/generic_graph_lower.py

@@ -37,22 +37,22 @@ def apply(self, ugraph):
     life_span = defaultdict(lambda: [None, None])
     for op_info in ugraph.ops_info.values():
       for tensor in op_info.input_tensors:
-        ref_cnts[tensor.name] += 1
+        ref_cnts[tensor] += 1
     for time_slot, op_name in enumerate(ugraph.topo_order):
       op_info = ugraph.ops_info[op_name]
       for tensor in op_info.output_tensors:
-        life_span[tensor.name][0] = time_slot
+        life_span[tensor][0] = time_slot
       for tensor in op_info.input_tensors:
-        ref_cnts[tensor.name] -= 1
-        if ref_cnts[tensor.name] == 0:
-          life_span[tensor.name][1] = time_slot
+        ref_cnts[tensor] -= 1
+        if ref_cnts[tensor] == 0:
+          life_span[tensor][1] = time_slot
     time_alloc_plan = {}
-    for tensor_name, (start, end) in life_span.items():
+    for tensor_info, (start, end) in life_span.items():
       time_alloc = TimeslotAllocation(
         time_slot_start=start,
         time_slot_end=end
       )
-      time_alloc_plan[tensor_name] = time_alloc
+      time_alloc_plan[tensor_info] = time_alloc
     logger.info('topo ordered tensor life span analysis done')
     ugraph.attributes[self.KWARGS_NAMESCOPE] = time_alloc_plan
 
@@ -126,7 +126,7 @@ def apply(self, ugraph):
         continue
       # all output tensor should not overlap with tensors that's still alive
       for out_tensor, known_tensor in product(op_info.output_tensors, tensors_to_schedule):
-        time_alloc = time_alloc_plan[known_tensor.name]
+        time_alloc = time_alloc_plan[known_tensor]
         if time_slot in time_alloc:
           nonoverlap_map[out_tensor].add(known_tensor)
       # all output tensors should not overlap with each other
@@ -178,11 +178,11 @@ def _solve_space_alloc(self, tensors_to_schedule, nonoverlap_map):
       var_end = model.NewIntVar(0, self.max_pool_size, '{}_end'.format(tensor.name))
       size = self._compute_tensor_bytes_size(tensor)
       intv_var = model.NewIntervalVar(var_start, size, var_end, '{}_alloc'.format(tensor.name))
-      inter_vars[tensor.name] = intv_var
-      tensor_allocs[tensor.name] = _VarMemorySpan(var_start, var_end, size)
+      inter_vars[tensor] = intv_var
+      tensor_allocs[tensor] = _VarMemorySpan(var_start, var_end, size)
     for tensor in tensors_to_schedule:
-      inter_var = inter_vars[tensor.name]
-      nonoverlap_vars = [inter_vars[t.name] for t in nonoverlap_map[tensor]]
+      inter_var = inter_vars[tensor]
+      nonoverlap_vars = [inter_vars[t] for t in nonoverlap_map[tensor]]
       for other in nonoverlap_vars:
           model.AddNoOverlap([inter_var, other])
     var_mempool_size = model.NewIntVar(0, self.max_pool_size, 'mempool_size')
@@ -195,8 +195,8 @@ def _solve_space_alloc(self, tensors_to_schedule, nonoverlap_map):
     opt_mempool_size = None
     if status == cp_model.OPTIMAL:
       opt_mempool_size = solver.Value(var_mempool_size)
-      for name, alloc in tensor_allocs.items():
-        alloc_plan[name] = SpaceAllocation(
+      for entity, alloc in tensor_allocs.items():
+        alloc_plan[entity] = SpaceAllocation(
           offset_start=solver.Value(alloc.start),
           size=alloc.size,
           data_alignment=self.data_alignment,
@@ -307,7 +307,7 @@ def apply(self, ugraph):
               max_offset_end = allocate_table[in_o.name]['offsetend']
             allocs.append(
               TimeSpaceAllocation(
-                entity_name=in_o.name,
+                entity=in_o,
                 time_alloc=time_alloc,
                 space_alloc=space_alloc,
               )
@@ -328,7 +328,8 @@ def apply(self, ugraph):
               max_offset_end = allocate_table[out_o.name]['offsetend']
             allocs.append(
               TimeSpaceAllocation(
-                entity_name=out_o.name,
+                entity=out_o
+                ,
                 time_alloc=time_alloc,
                 space_alloc=space_alloc,
               )