Deferred Allocation

dace/codegen/dispatcher.py

@@ -20,7 +20,7 @@
 @registry.extensible_enum
 class DefinedType(aenum.AutoNumberEnum):
     """ Data types for `DefinedMemlets`.
-    
+
         :see: DefinedMemlets
     """
     Pointer = ()  # Pointer
@@ -159,6 +159,8 @@ class TargetDispatcher(object):
     _state_dispatchers: List[Tuple[Callable, target.TargetCodeGenerator]]
     _generic_state_dispatcher: Optional[target.TargetCodeGenerator]
 
+    _generic_reallocate_dispatchers: Dict[dtypes.StorageType, target.TargetCodeGenerator]
+
     _declared_arrays: DefinedMemlets
     _defined_vars: DefinedMemlets
 
@@ -181,6 +183,7 @@ def __init__(self, framecode):
         self._node_dispatchers = []
         self._generic_node_dispatcher = None
         self._state_dispatchers = []
+        self._generic_reallocate_dispatchers = {}
         self._generic_state_dispatcher = None
 
         self._declared_arrays = DefinedMemlets()
@@ -189,7 +192,7 @@ def __init__(self, framecode):
     @property
     def declared_arrays(self) -> DefinedMemlets:
         """ Returns a list of declared variables.
-        
+
             This is used for variables that must have their declaration and
             allocation separate. It includes all such variables that have been
             declared by the dispatcher.
@@ -199,7 +202,7 @@ def declared_arrays(self) -> DefinedMemlets:
     @property
     def defined_vars(self) -> DefinedMemlets:
         """ Returns a list of defined variables.
-        
+
             This includes all variables defined by the dispatcher.
         """
         return self._defined_vars
@@ -354,6 +357,15 @@ def register_copy_dispatcher(self, src_storage: dtypes.StorageType, dst_storage:
 
         self._copy_dispatchers[dispatcher].append((predicate, func))
 
+    def register_reallocate_dispatcher(self, node_storage: dtypes.StorageType,
+                                      func: target.TargetCodeGenerator,
+                                      predicate: Optional[Callable] = None) -> None:
+
+        if not isinstance(node_storage, dtypes.StorageType): raise TypeError(node_storage, dtypes.StorageType, isinstance(node_storage, dtypes.StorageType))
+        dispatcher = node_storage
+        self._generic_reallocate_dispatchers[dispatcher] = func
+        return
+
     def get_state_dispatcher(self, sdfg: SDFG, state: SDFGState) -> target.TargetCodeGenerator:
         # Check if the state satisfies any predicates that delegate to a
         # specific code generator
@@ -594,6 +606,14 @@ def get_copy_dispatcher(self, src_node: Union[nodes.CodeNode, nodes.AccessNode],
 
         return target
 
+    def get_reallocate_dispatcher(self, node: Union[nodes.CodeNode, nodes.AccessNode],
+                            edge: MultiConnectorEdge[Memlet],
+                            sdfg: SDFG, state: SDFGState) -> Optional[target.TargetCodeGenerator]:
+        node_storage = sdfg.arrays[node.data].storage
+        target = self._generic_reallocate_dispatchers[node_storage]
+        return target
+
+
     def dispatch_copy(self, src_node: nodes.Node, dst_node: nodes.Node, edge: MultiConnectorEdge[Memlet], sdfg: SDFG,
                       cfg: ControlFlowRegion, dfg: StateSubgraphView, state_id: int, function_stream: CodeIOStream,
                       output_stream: CodeIOStream) -> None:
@@ -609,6 +629,20 @@ def dispatch_copy(self, src_node: nodes.Node, dst_node: nodes.Node, edge: MultiC
         self._used_targets.add(target)
         target.copy_memory(sdfg, cfg, dfg, state_id, src_node, dst_node, edge, function_stream, output_stream)
 
+    def dispatch_reallocate(self, src_node: nodes.Node, node: nodes.Node, edge: MultiConnectorEdge[Memlet], sdfg: SDFG,
+                      cfg: ControlFlowRegion, dfg: StateSubgraphView, state_id: int, function_stream: CodeIOStream,
+                      output_stream: CodeIOStream) -> None:
+        state = cfg.state(state_id)
+        target = self.get_reallocate_dispatcher(node, edge, sdfg, state)
+        assert target is not None
+        if target is None:
+            return
+
+        # Dispatch reallocate
+        self._used_targets.add(target)
+        target.reallocate(sdfg, cfg, dfg, state_id, src_node, node, edge, function_stream, output_stream)
+
+
     # Dispatches definition code for a memlet that is outgoing from a tasklet
     def dispatch_output_definition(self, src_node: nodes.Node, dst_node: nodes.Node, edge, sdfg: SDFG,
                                    cfg: ControlFlowRegion, dfg: StateSubgraphView, state_id: int,

dace/codegen/targets/cpp.py

@@ -9,6 +9,7 @@
 import itertools
 import math
 import numbers
+import re
 import sys
 import warnings
 
@@ -231,6 +232,14 @@ def memlet_copy_to_absolute_strides(dispatcher: 'TargetDispatcher',
         elif memlet.data == dst_node.data:
             copy_shape, src_strides = reshape_strides(dst_subset, dst_strides, src_strides, copy_shape)
 
+    def replace_dace_defer_dim(string, arrname):
+        pattern = r"__dace_defer_dim(\d+)"
+        return re.sub(pattern, r"A_size[\1]", string)
+
+    # TODO: do this better?
+    dst_expr = replace_dace_defer_dim(dst_expr, dst_node.data) if dst_expr is not None else None
+    src_expr = replace_dace_defer_dim(src_expr, src_node.data) if src_expr is not None else None
+
     return copy_shape, src_strides, dst_strides, src_expr, dst_expr
 
 
@@ -539,7 +548,8 @@ def ndcopy_to_strided_copy(
         return None
 
 
-def cpp_offset_expr(d: data.Data, subset_in: subsets.Subset, offset=None, packed_veclen=1, indices=None):
+def cpp_offset_expr(d: data.Data, subset_in: subsets.Subset, offset=None, packed_veclen=1, indices=None,
+                    deferred_size_names=None):
     """ Creates a C++ expression that can be added to a pointer in order
         to offset it to the beginning of the given subset and offset.
 
@@ -569,7 +579,16 @@ def cpp_offset_expr(d: data.Data, subset_in: subsets.Subset, offset=None, packed
     if packed_veclen > 1:
         index /= packed_veclen
 
-    return sym2cpp(index)
+    if not (deferred_size_names is None):
+        access_str_with_deferred_vars = sym2cpp(index)
+        def replace_pattern(match):
+            number = match.group(1)
+            return deferred_size_names[int(number)]
+        pattern = r'__dace_defer_dim(\d+)'
+        access_str = re.sub(pattern, replace_pattern, access_str_with_deferred_vars)
+        return access_str
+    else:
+        return sym2cpp(index)
 
 
 def cpp_array_expr(sdfg,
@@ -581,13 +600,15 @@ def cpp_array_expr(sdfg,
                    use_other_subset=False,
                    indices=None,
                    referenced_array=None,
-                   codegen=None):
+                   codegen=None,
+                   deferred_size_names=None):
     """ Converts an Indices/Range object to a C++ array access string. """
     subset = memlet.subset if not use_other_subset else memlet.other_subset
     s = subset if relative_offset else subsets.Indices(offset)
     o = offset if relative_offset else None
-    desc = (sdfg.arrays[memlet.data] if referenced_array is None else referenced_array)
-    offset_cppstr = cpp_offset_expr(desc, s, o, packed_veclen, indices=indices)
+    desc : dace.Data = (sdfg.arrays[memlet.data] if referenced_array is None else referenced_array)
+    desc_name = memlet.data
+    offset_cppstr = cpp_offset_expr(desc, s, o, packed_veclen, indices=indices, deferred_size_names=deferred_size_names)
 
     # NOTE: Are there any cases where a mix of '.' and '->' is needed when traversing nested structs?
     # TODO: Study this when changing Structures to be (optionally?) non-pointers.
@@ -763,7 +784,7 @@ def is_write_conflicted_with_reason(dfg, edge, datanode=None, sdfg_schedule=None
     Detects whether a write-conflict-resolving edge can be emitted without
     using atomics or critical sections, returning the node or SDFG that caused
     the decision.
-    
+
     :return: None if the conflict is nonatomic, otherwise returns the scope entry
              node or SDFG that caused the decision to be made.
     """