From 02a48e8227b2a28418007068bc9e439de0370ed5 Mon Sep 17 00:00:00 2001 From: Yakup Budanaz Date: Fri, 13 Dec 2024 17:04:06 +0100 Subject: [PATCH] Fixes to size desc name checks --- dace/codegen/targets/cpu.py | 9 ++++--- dace/codegen/targets/cuda.py | 27 ++++++++++--------- dace/sdfg/sdfg.py | 16 ++++++----- .../passes/array_elimination.py | 2 +- 4 files changed, 29 insertions(+), 25 deletions(-) diff --git a/dace/codegen/targets/cpu.py b/dace/codegen/targets/cpu.py index af7bb0502b..2d1c33df44 100644 --- a/dace/codegen/targets/cpu.py +++ b/dace/codegen/targets/cpu.py @@ -341,7 +341,7 @@ def declare_array(self, self._dispatcher.declared_arrays.add(name, DefinedType.Pointer, ctypedef) # Size desc is defined only for transient arrays - if nodedesc.transient and nodedesc.storage == dtypes.StorageType.CPU_Heap: + if nodedesc.transient and nodedesc.storage == dtypes.StorageType.CPU_Heap and type(nodedesc) == data.Array: size_desc_name = sdfg.arrays[name].size_desc_name if size_desc_name is not None: size_desc = sdfg.arrays[size_desc_name] @@ -698,13 +698,14 @@ def reallocate( data_name = dst_node.data new_size_array_name = src_node.data - data = sdfg.arrays[data_name] - size_array_name = data.size_desc_name + desc = sdfg.arrays[data_name] + assert type(data) == dt.Array + size_array_name = desc.size_desc_name dtype = sdfg.arrays[data_name].dtype size_assignment_strs, new_size_strs, _ = cpp._get_realloc_dimensions( - size_array_name, new_size_array_name, data.shape + size_array_name, new_size_array_name, desc.shape ) for size_assignment in size_assignment_strs: diff --git a/dace/codegen/targets/cuda.py b/dace/codegen/targets/cuda.py index 6d77daa219..abebe1119d 100644 --- a/dace/codegen/targets/cuda.py +++ b/dace/codegen/targets/cuda.py @@ -1582,7 +1582,7 @@ def generate_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: StateSub if aname in sdfg.arrays: arr = sdfg.arrays[aname] - if arr.transient and arr.storage == dtypes.StorageType.GPU_Global and arr.size_desc_name is not None: + if arr.transient and arr.storage == dtypes.StorageType.GPU_Global and type(arr) == dt.Array and arr.size_desc_name is not None: size_arr_name = data_desc.size_desc_name if size_arr_name is not None: size_arr = sdfg.arrays[size_arr_name] @@ -2795,13 +2795,14 @@ def reallocate( data_name = dst_node.data new_size_array_name = src_node.data - data = sdfg.arrays[data_name] - size_array_name = data.size_desc_name + desc = sdfg.arrays[data_name] + assert type(desc) == dt.Array + size_array_name = desc.size_desc_name dtype = sdfg.arrays[data_name].dtype size_assignment_strs, new_size_strs, old_size_strs = cpp._get_realloc_dimensions( - size_array_name, new_size_array_name, data.shape + size_array_name, new_size_array_name, desc.shape ) @@ -2814,23 +2815,23 @@ def reallocate( tmp_storage_name = "__tmp_realloc_move_storage" callsite_stream.write(f"if ({dst_node.data} == nullptr) {{", cfg, state_id, dst_node.guid) - if data.storage == dtypes.StorageType.GPU_Global: - self._alloc_gpu_global(dst_node, data, callsite_stream, data_name, new_size_str) + if desc.storage == dtypes.StorageType.GPU_Global: + self._alloc_gpu_global(dst_node, desc, callsite_stream, data_name, new_size_str) else: - assert data.storage == dtypes.StorageType.CPU_Pinned + assert desc.storage == dtypes.StorageType.CPU_Pinned callsite_stream.write(f"DACE_GPU_CHECK({self.backend}MallocHost(reinterpret_cast(&{data_name}), {new_size_str}));", cfg, state_id, dst_node.guid) callsite_stream.write("} else {\n", cfg, state_id, dst_node.guid) callsite_stream.write(f"{dtype}* {tmp_storage_name};") - if data.storage == dtypes.StorageType.GPU_Global: - self._alloc_gpu_global(None, data, callsite_stream, tmp_storage_name, new_size_str) + if desc.storage == dtypes.StorageType.GPU_Global: + self._alloc_gpu_global(None, desc, callsite_stream, tmp_storage_name, new_size_str) else: - assert data.storage == dtypes.StorageType.CPU_Pinned + assert desc.storage == dtypes.StorageType.CPU_Pinned callsite_stream.write(f"DACE_GPU_CHECK({self.backend}MallocHost(reinterpret_cast(&{tmp_storage_name}), {new_size_str}));", cfg, state_id, dst_node.guid) s = "" copy_size_str = f"Min({old_size_str}, {new_size_str})" - if data.storage == dtypes.StorageType.GPU_Global: - if not data.pool: # If pooled, will be freed somewhere else + if desc.storage == dtypes.StorageType.GPU_Global: + if not desc.pool: # If pooled, will be freed somewhere else s += f"DACE_GPU_CHECK({self.backend}Memcpy(static_cast({tmp_storage_name}), static_cast({data_name}), {copy_size_str}, cudaMemcpyDeviceToDevice));\n" s += f"DACE_GPU_CHECK({self.backend}Free({data_name}));\n" s += f"{data_name} = {tmp_storage_name};\n" @@ -2845,7 +2846,7 @@ def reallocate( self._emit_sync(callsite_stream) callsite_stream.write(f"{data_name} = {tmp_storage_name};\n") s = "" - elif data.storage == dtypes.StorageType.CPU_Pinned: + elif desc.storage == dtypes.StorageType.CPU_Pinned: s += f"DACE_GPU_CHECK({self.backend}Memcpy(static_cast({tmp_storage_name}), static_cast({data_name}), {copy_size_str}, cudaMemcpyHostToHost));\n" s += f"DACE_GPU_CHECK({self.backend}FreeHost({data_name}));\n" s += f"{data_name} = {tmp_storage_name};\n" diff --git a/dace/sdfg/sdfg.py b/dace/sdfg/sdfg.py index 58b9dfd854..e1d4962a02 100644 --- a/dace/sdfg/sdfg.py +++ b/dace/sdfg/sdfg.py @@ -774,11 +774,10 @@ def replace_dict(self, size_ararys_to_rm = set() for arr_name, size_desc_name in size_desc_map.items(): arr = self.arrays[arr_name] if arr_name in self.arrays else None - if arr is not None: + if arr is not None and type(arr) == dt.Array: size_desc_name_before = arr.size_desc_name # If we change the name of an array, then we need to change its size array accordingly - if (arr.transient and type(arr) == dt.Array and size_desc_name_before is not None - and size_desc_name is not None): + if (arr.transient and size_desc_name_before is not None): arr.size_desc_name = size_desc_name assert (arr.size_desc_name == size_desc_name) self.arrays[size_desc_name] = self.arrays.pop(size_desc_name_before) @@ -1199,10 +1198,13 @@ 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 - # If unused it might have been removed by optimization - if size_desc_name is not None and size_desc_name in self._arrays: - del self._arrays[size_desc_name] + # Check for size desc + if type(self._arrays[name]) == dt.Array: + size_desc_name = self._arrays[name].size_desc_name + # If unused it might have been removed by optimization + if size_desc_name is not None and size_desc_name in self._arrays: + del self._arrays[size_desc_name] + del self._arrays[name] diff --git a/dace/transformation/passes/array_elimination.py b/dace/transformation/passes/array_elimination.py index f7cf0ac64d..6a28877aa4 100644 --- a/dace/transformation/passes/array_elimination.py +++ b/dace/transformation/passes/array_elimination.py @@ -85,7 +85,7 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S # If node is completely removed from graph, erase data descriptor array_items = list(sdfg.arrays.items()) - size_descriptors = set([v.size_desc_name for v in sdfg.arrays.values() if v.size_desc_name is not None]) + size_descriptors = set([v.size_desc_name for v in sdfg.arrays.values() if type(v) == data.Array and v.size_desc_name is not None]) for aname, desc in array_items: # Remove size descriptors only if the original array is removed if aname in size_descriptors: