Error out cleanly for gt:gpu and CUDA < 12

GridTools · Sep 19, 2024 · 8fb7452 · 8fb7452
1 parent 2846713
commit 8fb7452
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)