Skip to content

Commit

Permalink
Split size and array storage
Browse files Browse the repository at this point in the history
  • Loading branch information
@ThrudPrimrose
ThrudPrimrose committed Dec 3, 2024
1 parent 97bc728 commit 08cb50c
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 19 deletions.
5 changes: 2 additions & 3 deletions dace/codegen/targets/cpu.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -344,7 +344,7 @@ def declare_array(self,
if nodedesc.transient and nodedesc.storage == dtypes.StorageType.CPU_Heap:
size_desc_name = sdfg.arrays[name].size_desc_name
if size_desc_name is not None:
size_desc = sdfg.arrays[size_desc_name]
size_desc = sdfg.size_arrays[size_desc_name]
size_ctypedef = dtypes.pointer(size_desc.dtype).ctype
self._dispatcher.declared_arrays.add(size_desc_name, DefinedType.Pointer, size_ctypedef)
return
Expand Down Expand Up @@ -514,7 +514,7 @@ def allocate_array(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSubgraphV
# Initialize size array
size_str = ",".join(["0" if cpp.sym2cpp(dim).startswith("__dace_defer") else cpp.sym2cpp(dim) for dim in nodedesc.shape])
size_desc_name = nodedesc.size_desc_name
size_nodedesc = sdfg.arrays[size_desc_name]
size_nodedesc = sdfg.size_arrays[size_desc_name]
declaration_stream.write(f'{size_nodedesc.dtype.ctype} {size_desc_name}[{size_nodedesc.shape[0]}]{{{size_str}}};\n', cfg, state_id, node)
if deferred_allocation:
allocation_stream.write(
Expand Down Expand Up @@ -708,7 +708,6 @@ def reallocate(
data = sdfg.arrays[data_name]
size_array_name = data.size_desc_name

new_size_array = sdfg.arrays[new_size_array_name]
dtype = sdfg.arrays[data_name].dtype

# Only consider the offsets with __dace_defer in original dim
Expand Down
5 changes: 2 additions & 3 deletions dace/codegen/targets/cuda.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -613,7 +613,7 @@ def allocate_array(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg: StateSubgraphV
if nodedesc.transient:
size_desc_name = nodedesc.size_desc_name
if size_desc_name is not None:
size_nodedesc = sdfg.arrays[size_desc_name]
size_nodedesc = sdfg.size_arrays[size_desc_name]
result_decl.write(f'{size_nodedesc.dtype.ctype} {size_desc_name}[{size_nodedesc.shape[0]}]{{{size_str}}};\n')
self._dispatcher.defined_vars.add(size_desc_name, DefinedType.Pointer, size_nodedesc.dtype.ctype)
self._dispatcher.defined_vars.add(dataname, DefinedType.Pointer, ctypedef)
Expand Down Expand Up @@ -1586,7 +1586,7 @@ def generate_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: StateSub
if aname in sdfg.arrays:
size_arr_name = data_desc.size_desc_name
if size_arr_name is not None:
size_arr = sdfg.arrays[data_desc.size_desc_name]
size_arr = sdfg.size_arrays[data_desc.size_desc_name]
host_size_args[size_arr_name] = size_arr

kernel_args_typed = [('const ' if k in const_params else '') + v.as_arg(name=k)
Expand Down Expand Up @@ -2796,7 +2796,6 @@ def reallocate(
data = sdfg.arrays[data_name]
size_array_name = data.size_desc_name

new_size_array = sdfg.arrays[new_size_array_name]
dtype = sdfg.arrays[data_name].dtype

# Only consider the offsets with __dace_defer in original dim
Expand Down
71 changes: 58 additions & 13 deletions dace/sdfg/sdfg.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,10 @@ class SDFG(ControlFlowRegion):
desc="Data descriptors for this SDFG",
to_json=_arrays_to_json,
from_json=_nested_arrays_from_json)
_size_arrays = Property(dtype=NestedDict,
desc="Data size descriptors for this SDFG",
to_json=_arrays_to_json,
from_json=_nested_arrays_from_json)
symbols = DictProperty(str, dtypes.typeclass, desc="Global symbols for this SDFG")

instrument = EnumProperty(dtype=dtypes.InstrumentationType,
Expand Down Expand Up @@ -496,6 +500,7 @@ def __init__(self,
self._parent_sdfg = None
self._parent_nsdfg_node = None
self._arrays = NestedDict() # type: Dict[str, dt.Array]
self._size_arrays = NestedDict()
self.arg_names = []
self._labels: Set[str] = set()
self.global_code = {'frame': CodeBlock("", dtypes.Language.CPP)}
Expand Down Expand Up @@ -683,6 +688,10 @@ def arrays(self):
"""
return self._arrays

@property
def size_arrays(self):
return self._size_arrays

@property
def process_grids(self):
""" Returns a dictionary of process-grid descriptors (`ProcessGrid` objects) used in this SDFG. """
Expand Down Expand Up @@ -746,6 +755,7 @@ def replace_dict(self,
for name, new_name in repldict_filtered.items():
if validate_name(new_name):
_replace_dict_keys(self._arrays, name, new_name)
_replace_dict_keys(self._size_arrays, name + "_size", new_name + "_size")
_replace_dict_keys(self.symbols, name, new_name)
_replace_dict_keys(self.constants_prop, name, new_name)
_replace_dict_keys(self.callback_mapping, name, new_name)
Expand Down Expand Up @@ -1151,7 +1161,10 @@ def remove_data(self, name, validate=True):
f"{name}: it is accessed by node "
f"{node} in state {state}.")

size_desc_name = self._arrays[name].size_desc_name
del self._arrays[name]
if size_desc_name is not None:
del self._size_arrays[size_desc_name]

def reset_sdfg_list(self):
"""
Expand Down Expand Up @@ -1676,13 +1689,15 @@ def _find_new_name(self, name: str):
""" Tries to find a new name by adding an underscore and a number. """

names = (self._arrays.keys() | self.constants_prop.keys() | self._pgrids.keys() | self._subarrays.keys()
| self._rdistrarrays.keys() | self.symbols.keys())
| self._rdistrarrays.keys() | self.symbols.keys() | self._size_arrays.keys())
return dt.find_new_name(name, names)

def is_name_used(self, name: str) -> bool:
""" Checks if `name` is already used inside the SDFG."""
if name in self._arrays:
return True
if name in self._size_arrays:
return True
if name in self.symbols:
return True
if name in self.constants_prop:
Expand Down Expand Up @@ -1787,7 +1802,7 @@ def add_array(self,
array_name = self.add_datadesc(name, desc, find_new_name=find_new_name)
if transient:
size_desc_name = f"{array_name}_size"
self.add_datadesc(size_desc_name, size_desc, find_new_name=False)
self.add_size_datadesc(size_desc_name, size_desc)
# In case find_new_name and a new name is returned
# we need to update the size descriptor name of the array
desc.size_desc_name = size_desc_name
Expand Down Expand Up @@ -2038,6 +2053,16 @@ def add_temp_transient_like(self, desc: Union[dt.Array, dt.Scalar], dtype=None,
newdesc.debuginfo = debuginfo
return self.add_datadesc(self.temp_data_name(), newdesc), newdesc

@staticmethod
def _add_symbols(sdfg, desc: dt.Data):
if isinstance(desc, dt.Structure):
for v in desc.members.values():
if isinstance(v, dt.Data):
SDFG._add_symbols(sdfg, v)
for sym in desc.free_symbols:
if sym.name not in sdfg.symbols:
sdfg.add_symbol(sym.name, sym.dtype)

def add_datadesc(self, name: str, datadesc: dt.Data, find_new_name=False) -> str:
""" Adds an existing data descriptor to the SDFG array store.
Expand Down Expand Up @@ -2067,7 +2092,7 @@ def add_datadesc(self, name: str, datadesc: dt.Data, find_new_name=False) -> str
else:
# We do not check for data constant, because there is a link between the constants and
# the data descriptors.
if name in self.arrays:
if name in self.arrays or name in self.size_arrays:
raise FileExistsError(f'Data descriptor "{name}" already exists in SDFG')
if name in self.symbols:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a symbol.')
Expand All @@ -2078,18 +2103,38 @@ def add_datadesc(self, name: str, datadesc: dt.Data, find_new_name=False) -> str
if name in self._pgrids:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a ProcessGrid.')

def _add_symbols(sdfg: SDFG, desc: dt.Data):
if isinstance(desc, dt.Structure):
for v in desc.members.values():
if isinstance(v, dt.Data):
_add_symbols(sdfg, v)
for sym in desc.free_symbols:
if sym.name not in sdfg.symbols:
sdfg.add_symbol(sym.name, sym.dtype)

# Add the data descriptor to the SDFG and all symbols that are not yet known.
self._arrays[name] = datadesc
_add_symbols(self, datadesc)
SDFG._add_symbols(self, datadesc)

return name

def add_size_datadesc(self, name: str, datadesc: dt.Data) -> str:
""" Adds an existing data descriptor to the SDFG array store.
:param name: Name to use.
:param datadesc: Data descriptor to add.
:param find_new_name: If True and data descriptor with this name
exists, finds a new name to add.
:return: Name of the new data descriptor
"""
if not isinstance(name, str):
raise TypeError("Data descriptor name must be a string. Got %s" % type(name).__name__)

if name in self.arrays or name in self.size_arrays:
raise FileExistsError(f'Data descriptor "{name}" already exists in SDFG')
if name in self.symbols:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a symbol.')
if name in self._subarrays:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a subarray.')
if name in self._rdistrarrays:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a RedistrArray.')
if name in self._pgrids:
raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a ProcessGrid.')

# Add the data descriptor to the SDFG and all symbols that are not yet known.
self._size_arrays[name] = datadesc
SDFG._add_symbols(self, datadesc)

return name

Expand Down

0 comments on commit 08cb50c

Please sign in to comment.