From b0a430f372d6cb9a501f76edde28b791584e3cc2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 8 Oct 2024 15:09:35 -0400 Subject: [PATCH 1/2] Protect K write against dynamic access that make extent analysis impossible --- src/gt4py/cartesian/frontend/gtscript_frontend.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index e2aa98f3cf..87012df980 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -1458,6 +1458,11 @@ def visit_Assign(self, node: ast.Assign) -> list: 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), ) + if isinstance(spatial_offset[2], nodes.Expr): + raise GTScriptSyntaxError( + message=f"Assignment to non-zero offsets in K can only be done with literals, got {spatial_offset[2]}", + loc=nodes.Location.from_ast_node(t), + ) if not self._is_known(name): if name in self.temp_decls: From 6c734bc362bf580a3a03a8f9f767f01d16bc8b56 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 8 Oct 2024 15:21:00 -0400 Subject: [PATCH 2/2] utest --- .../frontend_tests/test_gtscript_frontend.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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 e62f878746..d108cb88af 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 @@ -1429,6 +1429,21 @@ def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float module=self.__class__.__name__, ) + with pytest.raises( + gt_frontend.GTScriptSyntaxError, + match=r"Assignment to non-zero offsets in K can only be done with literals, got(.*)", + ): + + def func( + in_field: gtscript.Field[np.float_], + out_field: gtscript.Field[np.float_], + index: int, + ): + with computation(FORWARD), interval(...): + out_field[0, 0, index] = in_field + + parse_definition(func, name=inspect.stack()[0][3], module=self.__class__.__name__) + def test_return_to_subscript(self): @gtscript.function def func(a):