Fixes to size desc name checks

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:

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<void**>(&{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<void**>(&{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<void *>({tmp_storage_name}), static_cast<void *>({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<void *>({tmp_storage_name}), static_cast<void *>({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"

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]
 
 

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: