From b2e46a19d0b4392d086903a7ab3e327f616306db Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 13 Feb 2024 14:22:51 -0500 Subject: [PATCH 01/25] Allow frontend to accept K write offset for FORWARD/BACKWARD --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index 2df8c106ce..8a70b8e73e 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1417,9 +1417,14 @@ def visit_Assign(self, node: ast.Assign) -> list: for t in node.targets[0].elts if isinstance(node.targets[0], ast.Tuple) else node.targets: name, spatial_offset, data_index = self._parse_assign_target(t) if spatial_offset: - if any(offset != 0 for offset in spatial_offset): + if spatial_offset[0] != 0 or spatial_offset[1] != 0: raise GTScriptSyntaxError( - message="Assignment to non-zero offsets is not supported.", + message="Assignment to non-zero offsets is not supported in IJ.", + loc=nodes.Location.from_ast_node(t), + ) + if spatial_offset[2] != 0 and self.iteration_order == nodes.IterationOrder.PARALLEL: + raise GTScriptSyntaxError( + message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", loc=nodes.Location.from_ast_node(t), ) From 6b9d127d5c485fc3bce1e376ec42dbb8095c7a9a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 13 Feb 2024 14:23:03 -0500 Subject: [PATCH 02/25] Unit test on frontend & code generation --- .../test_code_generation.py | 53 +++++++++++++++++++ .../frontend_tests/test_gtscript_frontend.py | 34 +++++++++++- 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index ad5f6de41c..0474857002 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -586,3 +586,56 @@ def test(out: Field[np.float64], inp: Field[np.float64]): backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 2), dtype=np.float64 ) test(out, inp) + + +@pytest.mark.parametrize("backend", ["numpy"]) +def test_K_offset_write(backend): + @gtscript.stencil(backend=backend) + def simple(A: Field[np.float64], B: Field[np.float64]): + with computation(FORWARD), interval(...): + B[0, 0, 1] = A + + A = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + A[:, :, :] = 42.0 + B = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + simple(A, B) + assert (B[:, :, 0] == 0).all() + assert (B[:, :, 1:3] == 42.0).all() + + @gtscript.stencil(backend=backend) + def forward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): + with computation(FORWARD), interval(1, 3): + A[0, 0, -1] = scalar + B[0, 0, 1] = A + + A = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + A[:, :, :] = 42.0 + B = gt_storage.empty( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + forward(A, B, 2.0) + # B will have non-updated values of A + assert (B[:, :, 1:3] == 42.0).all() + + @gtscript.stencil(backend=backend) + def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): + with computation(BACKWARD), interval(1, 3): + A[0, 0, -1] = scalar + B[0, 0, 1] = A + + A = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + A[:, :, :] = 42.0 + B = gt_storage.empty( + backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + ) + backward(A, B, 2.0) + # B will have updated values of A + assert (B[:, :, 1:3] == 2.0).all() diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py index 00b91eb2cf..a92368d940 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py @@ -28,6 +28,7 @@ IJ, IJK, PARALLEL, + FORWARD, Field, I, J, @@ -1382,7 +1383,7 @@ def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float with pytest.raises( gt_frontend.GTScriptSyntaxError, - match="Assignment to non-zero offsets is not supported.", + match="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", ): parse_definition( func, @@ -1421,7 +1422,7 @@ def definition_func( with pytest.raises( gt_frontend.GTScriptSyntaxError, - match="Assignment to non-zero offsets is not supported.", + match="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", ): parse_definition( definition_func, name=inspect.stack()[0][3], module=self.__class__.__name__ @@ -1455,6 +1456,35 @@ def func(in_field: gtscript.Field[np.float_]): parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) + def test_K_offset_write(self): + def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): + with computation(FORWARD), interval(...): + out[0, 0, 1] = inp + + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) + + with pytest.raises( + gt_frontend.GTScriptSyntaxError, + match=r"(.*?)Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.(.*)", + ): + + def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): + with computation(PARALLEL), interval(...): + out[0, 0, 1] = inp + + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) + + with pytest.raises( + gt_frontend.GTScriptSyntaxError, + match=r"(.*?)Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.(.*)", + ): + + def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): + with computation(...), interval(...): + out[0, 0, 1] = inp + + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) + class TestNestedWithSyntax: def test_nested_with(self): From 1980fe8e17aa4e8448463e59f6da8dec6d03fa70 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 13 Feb 2024 14:37:57 -0500 Subject: [PATCH 03/25] lint --- .../unit_tests/frontend_tests/test_gtscript_frontend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py index a92368d940..9a96644a92 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py @@ -25,10 +25,10 @@ from gt4py.cartesian.frontend import gtscript_frontend as gt_frontend, nodes from gt4py.cartesian.gtscript import ( __INLINED, + FORWARD, IJ, IJK, PARALLEL, - FORWARD, Field, I, J, From 6e8c69f3182f47e393f023bfaaed7d86ce224588 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 13 Feb 2024 15:01:25 -0500 Subject: [PATCH 04/25] Remove debug code --- .../frontend_tests/test_gtscript_frontend.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py index 9a96644a92..f281a0a4a2 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py @@ -1474,17 +1474,6 @@ def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) - with pytest.raises( - gt_frontend.GTScriptSyntaxError, - match=r"(.*?)Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.(.*)", - ): - - def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): - with computation(...), interval(...): - out[0, 0, 1] = inp - - parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) - class TestNestedWithSyntax: def test_nested_with(self): From d07a2de6fee9b4fca7085c56e3941544f4f18906 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 20 Feb 2024 12:27:51 -0500 Subject: [PATCH 05/25] Add conditional example simplified from existing column physics --- .../test_code_generation.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 0474857002..290cdddde6 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -639,3 +639,26 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): backward(A, B, 2.0) # B will have updated values of A assert (B[:, :, 1:3] == 2.0).all() + + # Conditional with secoundary K index on a while loop, + # simplified from in-situ code of NOAA's SHiELD microphysics + @gtscript.stencil(backend=backend) + def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): + with computation(BACKWARD), interval(1, 3): + if A > 0 and B > 0: + A[0, 0, -1] = scalar + B[0, 0, 1] = A + lev = 1 + while A >= 0 and B >= 0: + A[0, 0, lev] = -1 + B = -1 + lev = lev + 1 + + A = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64 + ) + A[:, :, :] = 42.0 + B = gt_storage.ones(backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64) + column_physics_conditional(A, B, 2.0) + assert (A[0, 0, :] == [2, 2, -1, -1]).all() + assert (B[0, 0, :] == [1, -1, 2, 42]).all() From e90ec9434cef958c38ca3970ed4e62ca7e757291 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 22 Feb 2024 09:24:49 -0500 Subject: [PATCH 06/25] Restore ALL_BACKENDS, straighten utest. --- .../test_code_generation.py | 58 ++++++++++++------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 290cdddde6..bb235145d1 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -588,63 +588,77 @@ def test(out: Field[np.float64], inp: Field[np.float64]): test(out, inp) -@pytest.mark.parametrize("backend", ["numpy"]) +@pytest.mark.parametrize("backend", ALL_BACKENDS) def test_K_offset_write(backend): + array_shape = (1, 1, 4) + K_values = np.arange(start=40, stop=44) + + # Simple case of writing ot an offset. + # A is untouched + # B is written in K+1 and should have K_values, except for the first element (FORWARD) @gtscript.stencil(backend=backend) def simple(A: Field[np.float64], B: Field[np.float64]): with computation(FORWARD), interval(...): B[0, 0, 1] = A A = gt_storage.zeros( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) - A[:, :, :] = 42.0 + A[:, :, :] = K_values B = gt_storage.zeros( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) simple(A, B) assert (B[:, :, 0] == 0).all() - assert (B[:, :, 1:3] == 42.0).all() + assert (B[:, :, 1:3] == K_values[0:2]).all() + # Order of operations: FORWARD with negative offset + # means while A is update B will have non-updated values of A + # Because of the interval, value of B[0] is 0 @gtscript.stencil(backend=backend) def forward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): - with computation(FORWARD), interval(1, 3): + with computation(FORWARD), interval(1, None): A[0, 0, -1] = scalar - B[0, 0, 1] = A + B[0, 0, 0] = A A = gt_storage.zeros( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) - A[:, :, :] = 42.0 - B = gt_storage.empty( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + A[:, :, :] = K_values + B = gt_storage.zeros( + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) forward(A, B, 2.0) - # B will have non-updated values of A - assert (B[:, :, 1:3] == 42.0).all() + assert (A[:, :, 1:3] == 2.0).all() + assert (A[:, :, 3] == K_values[3]).all() + assert (B[:, :, 0] == 0).all() + assert (B[:, :, 1:] == K_values[1:]).all() + # Order of operations: BACKWARD with negative offset + # means A is update B will get the updated values of A @gtscript.stencil(backend=backend) def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): - with computation(BACKWARD), interval(1, 3): + with computation(BACKWARD), interval(1, None): A[0, 0, -1] = scalar - B[0, 0, 1] = A + B[0, 0, 0] = A A = gt_storage.zeros( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) - A[:, :, :] = 42.0 + A[:, :, :] = K_values B = gt_storage.empty( - backend=backend, aligned_index=(0, 0, 0), shape=(2, 2, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) backward(A, B, 2.0) - # B will have updated values of A - assert (B[:, :, 1:3] == 2.0).all() + assert (A[:, :, 1:3] == 2.0).all() + assert (A[:, :, 3] == K_values[3]).all() + assert (B[:, :, :] == A[:, :, :]).all() # Conditional with secoundary K index on a while loop, # simplified from in-situ code of NOAA's SHiELD microphysics @gtscript.stencil(backend=backend) def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): - with computation(BACKWARD), interval(1, 3): + with computation(BACKWARD), interval(1, None): if A > 0 and B > 0: A[0, 0, -1] = scalar B[0, 0, 1] = A @@ -657,7 +671,7 @@ def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scala A = gt_storage.zeros( backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64 ) - A[:, :, :] = 42.0 + A[:, :, :] = K_values B = gt_storage.ones(backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64) column_physics_conditional(A, B, 2.0) assert (A[0, 0, :] == [2, 2, -1, -1]).all() From 5fff068770d38a568eb2cac6adc7205c405d0700 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 22 Feb 2024 17:07:35 -0500 Subject: [PATCH 07/25] Verbose DaCeIR coding error when a tasklet variable cannot find the input/output memlet link --- src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py b/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py index 779fca0c8d..3383160022 100644 --- a/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py +++ b/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py @@ -85,7 +85,12 @@ def visit_IndexAccess( # if this node is not a target, it will still use the symbol of the write memlet if the # field was previously written in the same memlet. memlets = kwargs["read_memlets"] + kwargs["write_memlets"] - memlet = next(mem for mem in memlets if mem.connector == node.name) + try: + memlet = next(mem for mem in memlets if mem.connector == node.name) + except StopIteration: + raise ValueError( + "Memlet connector and tasklet variable mismatch, DaCe IR error." + ) from None index_strs = [] if node.offset is not None: From 2147811df41d6e74e6ec61abc7456262f2e5e86c Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 22 Feb 2024 17:07:56 -0500 Subject: [PATCH 08/25] Allow re-use of offset K write ins --- src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py index 9a214441ad..c73ce412ad 100644 --- a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py +++ b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py @@ -354,7 +354,11 @@ def visit_FieldAccess( ) else: res = dcir.ScalarAccess(name=name, dtype=node.dtype) - if is_target: + # Because we allow writing in K, we allow targets who have been + # writing in K to be omitted so the original connector can be re-used. + # Previous guardrails in `gtscript` restrict the writes + # to non-parallel K-loop so we don't have issues. + if is_target and node.offset.to_dict()["k"] == 0: targets.add(node.name) return res From 3c6b466765d092299349b8e622325f7f175dd373 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 23 Feb 2024 15:03:35 -0500 Subject: [PATCH 09/25] Split utest to acknowledge dace:cpu issue --- .../test_code_generation.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index bb235145d1..04fd3d5066 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -629,7 +629,7 @@ def forward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) forward(A, B, 2.0) - assert (A[:, :, 1:3] == 2.0).all() + assert (A[:, :, :3] == 2.0).all() assert (A[:, :, 3] == K_values[3]).all() assert (B[:, :, 0] == 0).all() assert (B[:, :, 1:] == K_values[1:]).all() @@ -650,12 +650,22 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) backward(A, B, 2.0) - assert (A[:, :, 1:3] == 2.0).all() + assert (A[:, :, :3] == 2.0).all() assert (A[:, :, 3] == K_values[3]).all() assert (B[:, :, :] == A[:, :, :]).all() - # Conditional with secoundary K index on a while loop, - # simplified from in-situ code of NOAA's SHiELD microphysics + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +def test_K_offset_write_conditional(backend): + # While loop have a bug in `dace:X` backends where + # the read-connector is used which means you can never + # update the field in a while. + if backend == "dace:cpu": + pytest.skip("Internal compiler error in GitHub action container") + + array_shape = (1, 1, 4) + K_values = np.arange(start=40, stop=44) + @gtscript.stencil(backend=backend) def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): with computation(BACKWARD), interval(1, None): @@ -669,10 +679,12 @@ def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scala lev = lev + 1 A = gt_storage.zeros( - backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64 + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) A[:, :, :] = K_values - B = gt_storage.ones(backend=backend, aligned_index=(0, 0, 0), shape=(1, 1, 4), dtype=np.float64) + B = gt_storage.ones( + backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 + ) column_physics_conditional(A, B, 2.0) assert (A[0, 0, :] == [2, 2, -1, -1]).all() assert (B[0, 0, :] == [1, -1, 2, 42]).all() From 239093c3e77ecd8e9874d02b9fa12805b3764464 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 23 Feb 2024 15:39:53 -0500 Subject: [PATCH 10/25] Add GPU init/read support in test Restrict the test to the working backends --- tests/cartesian_tests/definitions.py | 11 ++++++++++ .../feature_tests/test_field_layouts.py | 13 +----------- .../test_code_generation.py | 20 ++++++++++++------- 3 files changed, 25 insertions(+), 19 deletions(-) diff --git a/tests/cartesian_tests/definitions.py b/tests/cartesian_tests/definitions.py index 5ce4095e26..ae93a8f168 100644 --- a/tests/cartesian_tests/definitions.py +++ b/tests/cartesian_tests/definitions.py @@ -21,6 +21,7 @@ import datetime +import numpy as np import pytest from gt4py import cartesian as gt4pyc @@ -60,3 +61,13 @@ def _get_backends_with_storage_info(storage_info_kind: str): @pytest.fixture() def id_version(): return gt_utils.shashed_id(str(datetime.datetime.now())) + + +def _get_array_library(backend: str): + backend_cls = gt4pyc.backend.from_name(backend) + assert backend_cls is not None + if backend_cls.storage_info["device"] == "gpu": + assert cp is not None + return cp + else: + return np diff --git a/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py b/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py index c35594eb71..67b925162a 100644 --- a/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py +++ b/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py @@ -12,13 +12,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np import pytest from gt4py import cartesian as gt4pyc, storage as gt_storage from gt4py.cartesian import gtscript -from cartesian_tests.definitions import ALL_BACKENDS, PERFORMANCE_BACKENDS +from cartesian_tests.definitions import ALL_BACKENDS, PERFORMANCE_BACKENDS, _get_array_library from cartesian_tests.integration_tests.multi_feature_tests.stencil_definitions import copy_stencil @@ -28,16 +27,6 @@ cp = None -def _get_array_library(backend: str): - backend_cls = gt4pyc.backend.from_name(backend) - assert backend_cls is not None - if backend_cls.storage_info["device"] == "gpu": - assert cp is not None - return cp - else: - return np - - @pytest.mark.parametrize("backend", ALL_BACKENDS) @pytest.mark.parametrize("order", ["C", "F"]) def test_numpy_allocators(backend, order): diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 04fd3d5066..d9f8d1d13d 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -32,7 +32,7 @@ ) from gt4py.storage.cartesian import utils as storage_utils -from cartesian_tests.definitions import ALL_BACKENDS, CPU_BACKENDS +from cartesian_tests.definitions import ALL_BACKENDS, CPU_BACKENDS, _get_array_library from cartesian_tests.integration_tests.multi_feature_tests.stencil_definitions import ( EXTERNALS_REGISTRY as externals_registry, REGISTRY as stencil_definitions, @@ -590,8 +590,13 @@ def test(out: Field[np.float64], inp: Field[np.float64]): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_K_offset_write(backend): + # Cuda generates bad code for the K offset + if backend == "cuda": + pytest.skip("cuda K-offset write generates bad code") + + arraylib = _get_array_library(backend) array_shape = (1, 1, 4) - K_values = np.arange(start=40, stop=44) + K_values = arraylib.arange(start=40, stop=44) # Simple case of writing ot an offset. # A is untouched @@ -660,11 +665,12 @@ def test_K_offset_write_conditional(backend): # While loop have a bug in `dace:X` backends where # the read-connector is used which means you can never # update the field in a while. - if backend == "dace:cpu": - pytest.skip("Internal compiler error in GitHub action container") + if backend.startswith("dace") or backend == "cuda": + pytest.skip("DaCe backends have a bug when handling while loop") + arraylib = _get_array_library(backend) array_shape = (1, 1, 4) - K_values = np.arange(start=40, stop=44) + K_values = arraylib.arange(start=40, stop=44) @gtscript.stencil(backend=backend) def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): @@ -686,5 +692,5 @@ def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scala backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64 ) column_physics_conditional(A, B, 2.0) - assert (A[0, 0, :] == [2, 2, -1, -1]).all() - assert (B[0, 0, :] == [1, -1, 2, 42]).all() + assert (A[0, 0, :] == arraylib.array([2, 2, -1, -1])).all() + assert (B[0, 0, :] == arraylib.array([1, -1, 2, 42])).all() From d956202b4b9396099f21a2b8433b548c2ac53b5d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 15 Mar 2024 10:12:47 -0400 Subject: [PATCH 11/25] lint --- .../test_code_generation.py | 24 +- .../frontend_tests/test_gtscript_frontend.py | 294 +++++------------- 2 files changed, 88 insertions(+), 230 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 3b5c24754b..64059a2f68 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -65,9 +65,7 @@ def test_generation(name, backend): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_lazy_stencil(backend): @gtscript.lazy_stencil(backend=backend) - def definition( - field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_] - ): + def definition(field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): field_a = field_b @@ -150,9 +148,7 @@ def test_stage_merger_induced_interval_block_reordering(backend): ) @gtscript.stencil(backend=backend) - def stencil( - field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_] - ): + def stencil(field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_]): with computation(BACKWARD): with interval(-2, -1): # block 1 field_out = field_in @@ -297,9 +293,7 @@ def copy_2to3( with computation(FORWARD), interval(...): outp[0, 0, 0] = inp - inp_f = gt_storage.from_array( - np.random.randn(10, 10), aligned_index=(0, 0), backend=backend - ) + inp_f = gt_storage.from_array(np.random.randn(10, 10), aligned_index=(0, 0), backend=backend) outp_f = gt_storage.from_array( np.random.randn(10, 10, 10), aligned_index=(0, 0, 0), backend=backend ) @@ -484,9 +478,7 @@ def stencil( full_shape, backend=backend, aligned_index=aligned_index, dtype=INT32_VEC2 ) - gtscript.stencil(definition=stencil, backend=backend)( - input_field, output_field, index := 1 - ) + gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, index := 1) assert output_field[0, 0, 0, index] == 1 @@ -583,9 +575,7 @@ def k_to_ijk(outp: Field[np.float64], inp: Field[gtscript.K, np.float64]): inp = storage_utils.cpu_copy(inp) outp = storage_utils.cpu_copy(outp) np.testing.assert_allclose(data, inp) - np.testing.assert_allclose( - np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1] - ) + np.testing.assert_allclose(np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1]) np.testing.assert_allclose(0.0, outp[:, :, 0]) np.testing.assert_allclose(0.0, outp[:, :, -1]) @@ -693,9 +683,7 @@ def test_K_offset_write_conditional(backend): K_values = arraylib.arange(start=40, stop=44) @gtscript.stencil(backend=backend) - def column_physics_conditional( - A: Field[np.float64], B: Field[np.float64], scalar: np.float64 - ): + def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): with computation(BACKWARD), interval(1, None): if A > 0 and B > 0: A[0, 0, -1] = scalar diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py index fca1d22506..d456b54045 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py @@ -58,9 +58,7 @@ def parse_definition( rebuild=False, **kwargs, ): - original_annotations = gtscript._set_arg_dtypes( - definition_func, dtypes=dtypes or {} - ) + original_annotations = gtscript._set_arg_dtypes(definition_func, dtypes=dtypes or {}) build_options = gt_definitions.BuildOptions( name=name, @@ -87,9 +85,7 @@ def parse_definition( GLOBAL_BOOL_CONSTANT = True GLOBAL_CONSTANT = 1.0 GLOBAL_NESTED_CONSTANTS = types.SimpleNamespace(A=100, B=200) -GLOBAL_VERY_NESTED_CONSTANTS = types.SimpleNamespace( - nested=types.SimpleNamespace(A=1000, B=2000) -) +GLOBAL_VERY_NESTED_CONSTANTS = types.SimpleNamespace(nested=types.SimpleNamespace(A=1000, B=2000)) @gtscript.function @@ -121,9 +117,7 @@ def definition_func(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): inout_field = inout_field[0, 0, 0] + MISSING_CONSTANT - with pytest.raises( - gt_frontend.GTScriptSymbolError, match=r".*MISSING_CONSTANT.*" - ): + with pytest.raises(gt_frontend.GTScriptSymbolError, match=r".*MISSING_CONSTANT.*"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -220,9 +214,7 @@ def definition_func(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): inout_field = inout_field[0, 0, 0] + WRONG_VALUE_CONSTANT - with pytest.raises( - gt_frontend.GTScriptDefinitionError, match=r".*WRONG_VALUE_CONSTANT.*" - ): + with pytest.raises(gt_frontend.GTScriptDefinitionError, match=r".*WRONG_VALUE_CONSTANT.*"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -370,9 +362,7 @@ def definition_func(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): inout_field = func2(inout_field) - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="recursive function call" - ): + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="recursive function call"): parse_definition( definition_func, name="testf", @@ -402,9 +392,7 @@ def definition_func(phi: gtscript.Field[np.float64]): class TestAxisSyntax: def test_good_syntax(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): out_field = in_field[J - 1] + in_field[J] @@ -413,9 +401,7 @@ def definition_func( ) def test_good_syntax_external(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): from gt4py.cartesian.__externals__ import AXIS with computation(PARALLEL), interval(...): @@ -429,9 +415,7 @@ def definition_func( ) def test_good_syntax_external_value(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): from gt4py.cartesian.__externals__ import VALUE with computation(PARALLEL), interval(...): @@ -446,9 +430,7 @@ def definition_func( ) def test_bad_mul_syntax(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): out_field = in_field[I * 1] @@ -460,9 +442,7 @@ def definition_func( ) def test_bad_dup_add(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): out_field = in_field[I + 1 + I] @@ -474,9 +454,7 @@ def definition_func( ) def test_bad_dup_axis(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): out_field = in_field[I, I - 1] @@ -488,9 +466,7 @@ def definition_func( ) def test_bad_out_of_order(self): - def definition_func( - in_field: gtscript.Field[float], out_field: gtscript.Field[float] - ): + def definition_func(in_field: gtscript.Field[float], out_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): out_field = in_field[J, I - 1] @@ -541,9 +517,7 @@ def definition_func(inout_field: gtscript.Field[float]): ) def test_missing(self): - externals = dict( - CONSTANT=-2.0, NESTED_CONSTANTS=types.SimpleNamespace(A=-100, B=-200) - ) + externals = dict(CONSTANT=-2.0, NESTED_CONSTANTS=types.SimpleNamespace(A=-100, B=-200)) def definition_func(inout_field: gtscript.Field[float]): from gt4py.cartesian.__externals__ import MISSING_CONSTANT @@ -551,9 +525,7 @@ def definition_func(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): inout_field = inout_field[0, 0, 0] + MISSING_CONSTANT - with pytest.raises( - gt_frontend.GTScriptDefinitionError, match=r".*MISSING_CONSTANT.*" - ): + with pytest.raises(gt_frontend.GTScriptDefinitionError, match=r".*MISSING_CONSTANT.*"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -586,9 +558,7 @@ def definition_func(inout_field: gtscript.Field[float]): externals = dict(WRONG_VALUE_CONSTANT=value_type()) - with pytest.raises( - gt_frontend.GTScriptDefinitionError, match=r".*WRONG_VALUE_CONSTANT.*" - ): + with pytest.raises(gt_frontend.GTScriptDefinitionError, match=r".*WRONG_VALUE_CONSTANT.*"): parse_definition( definition_func, externals=externals, @@ -691,9 +661,7 @@ def definition_func(field: gtscript.Field[float]): with computation(PARALLEL), interval(K[2], -1): field = 0 - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="Two-argument syntax" - ): + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Two-argument syntax"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -723,9 +691,7 @@ def definition_func(field: gtscript.Field[float]): with interval(-1, None): field = 1 - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="Overlapping intervals" - ): + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Overlapping intervals"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -740,9 +706,7 @@ def definition_func(field: gtscript.Field[float]): with interval(2, None): field = 1 - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="Overlapping intervals" - ): + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Overlapping intervals"): parse_definition( definition_func, name=inspect.stack()[0][3], @@ -856,9 +820,7 @@ def stencil(in_f: gtscript.Field[np.float_]): in_f = 1.0 with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Unknown symbol"): - parse_definition( - stencil, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(stencil, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_error_nested(self): def stencil(in_f: gtscript.Field[np.float_]): @@ -873,18 +835,14 @@ def stencil(in_f: gtscript.Field[np.float_]): gt_frontend.GTScriptSyntaxError, match="Cannot nest `with` node inside horizontal region", ): - parse_definition( - stencil, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(stencil, name=inspect.stack()[0][3], module=self.__class__.__name__) class TestExternalsWithSubroutines: def test_all_legal_combinations(self): @gtscript.function def _stage_laplacian_x(dx, phi): - lap = add_external_const( - phi[-1, 0, 0] - 2.0 * phi[0, 0, 0] + phi[1, 0, 0] - ) / (dx * dx) + lap = add_external_const(phi[-1, 0, 0] - 2.0 * phi[0, 0, 0] + phi[1, 0, 0]) / (dx * dx) return lap @gtscript.function @@ -1039,9 +997,7 @@ def return_tuple(): tmp2 = 2 return tmp1, tmp2 - def definition_func( - res1: gtscript.Field[np.float64], res2: gtscript.Field[np.float64] - ): + def definition_func(res1: gtscript.Field[np.float64], res2: gtscript.Field[np.float64]): with computation(PARALLEL), interval(...): res1, res2 = return_tuple() @@ -1066,9 +1022,7 @@ def definition(inout_field: gtscript.Field[float]): module=self.__class__.__name__, ) - with pytest.raises( - gt_frontend.GTScriptAssertionError, match="Assertion failed" - ): + with pytest.raises(gt_frontend.GTScriptAssertionError, match="Assertion failed"): parse_definition( definition, externals={"EXTERNAL": 1}, @@ -1080,13 +1034,9 @@ def test_nested_attribute(self): def definition(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): compile_assert(GLOBAL_VERY_NESTED_CONSTANTS.nested.A > 1) - inout_field = ( - inout_field[0, 0, 0] + GLOBAL_VERY_NESTED_CONSTANTS.nested.A - ) + inout_field = inout_field[0, 0, 0] + GLOBAL_VERY_NESTED_CONSTANTS.nested.A - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_inside_func(self): @gtscript.function @@ -1098,9 +1048,7 @@ def definition(inout_field: gtscript.Field[float]): with computation(PARALLEL), interval(...): inout_field = assert_in_func(inout_field) - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_runtime_error(self): def definition(inout_field: gtscript.Field[float]): @@ -1111,9 +1059,7 @@ def definition(inout_field: gtscript.Field[float]): gt_frontend.GTScriptSyntaxError, match="Evaluation of compile_assert condition failed", ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) class TestReducedDimensions: @@ -1160,9 +1106,7 @@ def definition( gt_frontend.GTScriptSyntaxError, match="Incorrect offset specification detected. Found .* but the field has dimensions .*", ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_error_write_1d(self): def definition( @@ -1176,9 +1120,7 @@ def definition( gt_frontend.GTScriptSyntaxError, match="Cannot assign to field .* as all parallel axes .* are not present", ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_higher_dim_temp(self): def definition( @@ -1210,9 +1152,7 @@ def definition( gt_frontend.GTScriptSyntaxError, match="Found IJ, but only IJK is currently supported for temporaries", ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) class TestDataDimensions: @@ -1227,9 +1167,7 @@ def definition( field_out[0, 0, 0][1] = field_in field_out[0, 0, 0][2] = field_in[0, 0, 0] + another_field[0, 0, 0][2] - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_syntax_no_datadim(self): def definition( @@ -1241,12 +1179,8 @@ def definition( field_out[0, 0, 0][1] = field_in field_out[0, 0, 0][2] = field_in[0, 0, 0][0] - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="Incorrect data index length" - ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Incorrect data index length"): + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_syntax_out_bounds(self): def definition( @@ -1256,12 +1190,8 @@ def definition( with computation(PARALLEL), interval(...): field_out[0, 0, 0][3] = field_in[0, 0, 0] - with pytest.raises( - gt_frontend.GTScriptSyntaxError, match="Data index out of bounds" - ): - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + with pytest.raises(gt_frontend.GTScriptSyntaxError, match="Data index out of bounds"): + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_indirect_access_read(self): def definition( @@ -1275,9 +1205,7 @@ def definition( def_ir = parse_definition( definition, name=inspect.stack()[0][3], module=self.__class__.__name__ ) - assert isinstance( - def_ir.computations[0].body.stmts[0].value.data_index[0], nodes.VarRef - ) + assert isinstance(def_ir.computations[0].body.stmts[0].value.data_index[0], nodes.VarRef) def test_indirect_access_write(self): def definition( @@ -1291,9 +1219,7 @@ def definition( def_ir = parse_definition( definition, name=inspect.stack()[0][3], module=self.__class__.__name__ ) - assert isinstance( - def_ir.computations[0].body.stmts[0].target.data_index[0], nodes.VarRef - ) + assert isinstance(def_ir.computations[0].body.stmts[0].target.data_index[0], nodes.VarRef) class TestImports: @@ -1323,16 +1249,14 @@ def definition_func(inout_field: gtscript.Field[float]): @pytest.mark.parametrize( "id_case,import_line", list( - enumerate( - [ - "import gt4py", - "from externals import EXTERNAL", - "from gt4py.cartesian import __gtscript__", - "from gt4py.cartesian import __externals__", - "from gt4py.cartesian.gtscript import computation", - "from gt4py.cartesian.externals import EXTERNAL", - ] - ) + enumerate([ + "import gt4py", + "from externals import EXTERNAL", + "from gt4py.cartesian import __gtscript__", + "from gt4py.cartesian import __externals__", + "from gt4py.cartesian.gtscript import computation", + "from gt4py.cartesian.externals import EXTERNAL", + ]) ), ) def test_wrong_imports(self, id_case, import_line): @@ -1363,25 +1287,21 @@ class TestDTypes: @pytest.mark.parametrize( "id_case,test_dtype", list( - enumerate( - [ - bool, - np.bool_, - int, - np.int32, - np.int64, - float, - np.float32, - np.float64, - np.dtype((np.float32, (3,))), - ] - ) + enumerate([ + bool, + np.bool_, + int, + np.int32, + np.int64, + float, + np.float32, + np.float64, + np.dtype((np.float32, (3,))), + ]) ), ) def test_all_legal_dtypes_instance(self, id_case, test_dtype): - test_base_dtype = ( - test_dtype.base.type if isinstance(test_dtype, np.dtype) else test_dtype - ) + test_base_dtype = test_dtype.base.type if isinstance(test_dtype, np.dtype) else test_dtype def definition_func( in_field: gtscript.Field[test_dtype], @@ -1475,26 +1395,18 @@ def test_literal_floating_parametrization(self, the_float): class TestAssignmentSyntax: def test_ellipsis(self): - def func( - in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_] - ): + def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): out_field[...] = in_field - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_offset(self): - def func( - in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_] - ): + def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): out_field[0, 0, 0] = in_field - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) with pytest.raises(gt_frontend.GTScriptSyntaxError): @@ -1505,13 +1417,9 @@ def func( with computation(PARALLEL), interval(...): out_field[0, 0, 1] = in_field - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) - def func( - in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_] - ): + def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_]): from gt4py.cartesian.__externals__ import offset with computation(PARALLEL), interval(...): @@ -1583,9 +1491,7 @@ def func( with computation(PARALLEL), interval(...): out_field[:, :, :] = in_field - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_string(self): with pytest.raises(gt_frontend.GTScriptSyntaxError): @@ -1597,9 +1503,7 @@ def func( with computation(PARALLEL), interval(...): out_field["a_key"] = in_field - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_augmented(self): def func(in_field: gtscript.Field[np.float_]): @@ -1609,18 +1513,14 @@ def func(in_field: gtscript.Field[np.float_]): in_field /= 0.5 in_field *= 4.0 - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_K_offset_write(self): def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): with computation(FORWARD), interval(...): out[0, 0, 1] = inp - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) with pytest.raises( gt_frontend.GTScriptSyntaxError, @@ -1631,9 +1531,7 @@ def func(out: gtscript.Field[np.float64], inp: gtscript.Field[np.float64]): with computation(PARALLEL), interval(...): out[0, 0, 1] = inp - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_datadims_direct_access(self): # Check classic data dimensions are working @@ -1644,9 +1542,7 @@ def data_dims( with computation(PARALLEL), interval(...): out_field = global_field[0, 0, 0][1, 0, 2] - parse_definition( - data_dims, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(data_dims, name=inspect.stack()[0][3], module=self.__class__.__name__) # Check .A on read def at_read( @@ -1656,9 +1552,7 @@ def at_read( with computation(PARALLEL), interval(...): out_field = global_field.A[1, 0, 2, 2] - parse_definition( - at_read, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(at_read, name=inspect.stack()[0][3], module=self.__class__.__name__) # Can't write to the field def at_write( @@ -1672,9 +1566,7 @@ def at_write( gt_frontend.GTScriptSyntaxError, match="writing to an GlobalTable \('A' global indexation\) is forbidden", ): - parse_definition( - at_write, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(at_write, name=inspect.stack()[0][3], module=self.__class__.__name__) # Can't index cartesian style def GlobalTable_access_as_IJK( @@ -1711,16 +1603,12 @@ def data_dims_with_at( class TestNestedWithSyntax: def test_nested_with(self): - def definition( - in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_] - ): + def definition(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_]): with computation(PARALLEL): with interval(...): in_field = out_field - parse_definition( - definition, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(definition, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_nested_with_ordering(self): def definition_fw( @@ -1763,45 +1651,35 @@ def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += sin(in_field) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_offset_arg(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += sin(in_field[1, 0, 0]) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_nested_calls(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += sin(abs(in_field)) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_nested_external_call(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += sin(add_external_const(in_field)) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_multi_nested_calls(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += min(abs(sin(add_external_const(in_field))), -0.5) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_native_in_function(self): @gtscript.function @@ -1812,36 +1690,28 @@ def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field += sinus(in_field) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_native_function_unary(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field = not isfinite(in_field) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_native_function_binary(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field = asin(in_field) + 1 - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) def test_native_function_ternary(self): def func(in_field: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): in_field = asin(in_field) + 1 if 1 < in_field else sin(in_field) - parse_definition( - func, name=inspect.stack()[0][3], module=self.__class__.__name__ - ) + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) class TestWarnInlined: From 899a5dae8e96283123a2416c549a1d88dfc25945 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 15 Mar 2024 10:21:25 -0400 Subject: [PATCH 12/25] Verbose current `dace:X` bug around while loop --- .../test_code_generation.py | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 64059a2f68..a36b178b95 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -65,7 +65,9 @@ def test_generation(name, backend): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_lazy_stencil(backend): @gtscript.lazy_stencil(backend=backend) - def definition(field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_]): + def definition( + field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_] + ): with computation(PARALLEL), interval(...): field_a = field_b @@ -148,7 +150,9 @@ def test_stage_merger_induced_interval_block_reordering(backend): ) @gtscript.stencil(backend=backend) - def stencil(field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_]): + def stencil( + field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_] + ): with computation(BACKWARD): with interval(-2, -1): # block 1 field_out = field_in @@ -293,7 +297,9 @@ def copy_2to3( with computation(FORWARD), interval(...): outp[0, 0, 0] = inp - inp_f = gt_storage.from_array(np.random.randn(10, 10), aligned_index=(0, 0), backend=backend) + inp_f = gt_storage.from_array( + np.random.randn(10, 10), aligned_index=(0, 0), backend=backend + ) outp_f = gt_storage.from_array( np.random.randn(10, 10, 10), aligned_index=(0, 0, 0), backend=backend ) @@ -478,7 +484,9 @@ def stencil( full_shape, backend=backend, aligned_index=aligned_index, dtype=INT32_VEC2 ) - gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, index := 1) + gtscript.stencil(definition=stencil, backend=backend)( + input_field, output_field, index := 1 + ) assert output_field[0, 0, 0, index] == 1 @@ -575,7 +583,9 @@ def k_to_ijk(outp: Field[np.float64], inp: Field[gtscript.K, np.float64]): inp = storage_utils.cpu_copy(inp) outp = storage_utils.cpu_copy(outp) np.testing.assert_allclose(data, inp) - np.testing.assert_allclose(np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1]) + np.testing.assert_allclose( + np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1] + ) np.testing.assert_allclose(0.0, outp[:, :, 0]) np.testing.assert_allclose(0.0, outp[:, :, -1]) @@ -675,6 +685,7 @@ def test_K_offset_write_conditional(backend): # While loop have a bug in `dace:X` backends where # the read-connector is used which means you can never # update the field in a while. + # Logged in: https://github.com/GridTools/gt4py/issues/1496 if backend.startswith("dace") or backend == "cuda": pytest.skip("DaCe backends have a bug when handling while loop.") @@ -683,7 +694,9 @@ def test_K_offset_write_conditional(backend): K_values = arraylib.arange(start=40, stop=44) @gtscript.stencil(backend=backend) - def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): + def column_physics_conditional( + A: Field[np.float64], B: Field[np.float64], scalar: np.float64 + ): with computation(BACKWARD), interval(1, None): if A > 0 and B > 0: A[0, 0, -1] = scalar From b5f112a3c9a2f74b8af74a6f27f1524edb03fc45 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 15 Mar 2024 10:25:30 -0400 Subject: [PATCH 13/25] lint --- .../test_code_generation.py | 24 +++++-------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index a36b178b95..cad6eda7c5 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -65,9 +65,7 @@ def test_generation(name, backend): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_lazy_stencil(backend): @gtscript.lazy_stencil(backend=backend) - def definition( - field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_] - ): + def definition(field_a: gtscript.Field[np.float_], field_b: gtscript.Field[np.float_]): with computation(PARALLEL), interval(...): field_a = field_b @@ -150,9 +148,7 @@ def test_stage_merger_induced_interval_block_reordering(backend): ) @gtscript.stencil(backend=backend) - def stencil( - field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_] - ): + def stencil(field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_]): with computation(BACKWARD): with interval(-2, -1): # block 1 field_out = field_in @@ -297,9 +293,7 @@ def copy_2to3( with computation(FORWARD), interval(...): outp[0, 0, 0] = inp - inp_f = gt_storage.from_array( - np.random.randn(10, 10), aligned_index=(0, 0), backend=backend - ) + inp_f = gt_storage.from_array(np.random.randn(10, 10), aligned_index=(0, 0), backend=backend) outp_f = gt_storage.from_array( np.random.randn(10, 10, 10), aligned_index=(0, 0, 0), backend=backend ) @@ -484,9 +478,7 @@ def stencil( full_shape, backend=backend, aligned_index=aligned_index, dtype=INT32_VEC2 ) - gtscript.stencil(definition=stencil, backend=backend)( - input_field, output_field, index := 1 - ) + gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, index := 1) assert output_field[0, 0, 0, index] == 1 @@ -583,9 +575,7 @@ def k_to_ijk(outp: Field[np.float64], inp: Field[gtscript.K, np.float64]): inp = storage_utils.cpu_copy(inp) outp = storage_utils.cpu_copy(outp) np.testing.assert_allclose(data, inp) - np.testing.assert_allclose( - np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1] - ) + np.testing.assert_allclose(np.broadcast_to(data[2:], shape=(2, 2, 8)), outp[:, :, 1:-1]) np.testing.assert_allclose(0.0, outp[:, :, 0]) np.testing.assert_allclose(0.0, outp[:, :, -1]) @@ -694,9 +684,7 @@ def test_K_offset_write_conditional(backend): K_values = arraylib.arange(start=40, stop=44) @gtscript.stencil(backend=backend) - def column_physics_conditional( - A: Field[np.float64], B: Field[np.float64], scalar: np.float64 - ): + def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): with computation(BACKWARD), interval(1, None): if A > 0 and B > 0: A[0, 0, -1] = scalar From 6916121aaf654b00e92e8e87947ccad4d2f9db80 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 30 Apr 2024 13:52:44 -0400 Subject: [PATCH 14/25] Allow K offset writes in `dace` backends Rename connector for better tracking in SDFG var_offset_fields -> offset_in_K_fields --- .../gtc/dace/expansion/daceir_builder.py | 60 ++++++++++++------- src/gt4py/cartesian/gtc/dace/utils.py | 10 +++- .../test_code_generation.py | 7 +-- 3 files changed, 48 insertions(+), 29 deletions(-) diff --git a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py index 8435d32158..a9c3a6fc05 100644 --- a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py +++ b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py @@ -326,23 +326,34 @@ def visit_FieldAccess( *, is_target: bool, targets: Set[eve.SymbolRef], - var_offset_fields: Set[eve.SymbolRef], + offset_in_K_fields: Set[eve.SymbolRef], **kwargs: Any, ) -> Union[dcir.IndexAccess, dcir.ScalarAccess]: + """Generate the relevant accessor to match the memlet that was previously setup + + This function was written when offset writing was forbidden. It has been refactor + to allow offset write in K with minimum changes. + """ + res: Union[dcir.IndexAccess, dcir.ScalarAccess] - if node.name in var_offset_fields: - res = dcir.IndexAccess( - name=node.name + "__", - offset=self.visit( - node.offset, - is_target=False, - targets=targets, - var_offset_fields=var_offset_fields, - **kwargs, - ), - data_index=node.data_index, - dtype=node.dtype, - ) + if node.name in offset_in_K_fields: + is_target = is_target or node.name in targets + name = get_tasklet_symbol(node.name, node.offset, is_target=is_target) + if is_target: + res = dcir.IndexAccess( + name=name, + offset=self.visit( + node.offset, + is_target=is_target, + targets=targets, + offset_in_K_fields=offset_in_K_fields, + **kwargs, + ), + data_index=node.data_index, + dtype=node.dtype, + ) + else: + res = dcir.ScalarAccess(name=name, dtype=node.dtype) else: is_target = is_target or ( node.name in targets and node.offset == common.CartesianOffset.zero() @@ -354,11 +365,7 @@ def visit_FieldAccess( ) else: res = dcir.ScalarAccess(name=name, dtype=node.dtype) - # Because we allow writing in K, we allow targets who have been - # writing in K to be omitted so the original connector can be re-used. - # Previous guardrails in `gtscript` restrict the writes - # to non-parallel K-loop so we don't have issues. - if is_target and node.offset.to_dict()["k"] == 0: + if is_target: targets.add(node.name) return res @@ -826,11 +833,22 @@ def visit_VerticalLoop( ) ) - var_offset_fields = { + # Offsets in K can be both: + # - read indexed via an expression, + # - write offset in K (both in expression and on scalar) + # We get all indexed via expression + offset_in_K_fields = { acc.name for acc in node.walk_values().if_isinstance(oir.FieldAccess) if isinstance(acc.offset, oir.VariableKOffset) } + # We add write offset to K + for assign_node in node.walk_values().if_isinstance(oir.AssignStmt): + if isinstance(assign_node.left, oir.FieldAccess): + acc = assign_node.left + if isinstance(acc.offset, common.CartesianOffset) and acc.offset.k != 0: + offset_in_K_fields.add(acc.name) + sections_idx = next( idx for idx, item in enumerate(global_ctx.library_node.expansion_specification) @@ -847,7 +865,7 @@ def visit_VerticalLoop( global_ctx=global_ctx, iteration_ctx=iteration_ctx, symbol_collector=symbol_collector, - var_offset_fields=var_offset_fields, + offset_in_K_fields=offset_in_K_fields, **kwargs, ) ) diff --git a/src/gt4py/cartesian/gtc/dace/utils.py b/src/gt4py/cartesian/gtc/dace/utils.py index dac0c8acc5..7bb9f8d571 100644 --- a/src/gt4py/cartesian/gtc/dace/utils.py +++ b/src/gt4py/cartesian/gtc/dace/utils.py @@ -62,9 +62,9 @@ def replace_strides(arrays, get_layout_map): def get_tasklet_symbol(name, offset, is_target): if is_target: - return f"__{name}" + return f"gtOUT__{name}" - acc_name = name + "__" + acc_name = f"gtIN__{name}" if offset is not None: offset_strs = [] for axis in dcir.Axis.dims_3d(): @@ -231,9 +231,12 @@ def _make_access_info( region, he_grid, grid_subset, + is_write, ) -> "dcir.FieldAccessInfo": + # Check we have expression offsets in K + # OR write offsets in K offset = [offset_node.to_dict()[k] for k in "ijk"] - if isinstance(offset_node, oir.VariableKOffset): + if isinstance(offset_node, oir.VariableKOffset) or (offset[2] != 0 and is_write): variable_offset_axes = [dcir.Axis.K] else: variable_offset_axes = [] @@ -292,6 +295,7 @@ def visit_FieldAccess( region=region, he_grid=he_grid, grid_subset=grid_subset, + is_write=is_write, ) ctx.access_infos[node.name] = access_info.union( ctx.access_infos.get(node.name, access_info) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index cad6eda7c5..f49699453f 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -673,11 +673,8 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_K_offset_write_conditional(backend): # While loop have a bug in `dace:X` backends where - # the read-connector is used which means you can never - # update the field in a while. - # Logged in: https://github.com/GridTools/gt4py/issues/1496 - if backend.startswith("dace") or backend == "cuda": - pytest.skip("DaCe backends have a bug when handling while loop.") + if backend == "cuda": + pytest.skip("Cuda backend is not capable of K offset write") arraylib = _get_array_library(backend) array_shape = (1, 1, 4) From 9803336374eaed537d15e1743d004cee38105461 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 1 May 2024 12:20:09 -0400 Subject: [PATCH 15/25] Fix Assign guard on K for I and IJ fields --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index 6072edb1e0..30933baf64 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1447,7 +1447,11 @@ def visit_Assign(self, node: ast.Assign) -> list: message="Assignment to non-zero offsets is not supported in IJ.", loc=nodes.Location.from_ast_node(t), ) - if spatial_offset[2] != 0 and self.iteration_order == nodes.IterationOrder.PARALLEL: + if ( + len(spatial_offset) == 3 + and spatial_offset[2] != 0 + and self.iteration_order == nodes.IterationOrder.PARALLEL + ): raise GTScriptSyntaxError( message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", loc=nodes.Location.from_ast_node(t), From 5e2a622bbd949656e9aae4acdb14a8da648ccc18 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 14 May 2024 12:44:19 -0400 Subject: [PATCH 16/25] Make `get_array_library` public Remove dead comment --- tests/cartesian_tests/definitions.py | 3 ++- .../integration_tests/feature_tests/test_field_layouts.py | 6 +++--- .../multi_feature_tests/test_code_generation.py | 7 +++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/cartesian_tests/definitions.py b/tests/cartesian_tests/definitions.py index a17e89675b..c809830bc4 100644 --- a/tests/cartesian_tests/definitions.py +++ b/tests/cartesian_tests/definitions.py @@ -63,7 +63,8 @@ def id_version(): return gt_utils.shashed_id(str(datetime.datetime.now())) -def _get_array_library(backend: str): +def get_array_library(backend: str): + """Return device ready array maker library""" backend_cls = gt4pyc.backend.from_name(backend) assert backend_cls is not None if backend_cls.storage_info["device"] == "gpu": diff --git a/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py b/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py index 67b925162a..ab79d06c6f 100644 --- a/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py +++ b/tests/cartesian_tests/integration_tests/feature_tests/test_field_layouts.py @@ -17,7 +17,7 @@ from gt4py import cartesian as gt4pyc, storage as gt_storage from gt4py.cartesian import gtscript -from cartesian_tests.definitions import ALL_BACKENDS, PERFORMANCE_BACKENDS, _get_array_library +from cartesian_tests.definitions import ALL_BACKENDS, PERFORMANCE_BACKENDS, get_array_library from cartesian_tests.integration_tests.multi_feature_tests.stencil_definitions import copy_stencil @@ -30,7 +30,7 @@ @pytest.mark.parametrize("backend", ALL_BACKENDS) @pytest.mark.parametrize("order", ["C", "F"]) def test_numpy_allocators(backend, order): - xp = _get_array_library(backend) + xp = get_array_library(backend) shape = (20, 10, 5) inp = xp.array(xp.random.randn(*shape), order=order, dtype=xp.float_) outp = xp.zeros(shape=shape, order=order, dtype=xp.float_) @@ -43,7 +43,7 @@ def test_numpy_allocators(backend, order): @pytest.mark.parametrize("backend", PERFORMANCE_BACKENDS) def test_bad_layout_warns(backend): - xp = _get_array_library(backend) + xp = get_array_library(backend) backend_cls = gt4pyc.backend.from_name(backend) assert backend_cls is not None diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 1a56686204..e1f2a8fcc3 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -33,7 +33,7 @@ ) from gt4py.storage.cartesian import utils as storage_utils -from cartesian_tests.definitions import ALL_BACKENDS, CPU_BACKENDS, _get_array_library +from cartesian_tests.definitions import ALL_BACKENDS, CPU_BACKENDS, get_array_library from cartesian_tests.integration_tests.multi_feature_tests.stencil_definitions import ( EXTERNALS_REGISTRY as externals_registry, REGISTRY as stencil_definitions, @@ -595,7 +595,7 @@ def test_K_offset_write(backend): if backend == "cuda": pytest.skip("cuda K-offset write generates bad code") - arraylib = _get_array_library(backend) + arraylib = get_array_library(backend) array_shape = (1, 1, 4) K_values = arraylib.arange(start=40, stop=44) @@ -663,11 +663,10 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): @pytest.mark.parametrize("backend", ALL_BACKENDS) def test_K_offset_write_conditional(backend): - # While loop have a bug in `dace:X` backends where if backend == "cuda": pytest.skip("Cuda backend is not capable of K offset write") - arraylib = _get_array_library(backend) + arraylib = get_array_library(backend) array_shape = (1, 1, 4) K_values = arraylib.arange(start=40, stop=44) From 8fb7452a3cf2b45268ab98a654259e069659f5ca Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Sep 2024 14:02:11 -0400 Subject: [PATCH 17/25] Error out cleanly for `gt:gpu` and CUDA < 12 --- .../cartesian/frontend/gtscript_frontend.py | 19 ++++++++++++++----- src/gt4py/cartesian/stencil_builder.py | 1 + .../test_code_generation.py | 8 +++++++- .../frontend_tests/test_ir_maker.py | 2 +- 4 files changed, 23 insertions(+), 7 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index 60846912b3..f9fc47168e 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -708,6 +708,7 @@ def __init__( fields: dict, parameters: dict, local_symbols: dict, + backend_name: str, *, domain: nodes.Domain, temp_decls: Optional[Dict[str, nodes.FieldDecl]] = None, @@ -721,6 +722,7 @@ def __init__( isinstance(value, (type, np.dtype)) for value in local_symbols.values() ) + self.backend_name = backend_name self.fields = fields self.parameters = parameters self.local_symbols = local_symbols @@ -1437,15 +1439,21 @@ def visit_Assign(self, node: ast.Assign) -> list: message="Assignment to non-zero offsets is not supported in IJ.", loc=nodes.Location.from_ast_node(t), ) + # Case of K-offset if ( len(spatial_offset) == 3 and spatial_offset[2] != 0 - and self.iteration_order == nodes.IterationOrder.PARALLEL ): - raise GTScriptSyntaxError( - message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", - loc=nodes.Location.from_ast_node(t), - ) + if (self.iteration_order == nodes.IterationOrder.PARALLEL): + raise GTScriptSyntaxError( + message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", + loc=nodes.Location.from_ast_node(t)) + if self.backend_name == "gt:gpu": + import cupy as cp + if cp.cuda.get_local_runtime_version() < 12000: + raise GTScriptSyntaxError( + message="Assignment to non-zero offsets in K is not available in gt:gpu for CUDA<12. Please update CUDA.", + loc=nodes.Location.from_ast_node(t)) if not self._is_known(name): if name in self.temp_decls: @@ -2064,6 +2072,7 @@ def run(self): fields=fields_decls, parameters=parameter_decls, local_symbols={}, # Not used + backend_name=self.options.backend_opts["backend_name"], domain=domain, temp_decls=temp_decls, dtypes=self.dtypes, diff --git a/src/gt4py/cartesian/stencil_builder.py b/src/gt4py/cartesian/stencil_builder.py index 07d58f25f5..3fa32d48f4 100644 --- a/src/gt4py/cartesian/stencil_builder.py +++ b/src/gt4py/cartesian/stencil_builder.py @@ -274,6 +274,7 @@ def pkg_path(self) -> pathlib.Path: @property def gtir_pipeline(self) -> GtirPipeline: + self.options.backend_opts["backend_name"] = self.backend.name return self._build_data.get("gtir_pipeline") or self._build_data.setdefault( "gtir_pipeline", GtirPipeline( diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 86bd00874d..cebd55c5c1 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -655,10 +655,16 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): assert (B[:, :, :] == A[:, :, :]).all() -@pytest.mark.parametrize("backend", ALL_BACKENDS) +@pytest.mark.parametrize("backend", ["gt:gpu"]) def test_K_offset_write_conditional(backend): if backend == "cuda": pytest.skip("Cuda backend is not capable of K offset write") + if backend == "gt:gpu": + import cupy as cp + if cp.cuda.get_local_runtime_version() < 12000: + pytest.skip( + "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" + ) arraylib = get_array_library(backend) array_shape = (1, 1, 4) diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_ir_maker.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_ir_maker.py index e054b7a715..7857ae3e00 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_ir_maker.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_ir_maker.py @@ -13,7 +13,7 @@ def test_AugAssign(): - ir_maker = IRMaker(None, None, None, domain=None) + ir_maker = IRMaker(None, None, None, None, domain=None) aug_assign = ast.parse("a += 1", feature_version=PYTHON_AST_VERSION).body[0] _, result = ir_maker.visit_AugAssign(aug_assign) From 0a376c0bcf8fee487dbc4d670ef0742ea9dd076a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Sep 2024 14:06:33 -0400 Subject: [PATCH 18/25] Lint --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 14 +++++++------- .../multi_feature_tests/test_code_generation.py | 1 + 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index f9fc47168e..d277afacd1 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1440,20 +1440,20 @@ def visit_Assign(self, node: ast.Assign) -> list: loc=nodes.Location.from_ast_node(t), ) # Case of K-offset - if ( - len(spatial_offset) == 3 - and spatial_offset[2] != 0 - ): - if (self.iteration_order == nodes.IterationOrder.PARALLEL): + if len(spatial_offset) == 3 and spatial_offset[2] != 0: + if self.iteration_order == nodes.IterationOrder.PARALLEL: raise GTScriptSyntaxError( message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", - loc=nodes.Location.from_ast_node(t)) + loc=nodes.Location.from_ast_node(t), + ) if self.backend_name == "gt:gpu": import cupy as cp + if cp.cuda.get_local_runtime_version() < 12000: raise GTScriptSyntaxError( message="Assignment to non-zero offsets in K is not available in gt:gpu for CUDA<12. Please update CUDA.", - loc=nodes.Location.from_ast_node(t)) + loc=nodes.Location.from_ast_node(t), + ) if not self._is_known(name): if name in self.temp_decls: diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index cebd55c5c1..108d27a632 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -661,6 +661,7 @@ def test_K_offset_write_conditional(backend): pytest.skip("Cuda backend is not capable of K offset write") if backend == "gt:gpu": import cupy as cp + if cp.cuda.get_local_runtime_version() < 12000: pytest.skip( "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" From a07ea46112af8eac033b6712bfa070958a1388ed Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Sep 2024 14:37:03 -0400 Subject: [PATCH 19/25] Change cupy call to lower level Adapt backend_name to `run` on frontend --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 10 +++++----- src/gt4py/cartesian/stencil_builder.py | 5 +++-- .../multi_feature_tests/test_code_generation.py | 2 +- .../frontend_tests/test_gtscript_frontend.py | 2 +- 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index d277afacd1..fca48e3e80 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1449,7 +1449,7 @@ def visit_Assign(self, node: ast.Assign) -> list: if self.backend_name == "gt:gpu": import cupy as cp - if cp.cuda.get_local_runtime_version() < 12000: + if cp.cuda.runtime.runtimeGetVersion() < 12000: raise GTScriptSyntaxError( message="Assignment to non-zero offsets in K is not available in gt:gpu for CUDA<12. Please update CUDA.", loc=nodes.Location.from_ast_node(t), @@ -2014,7 +2014,7 @@ def extract_arg_descriptors(self): return api_signature, fields_decls, parameter_decls - def run(self): + def run(self, backend_name: str): assert ( isinstance(self.ast_root, ast.Module) and "body" in self.ast_root._fields @@ -2072,7 +2072,7 @@ def run(self): fields=fields_decls, parameters=parameter_decls, local_symbols={}, # Not used - backend_name=self.options.backend_opts["backend_name"], + backend_name=backend_name, domain=domain, temp_decls=temp_decls, dtypes=self.dtypes, @@ -2128,14 +2128,14 @@ def prepare_stencil_definition(cls, definition, externals): return GTScriptParser.annotate_definition(definition, externals) @classmethod - def generate(cls, definition, externals, dtypes, options): + def generate(cls, definition, externals, dtypes, options, backend_name): if options.build_info is not None: start_time = time.perf_counter() if not hasattr(definition, "_gtscript_"): cls.prepare_stencil_definition(definition, externals) translator = GTScriptParser(definition, externals=externals, dtypes=dtypes, options=options) - definition_ir = translator.run() + definition_ir = translator.run(backend_name) # GTIR only supports LatLonGrids if definition_ir.domain != nodes.Domain.LatLonGrid(): diff --git a/src/gt4py/cartesian/stencil_builder.py b/src/gt4py/cartesian/stencil_builder.py index 3fa32d48f4..c0f58c0bc9 100644 --- a/src/gt4py/cartesian/stencil_builder.py +++ b/src/gt4py/cartesian/stencil_builder.py @@ -274,11 +274,12 @@ def pkg_path(self) -> pathlib.Path: @property def gtir_pipeline(self) -> GtirPipeline: - self.options.backend_opts["backend_name"] = self.backend.name return self._build_data.get("gtir_pipeline") or self._build_data.setdefault( "gtir_pipeline", GtirPipeline( - self.frontend.generate(self.definition, self.externals, self.dtypes, self.options), + self.frontend.generate( + self.definition, self.externals, self.dtypes, self.options, self.backend.name + ), self.stencil_id, ), ) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 108d27a632..3750a71070 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -662,7 +662,7 @@ def test_K_offset_write_conditional(backend): if backend == "gt:gpu": import cupy as cp - if cp.cuda.get_local_runtime_version() < 12000: + if cp.cuda.runtime.runtimeGetVersion() < 12000: pytest.skip( "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" ) diff --git a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py index 31db760a99..e62f878746 100644 --- a/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py +++ b/tests/cartesian_tests/unit_tests/frontend_tests/test_gtscript_frontend.py @@ -63,7 +63,7 @@ def parse_definition( ) definition_ir = gt_frontend.GTScriptParser( definition_func, externals=externals or {}, options=build_options, dtypes=dtypes - ).run() + ).run("numpy") setattr(definition_func, "__annotations__", original_annotations) From d5f3c25d5ac118c0844933a39597f88895647ce6 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Sep 2024 14:53:05 -0400 Subject: [PATCH 20/25] Missed update to base class, skip test for gt:gpu --- src/gt4py/cartesian/frontend/base.py | 1 + .../multi_feature_tests/test_code_generation.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/src/gt4py/cartesian/frontend/base.py b/src/gt4py/cartesian/frontend/base.py index 3ba54f3356..5e542cd36d 100644 --- a/src/gt4py/cartesian/frontend/base.py +++ b/src/gt4py/cartesian/frontend/base.py @@ -74,6 +74,7 @@ def generate( externals: Dict[str, Any], dtypes: Dict[Type, Type], options: BuildOptions, + backend_name: str, ) -> gtir.Stencil: """ Generate a StencilDefinition from a stencil Python function. diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 3750a71070..2128b3c742 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -588,7 +588,14 @@ def test_K_offset_write(backend): # Cuda generates bad code for the K offset if backend == "cuda": pytest.skip("cuda K-offset write generates bad code") + if backend == "gt:gpu": + import cupy as cp + if cp.cuda.runtime.runtimeGetVersion() < 12000: + pytest.skip( + "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" + ) + arraylib = get_array_library(backend) array_shape = (1, 1, 4) K_values = arraylib.arange(start=40, stop=44) From 0f15c86df7e7cabc94cb38c80119696866e75cd5 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Sep 2024 14:57:41 -0400 Subject: [PATCH 21/25] lint --- .../multi_feature_tests/test_code_generation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 2128b3c742..9a00876009 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -595,7 +595,7 @@ def test_K_offset_write(backend): pytest.skip( "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" ) - + arraylib = get_array_library(backend) array_shape = (1, 1, 4) K_values = arraylib.arange(start=40, stop=44) @@ -662,7 +662,7 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): assert (B[:, :, :] == A[:, :, :]).all() -@pytest.mark.parametrize("backend", ["gt:gpu"]) +@pytest.mark.parametrize("backend", ALL_BACKENDS) def test_K_offset_write_conditional(backend): if backend == "cuda": pytest.skip("Cuda backend is not capable of K offset write") From ffab0d2968589e5af0dc93d60ee96ebd30bcb045 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Sep 2024 12:02:36 -0400 Subject: [PATCH 22/25] GCC 10.3 + NVCC brings a bug to DaCe. Side step it. --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 4 ++-- .../multi_feature_tests/test_code_generation.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index fca48e3e80..fd3c350563 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1446,12 +1446,12 @@ def visit_Assign(self, node: ast.Assign) -> list: message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", loc=nodes.Location.from_ast_node(t), ) - if self.backend_name == "gt:gpu": + if self.backend_name == "gt:gpu" or self.backend_name == "dace:gpu": import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: raise GTScriptSyntaxError( - message="Assignment to non-zero offsets in K is not available in gt:gpu for CUDA<12. Please update CUDA.", + message=f"Assignment to non-zero offsets in K is not available in {self.backend_name} for CUDA<12. Please update CUDA.", loc=nodes.Location.from_ast_node(t), ) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 9a00876009..4952dc14df 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -666,12 +666,12 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): def test_K_offset_write_conditional(backend): if backend == "cuda": pytest.skip("Cuda backend is not capable of K offset write") - if backend == "gt:gpu": + if backend == "gt:gpu" or "dace:gpu": import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: pytest.skip( - "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" + f"{backend} backend with CUDA 11 and/or GCC 10.3 is not capable of K offset write, update CUDA/GCC if possible" ) arraylib = get_array_library(backend) From c15cb13edaf2a2ee52f7e3ca5660dcdab129d191 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Sep 2024 12:04:17 -0400 Subject: [PATCH 23/25] Restore previous variable offset. tracking and add write in K tracking --- .../gtc/dace/expansion/daceir_builder.py | 65 ++++++++++--------- 1 file changed, 34 insertions(+), 31 deletions(-) diff --git a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py index 33362ea077..a8a3a3cb54 100644 --- a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py +++ b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py @@ -320,34 +320,35 @@ def visit_FieldAccess( *, is_target: bool, targets: Set[eve.SymbolRef], - offset_in_K_fields: Set[eve.SymbolRef], + var_offset_fields: Set[eve.SymbolRef], + K_write_with_offset: Set[eve.SymbolRef], **kwargs: Any, ) -> Union[dcir.IndexAccess, dcir.ScalarAccess]: - """Generate the relevant accessor to match the memlet that was previously setup + """Generate the relevant accessor to match the memlet that was previously setup. - This function was written when offset writing was forbidden. It has been refactor - to allow offset write in K with minimum changes. + When a Field is written in K, we force the usage of the OUT memlet throughout the stencil + to make sure all side effects are being properly resolved. Frontend checks ensure that no + parallel code issues sips here. """ res: Union[dcir.IndexAccess, dcir.ScalarAccess] - if node.name in offset_in_K_fields: - is_target = is_target or node.name in targets + if node.name in var_offset_fields.union(K_write_with_offset): + # If write in K, we consider the variable to always be a target + is_target = is_target or node.name in targets or node.name in K_write_with_offset name = get_tasklet_symbol(node.name, node.offset, is_target=is_target) - if is_target: - res = dcir.IndexAccess( - name=name, - offset=self.visit( - node.offset, - is_target=is_target, - targets=targets, - offset_in_K_fields=offset_in_K_fields, - **kwargs, - ), - data_index=node.data_index, - dtype=node.dtype, - ) - else: - res = dcir.ScalarAccess(name=name, dtype=node.dtype) + res = dcir.IndexAccess( + name=name, + offset=self.visit( + node.offset, + is_target=is_target, + targets=targets, + var_offset_fields=var_offset_fields, + K_write_with_offset=K_write_with_offset, + **kwargs, + ), + data_index=node.data_index, + dtype=node.dtype, + ) else: is_target = is_target or ( node.name in targets and node.offset == common.CartesianOffset.zero() @@ -810,21 +811,22 @@ def visit_VerticalLoop( ) ) - # Offsets in K can be both: - # - read indexed via an expression, - # - write offset in K (both in expression and on scalar) - # We get all indexed via expression - offset_in_K_fields = { + # Variable offsets + var_offset_fields = { acc.name for acc in node.walk_values().if_isinstance(oir.FieldAccess) if isinstance(acc.offset, oir.VariableKOffset) } - # We add write offset to K + + # We book keep - all write offset to K + K_write_with_offset = set() for assign_node in node.walk_values().if_isinstance(oir.AssignStmt): if isinstance(assign_node.left, oir.FieldAccess): - acc = assign_node.left - if isinstance(acc.offset, common.CartesianOffset) and acc.offset.k != 0: - offset_in_K_fields.add(acc.name) + if ( + isinstance(assign_node.left.offset, common.CartesianOffset) + and assign_node.left.offset.k != 0 + ): + K_write_with_offset.add(assign_node.left.name) sections_idx = next( idx @@ -842,7 +844,8 @@ def visit_VerticalLoop( global_ctx=global_ctx, iteration_ctx=iteration_ctx, symbol_collector=symbol_collector, - offset_in_K_fields=offset_in_K_fields, + var_offset_fields=var_offset_fields, + K_write_with_offset=K_write_with_offset, **kwargs, ) ) From 8e9ffe8b33c425056deb13779ce5bca74c0c82c3 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Sep 2024 12:26:56 -0400 Subject: [PATCH 24/25] Deactivate K write test for CUDA11+dace --- .../multi_feature_tests/test_code_generation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 4952dc14df..8ae1b99176 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -588,12 +588,12 @@ def test_K_offset_write(backend): # Cuda generates bad code for the K offset if backend == "cuda": pytest.skip("cuda K-offset write generates bad code") - if backend == "gt:gpu": + if backend == "gt:gpu" or "dace:gpu": import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: pytest.skip( - "gt:gpu backend with CUDA 11 is not capable of K offset write, update CUDA if possible" + f"{backend} backend with CUDA 11 and/or GCC 10.3 is not capable of K offset write, update CUDA/GCC if possible" ) arraylib = get_array_library(backend) From c6b30fdff61e3dd1a3b062e43c3dba42ea30529a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Sep 2024 12:58:54 -0400 Subject: [PATCH 25/25] Fix backend checks --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 2 +- .../multi_feature_tests/test_code_generation.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index fd3c350563..962d175eb1 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1446,7 +1446,7 @@ def visit_Assign(self, node: ast.Assign) -> list: message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.", loc=nodes.Location.from_ast_node(t), ) - if self.backend_name == "gt:gpu" or self.backend_name == "dace:gpu": + if self.backend_name in ["gt:gpu", "dace:gpu"]: import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py index 8ae1b99176..976f9a89af 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py @@ -588,7 +588,7 @@ def test_K_offset_write(backend): # Cuda generates bad code for the K offset if backend == "cuda": pytest.skip("cuda K-offset write generates bad code") - if backend == "gt:gpu" or "dace:gpu": + if backend in ["gt:gpu", "dace:gpu"]: import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: @@ -666,7 +666,7 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64): def test_K_offset_write_conditional(backend): if backend == "cuda": pytest.skip("Cuda backend is not capable of K offset write") - if backend == "gt:gpu" or "dace:gpu": + if backend in ["gt:gpu", "dace:gpu"]: import cupy as cp if cp.cuda.runtime.runtimeGetVersion() < 12000: