refactor[cartesian]: use DeviceType from core_defn in LayoutInfo

src/gt4py/cartesian/backend/base.py

@@ -45,9 +45,9 @@
 REGISTRY = gt_utils.Registry()
 
 
-def from_name(name: str) -> Optional[Type[Backend]]:
+def from_name(name: str) -> Type[Backend]:
     backend = REGISTRY.get(name, None)
-    if not backend:
+    if backend is None:
         raise NotImplementedError(
             f"Backend {name} is not implemented, options are: {REGISTRY.names}"
         )
@@ -84,7 +84,7 @@ class Backend(abc.ABC):
     #: Backend-specific storage parametrization:
     #:
     #:  - "alignment": in bytes
-    #:  - "device": "cpu" | "gpu"
+    #:  - "device": core_defs.DeviceType | None
     #:  - "layout_map": callback converting a mask to a layout
     #:  - "is_optimal_layout": callback checking if a storage has compatible layout
     storage_info: ClassVar[gt_storage.layout.LayoutInfo]
@@ -435,21 +435,20 @@ def disabled(message: str, *, enabled_env_var: str) -> Callable[[Type[Backend]],
     enabled = bool(int(os.environ.get(enabled_env_var, "0")))
     if enabled:
         return deprecated(message)
-    else:
 
-        def _decorator(cls: Type[Backend]) -> Type[Backend]:
-            def _no_generate(obj) -> Type[StencilObject]:
-                raise NotImplementedError(
-                    f"Disabled '{cls.name}' backend: 'f{message}'\n",
-                    f"You can still enable the backend by hand using the environment variable '{enabled_env_var}=1'",
-                )
-
-            # Replace generate method with raise
-            if not hasattr(cls, "generate"):
-                raise ValueError(f"Coding error. Expected a generate method on {cls}")
-            # Flag that it got disabled for register lookup
-            cls.disabled = True  # type: ignore
-            cls.generate = _no_generate  # type: ignore
-            return cls
-
-        return _decorator
+    def _decorator(cls: Type[Backend]) -> Type[Backend]:
+        def _no_generate(obj) -> Type[StencilObject]:
+            raise NotImplementedError(
+                f"Disabled '{cls.name}' backend: 'f{message}'\n",
+                f"You can still enable the backend by hand using the environment variable '{enabled_env_var}=1'",
+            )
+
+        # Replace generate method with raise
+        if not hasattr(cls, "generate"):
+            raise ValueError(f"Coding error. Expected a generate method on {cls}")
+        # Flag that it got disabled for register lookup
+        cls.disabled = True  # type: ignore
+        cls.generate = _no_generate  # type: ignore
+        return cls
+
+    return _decorator

src/gt4py/cartesian/backend/dace_backend.py

@@ -421,7 +421,7 @@ def __call__(self, stencil_ir: gtir.Stencil) -> Dict[str, Dict[str, str]]:
             stencil_ir, sdfg, module_name=self.module_name, backend=self.backend
         )
 
-        bindings_ext = "cu" if self.backend.storage_info["device"] == "gpu" else "cpp"
+        bindings_ext = "cu" if self.backend.name.endswith(":gpu") else "cpp"
         sources = {
             "computation": {"computation.hpp": implementation},
             "bindings": {f"bindings.{bindings_ext}": bindings},
@@ -674,9 +674,7 @@ def generate_entry_params(self, stencil_ir: gtir.Stencil, sdfg: dace.SDFG) -> Li
                 assert isinstance(data, dace.data.Array)
                 res[name] = (
                     "py::{pybind_type} {name}, std::array<gt::int_t,{ndim}> {name}_origin".format(
-                        pybind_type=(
-                            "object" if self.backend.storage_info["device"] == "gpu" else "buffer"
-                        ),
+                        pybind_type=("object" if self.backend.name.endswith(":gpu") else "buffer"),
                         name=name,
                         ndim=len(data.shape),
                     )
@@ -777,14 +775,7 @@ def generate(self) -> Type[StencilObject]:
 class DaceCPUBackend(BaseDaceBackend):
     name = "dace:cpu"
     languages = {"computation": "c++", "bindings": ["python"]}
-    storage_info = {
-        "alignment": 1,
-        "device": "cpu",
-        "layout_map": gt_storage.layout.layout_maker_factory((0, 1, 2)),
-        "is_optimal_layout": gt_storage.layout.layout_checker_factory(
-            gt_storage.layout.layout_maker_factory((0, 1, 2))
-        ),
-    }
+    storage_info = gt_storage.layout.DaceCPULayout
     MODULE_GENERATOR_CLASS = DaCePyExtModuleGenerator
 
     options = BaseGTBackend.GT_BACKEND_OPTS
@@ -799,14 +790,7 @@ class DaceGPUBackend(BaseDaceBackend):
 
     name = "dace:gpu"
     languages = {"computation": "cuda", "bindings": ["python"]}
-    storage_info = {
-        "alignment": 32,
-        "device": "gpu",
-        "layout_map": gt_storage.layout.layout_maker_factory((2, 1, 0)),
-        "is_optimal_layout": gt_storage.layout.layout_checker_factory(
-            gt_storage.layout.layout_maker_factory((2, 1, 0))
-        ),
-    }
+    storage_info = gt_storage.layout.DaceGPULayout
     MODULE_GENERATOR_CLASS = DaCeCUDAPyExtModuleGenerator
     options = {**BaseGTBackend.GT_BACKEND_OPTS, "device_sync": {"versioning": True, "type": bool}}
 

src/gt4py/cartesian/backend/dace_stencil_object.py

@@ -26,12 +26,11 @@
 from gt4py.cartesian.utils import shash
 
 
-def _extract_array_infos(field_args, device) -> Dict[str, Optional[ArgsInfo]]:
+def _extract_array_infos(field_args) -> Dict[str, Optional[ArgsInfo]]:
     return {
         name: ArgsInfo(
             array=arg,
             dimensions=getattr(arg, "__gt_dims__", None),
-            device=device,
             origin=getattr(arg, "__gt_origin__", None),
         )
         for name, arg in field_args.items()
@@ -186,9 +185,7 @@ def normalize_args(
         args_as_kwargs = {
             name: (kwargs[name] if name in kwargs else next(args_iter)) for name in arg_names
         }
-        arg_infos = _extract_array_infos(
-            field_args=args_as_kwargs, device=backend_cls.storage_info["device"]
-        )
+        arg_infos = _extract_array_infos(field_args=args_as_kwargs)
 
         origin = DaCeStencilObject._normalize_origins(arg_infos, field_info, origin)
 

src/gt4py/cartesian/backend/gtc_common.py

@@ -23,6 +23,7 @@
 from gt4py.cartesian.gtc.passes.gtir_pipeline import GtirPipeline
 from gt4py.cartesian.gtc.passes.oir_pipeline import OirPipeline
 from gt4py.eve.codegen import MakoTemplate as as_mako
+from gt4py.storage.cartesian.layout import is_gpu_device
 
 
 if TYPE_CHECKING:
@@ -51,7 +52,7 @@ def pybuffer_to_sid(
     domain_ndim = domain_dim_flags.count(True)
     sid_ndim = domain_ndim + data_ndim
 
-    as_sid = "as_cuda_sid" if backend.storage_info["device"] == "gpu" else "as_sid"
+    as_sid = "as_cuda_sid" if is_gpu_device(backend.storage_info) else "as_sid"
 
     sid_def = """gt::{as_sid}<{ctype}, {sid_ndim},
         gt::integral_constant<int, {unique_index}>>({name})""".format(

src/gt4py/cartesian/backend/gtcpp_backend.py

@@ -25,6 +25,7 @@
 from gt4py.cartesian.gtc.passes.gtir_pipeline import GtirPipeline
 from gt4py.cartesian.gtc.passes.oir_pipeline import DefaultPipeline
 from gt4py.eve import codegen
+from gt4py.storage.cartesian.layout import is_gpu_device
 
 from .gtc_common import BaseGTBackend, CUDAPyExtModuleGenerator
 
@@ -85,7 +86,7 @@ def visit_FieldDecl(self, node: gtcpp.FieldDecl, **kwargs):
             if kwargs["external_arg"]:
                 return "py::{pybind_type} {name}, std::array<gt::int_t,{sid_ndim}> {name}_origin".format(
                     pybind_type=(
-                        "object" if self.backend.storage_info["device"] == "gpu" else "buffer"
+                        "object" if is_gpu_device(self.backend.storage_info) else "buffer"
                     ),
                     name=node.name,
                     sid_ndim=sid_ndim,

src/gt4py/cartesian/stencil_object.py

@@ -15,7 +15,7 @@
 from dataclasses import dataclass
 from numbers import Number
 from pickle import dumps
-from typing import Any, Callable, ClassVar, Dict, Literal, Optional, Tuple, Union, cast
+from typing import Any, Callable, ClassVar, Dict, Optional, Tuple, Union, cast
 
 import numpy as np
 
@@ -24,6 +24,7 @@
 from gt4py import cartesian as gt4pyc
 from gt4py.cartesian.definitions import AccessKind, DomainInfo, FieldInfo, ParameterInfo
 from gt4py.cartesian.gtc.definitions import Index, Shape
+from gt4py.storage.cartesian.layout import LayoutInfo
 
 
 try:
@@ -51,22 +52,21 @@ def _compute_domain_origin_cache_key(
 
 @dataclass
 class ArgsInfo:
-    device: str
     array: FieldType
     original_object: Optional[Any] = None
     origin: Optional[Tuple[int, ...]] = None
     dimensions: Optional[Tuple[str, ...]] = None
 
 
 def _extract_array_infos(
-    field_args: Dict[str, Optional[FieldType]], device: Literal["cpu", "gpu"]
+    field_args: Dict[str, Optional[FieldType]], layout_info: LayoutInfo
 ) -> Dict[str, Optional[ArgsInfo]]:
     array_infos: Dict[str, Optional[ArgsInfo]] = {}
     for name, arg in field_args.items():
         if arg is None:
             array_infos[name] = None
         else:
-            array = storage_utils.asarray(arg, device=device)
+            array = storage_utils.asarray(arg, layout_info=layout_info)
             dimensions = storage_utils.get_dims(arg)
             if dimensions is not None:
                 sorted_dimensions = [d for d in "IJK" if d in dimensions]
@@ -79,7 +79,6 @@ def _extract_array_infos(
                 array=array,
                 original_object=arg,
                 dimensions=dimensions,
-                device=device,
                 origin=storage_utils.get_origin(arg),
             )
     return array_infos
@@ -562,8 +561,7 @@ def _call_run(
             exec_info["call_run_start_time"] = time.perf_counter()
         backend_cls = gt4pyc.backend.from_name(self.backend)
         assert backend_cls is not None
-        device = backend_cls.storage_info["device"]
-        array_infos = _extract_array_infos(field_args, device)
+        array_infos = _extract_array_infos(field_args, backend_cls.storage_info)
 
         cache_key = _compute_domain_origin_cache_key(array_infos, parameter_args, domain, origin)
         if cache_key not in self._domain_origin_cache:

src/gt4py/cartesian/testing/definitions.py

@@ -0,0 +1,24 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from gt4py.cartesian.backend.base import REGISTRY as BACKEND_REGISTRY
+
+
+# TODO (romanc)
+# This file can move to the tests/ folder once we refactor the `StencilTestSuite`
+# class (i.e. https://github.com/GEOS-ESM/NDSL/issues/72). The stencil test suites
+# use the `GPU_BACKEND_NAMES` and I didn't wanna have `gt4py/cartesian/testing` depend
+# on the `tests/` directory. That sounded the wrong way around, so I moved them here
+# for now.
+
+ALL_BACKEND_NAMES = list(BACKEND_REGISTRY.keys())
+
+GPU_BACKEND_NAMES = ["cuda", "gt:gpu", "dace:gpu"]
+CPU_BACKEND_NAMES = [name for name in ALL_BACKEND_NAMES if name not in GPU_BACKEND_NAMES]
+
+PERFORMANCE_BACKEND_NAMES = [name for name in ALL_BACKEND_NAMES if name not in ("numpy", "cuda")]

src/gt4py/cartesian/testing/suites.py

@@ -23,6 +23,7 @@
 from gt4py.cartesian.definitions import AccessKind, FieldInfo
 from gt4py.cartesian.gtc.definitions import Boundary, CartesianSpace, Index, Shape
 from gt4py.cartesian.stencil_object import StencilObject
+from gt4py.cartesian.testing.definitions import GPU_BACKEND_NAMES
 from gt4py.storage.cartesian import utils as storage_utils
 
 from .input_strategies import (
@@ -206,7 +207,7 @@ def hyp_wrapper(test_hyp, hypothesis_data):
                 )
 
                 marks = test["marks"].copy()
-                if gt4pyc.backend.from_name(test["backend"]).storage_info["device"] == "gpu":
+                if test["backend"] in GPU_BACKEND_NAMES:
                     marks.append(pytest.mark.requires_gpu)
                 # Run generation and implementation tests in the same group to ensure
                 # (thread-) safe parallelization of stencil tests.
@@ -240,7 +241,7 @@ def hyp_wrapper(test_hyp, hypothesis_data):
                 )
 
                 marks = test["marks"].copy()
-                if gt4pyc.backend.from_name(test["backend"]).storage_info["device"] == "gpu":
+                if test["backend"] in GPU_BACKEND_NAMES:
                     marks.append(pytest.mark.requires_gpu)
                 # Run generation and implementation tests in the same group to ensure
                 # (thread-) safe parallelization of stencil tests.

src/gt4py/storage/cartesian/cupy_device.py

@@ -0,0 +1,24 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from typing import Final, Literal
+
+from gt4py._core.definitions import DeviceType
+
+
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+
+CUPY_DEVICE: Final[Literal[DeviceType.CUDA, DeviceType.ROCM] | None] = (
+    None if not cp else (DeviceType.ROCM if cp.cuda.get_hipcc_path() else DeviceType.CUDA)
+)

src/gt4py/storage/cartesian/interface.py

@@ -81,10 +81,15 @@ def empty(
     _error_on_invalid_preset(backend)
     storage_info = layout.from_name(backend)
     assert storage_info is not None
-    if storage_info["device"] == "gpu":
+
+    if storage_info["device"] is None:
+        raise ValueError("device is None")
+    elif layout.is_gpu_device(storage_info):
         allocate_f = storage_utils.allocate_gpu
-    else:
+    elif layout.is_cpu_device(storage_info):
         allocate_f = storage_utils.allocate_cpu
+    else:
+        raise ValueError("Unknown device")
 
     aligned_index, shape, dtype, dimensions = storage_utils.normalize_storage_spec(
         aligned_index, shape, dtype, dimensions
@@ -322,6 +327,6 @@ def from_array(
 
     layout_info = layout.from_name(backend)
     assert layout_info is not None
-    storage[...] = storage_utils.asarray(data, device=layout_info["device"])
+    storage[...] = storage_utils.asarray(data, layout_info=layout_info)
 
     return storage