From 97bc728c8aa19cb9fddfb073d6b290b95a9ad60b Mon Sep 17 00:00:00 2001 From: Yakup Budanaz Date: Tue, 3 Dec 2024 12:30:25 +0100 Subject: [PATCH] More codegen fixes --- dace/sdfg/validation.py | 9 +++++---- tests/deferred_alloc_test.py | 34 ++++++++++++++++++++++++++++------ 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/dace/sdfg/validation.py b/dace/sdfg/validation.py index 2e01b5883f..5212147c03 100644 --- a/dace/sdfg/validation.py +++ b/dace/sdfg/validation.py @@ -747,8 +747,8 @@ def validate_state(state: 'dace.sdfg.SDFGState', if isinstance(dst_node, nd.AccessNode) and isinstance(sdfg.arrays[dst_node.data], dt.Structure): name = None # Special case: if the name is the size array of the src_node, then it is ok, checked with the "size_desc_name" - src_size_access = isinstance(src_node, nd.AccessNode) and name == sdfg.arrays[src_node.data].size_desc_name - dst_size_access = isinstance(dst_node, nd.AccessNode) and name == sdfg.arrays[dst_node.data].size_desc_name + src_size_access = isinstance(src_node, nd.AccessNode) and isinstance(sdfg.arrays[src_node.data], dt.Array) and name is not None and name == sdfg.arrays[src_node.data].size_desc_name + dst_size_access = isinstance(dst_node, nd.AccessNode) and isinstance(sdfg.arrays[dst_node.data], dt.Array) and name is not None and name == sdfg.arrays[dst_node.data].size_desc_name sdict = state.scope_dict() if src_size_access and dst_size_access: raise InvalidSDFGEdgeError( @@ -766,9 +766,10 @@ def validate_state(state: 'dace.sdfg.SDFGState', ) if dst_size_access: dst_arr = sdfg.arrays[dst_node.data] - if dst_arr.storage != dace.dtypes.StorageType.GPU_Global or dst_arr.storage != dace.dtypes.StorageType.CPU_Heap: + if (dst_arr.storage != dtypes.StorageType.GPU_Global and + dst_arr.storage != dtypes.StorageType.CPU_Heap): raise InvalidSDFGEdgeError( - "Reallocating data (writing to the size connector) within a scope is not valid", + f"Reallocating data is allowed only to GPU_Global or CPU_Heap, the storage type is {dst_arr.storage}", sdfg, state_id, eid, diff --git a/tests/deferred_alloc_test.py b/tests/deferred_alloc_test.py index d2cf87168f..6459ee6105 100644 --- a/tests/deferred_alloc_test.py +++ b/tests/deferred_alloc_test.py @@ -1,6 +1,22 @@ import dace import numpy import cupy +import pytest + +@pytest.fixture(params=[dace.dtypes.StorageType.CPU_Heap, dace.dtypes.StorageType.GPU_Global]) +def storage_type(request): + return request.param + +@pytest.fixture(params=[True, False]) +def transient(request): + return request.param + +@pytest.fixture +def schedule_type(storage_type): + if storage_type == dace.dtypes.StorageType.CPU_Heap: + return dace.dtypes.ScheduleType.Sequential + elif storage_type == dace.dtypes.StorageType.GPU_Global: + return dace.dtypes.ScheduleType.GPU_Device def _get_trivial_alloc_sdfg(storage_type: dace.dtypes.StorageType, transient: bool, write_size="0:2"): sdfg = dace.sdfg.SDFG(name="deferred_alloc_test") @@ -126,26 +142,32 @@ def test_realloc_use(storage_type: dace.dtypes.StorageType, transient: bool, sch assert ( arr.get()[0] == 3.0 ) +def test_realloc_inside_map(): + pass + + +def test_all_combinations(storage_type, transient, schedule_type): + test_trivial_realloc(storage_type, transient) + test_realloc_use(storage_type, transient, schedule_type) + def test_incomplete_write_dimensions_1(): - sdfg = _get_trivial_alloc_sdfg(dace.dtypes.StorageType.CPU_Heap, True, "1:2") + sdfg = _get_trivial_alloc_sdfg(dace.dtypes.StorageType.CPU_Heap, True, "1:2") try: sdfg.validate() except Exception: return - raise AssertionError("Realloc-use with transient data and incomplete write did not fail when it was expected to.") + pytest.fail("Realloc-use with transient data and incomplete write did not fail when it was expected to.") def test_incomplete_write_dimensions_2(): - sdfg = _get_trivial_alloc_sdfg(dace.dtypes.StorageType.CPU_Heap, False, "1:2") + sdfg = _get_trivial_alloc_sdfg(dace.dtypes.StorageType.CPU_Heap, False, "1:2") try: sdfg.validate() except Exception: return - raise AssertionError("Realloc-use with non-transient data and incomplete write did not fail when it was expected to.") + pytest.fail("Realloc-use with non-transient data and incomplete write did not fail when it was expected to.") -def test_realloc_inside_map(): - pass if __name__ == "__main__": for storage_type, schedule_type in [(dace.dtypes.StorageType.CPU_Heap, dace.dtypes.ScheduleType.Sequential),