Enum for StorageDevice in LayoutInfo

src/gt4py/cartesian/backend/base.py

@@ -84,7 +84,7 @@ class Backend(abc.ABC):
     #: Backend-specific storage parametrization:
     #:
     #:  - "alignment": in bytes
-    #:  - "device": "cpu" | "gpu"
+    #:  - "device": StorageDevice.CPU | StorageDevice.GPU
     #:  - "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]

src/gt4py/cartesian/backend/dace_backend.py

@@ -47,6 +47,7 @@
 from gt4py.cartesian.utils import shash
 from gt4py.eve import codegen
 from gt4py.eve.codegen import MakoTemplate as as_mako
+from gt4py.storage.cartesian.layout import StorageDevice
 
 
 if TYPE_CHECKING:
@@ -408,7 +409,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.storage_info["device"] == StorageDevice.GPU else "cpp"
         sources = {
             "computation": {"computation.hpp": implementation},
             "bindings": {f"bindings.{bindings_ext}": bindings},
@@ -662,7 +663,9 @@ def generate_entry_params(self, stencil_ir: gtir.Stencil, sdfg: dace.SDFG) -> Li
                 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"
+                            "object"
+                            if self.backend.storage_info["device"] == StorageDevice.GPU
+                            else "buffer"
                         ),
                         name=name,
                         ndim=len(data.shape),
@@ -766,7 +769,7 @@ class DaceCPUBackend(BaseDaceBackend):
     languages = {"computation": "c++", "bindings": ["python"]}
     storage_info = {
         "alignment": 1,
-        "device": "cpu",
+        "device": StorageDevice.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))
@@ -788,7 +791,7 @@ class DaceGPUBackend(BaseDaceBackend):
     languages = {"computation": "cuda", "bindings": ["python"]}
     storage_info = {
         "alignment": 32,
-        "device": "gpu",
+        "device": StorageDevice.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))

src/gt4py/cartesian/backend/dace_stencil_object.py

@@ -24,9 +24,10 @@
 from gt4py.cartesian.definitions import AccessKind, DomainInfo, FieldInfo
 from gt4py.cartesian.stencil_object import ArgsInfo, FrozenStencil, StencilObject
 from gt4py.cartesian.utils import shash
+from gt4py.storage.cartesian.layout import StorageDevice
 
 
-def _extract_array_infos(field_args, device) -> Dict[str, Optional[ArgsInfo]]:
+def _extract_array_infos(field_args, device: StorageDevice) -> Dict[str, Optional[ArgsInfo]]:
     return {
         name: ArgsInfo(
             array=arg,

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 StorageDevice
 
 
 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 backend.storage_info["device"] == StorageDevice.GPU 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 StorageDevice
 
 from .gtc_common import BaseGTBackend, CUDAPyExtModuleGenerator
 
@@ -85,7 +86,9 @@ 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 self.backend.storage_info["device"] == StorageDevice.GPU
+                        else "buffer"
                     ),
                     name=node.name,
                     sid_ndim=sid_ndim,

src/gt4py/cartesian/gtc/dace/oir_to_dace.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, Literal
+from typing import Dict
 
 import dace
 import dace.properties
@@ -31,16 +31,17 @@
     AccessCollector,
     compute_horizontal_block_extents,
 )
+from gt4py.storage.cartesian.layout import StorageDevice
 
 
-transient_storage_per_device: Dict[Literal["cpu", "gpu"], dace.StorageType] = {
-    "cpu": dace.StorageType.Default,
-    "gpu": dace.StorageType.GPU_Global,
+transient_storage_per_device: Dict[StorageDevice, dace.StorageType] = {
+    StorageDevice.CPU: dace.StorageType.Default,
+    StorageDevice.GPU: dace.StorageType.GPU_Global,
 }
 
-device_type_per_device: Dict[Literal["cpu", "gpu"], dace.DeviceType] = {
-    "cpu": dace.DeviceType.CPU,
-    "gpu": dace.DeviceType.GPU,
+device_type_per_device: Dict[StorageDevice, dace.DeviceType] = {
+    StorageDevice.CPU: dace.DeviceType.CPU,
+    StorageDevice.GPU: dace.DeviceType.GPU,
 }
 
 
@@ -114,7 +115,7 @@ def visit_VerticalLoop(
         node: oir.VerticalLoop,
         *,
         ctx: OirSDFGBuilder.SDFGContext,
-        device: Literal["cpu", "gpu"],
+        device: StorageDevice,
     ) -> None:
         declarations = {
             acc.name: ctx.decls[acc.name]
@@ -155,7 +156,7 @@ def visit_VerticalLoop(
                 library_node, connector_name, access_node, None, dace.Memlet(field, subset=subset)
             )
 
-    def visit_Stencil(self, node: oir.Stencil, *, device: Literal["cpu", "gpu"]) -> dace.SDFG:
+    def visit_Stencil(self, node: oir.Stencil, *, device: StorageDevice) -> dace.SDFG:
         ctx = OirSDFGBuilder.SDFGContext(node)
         for param in node.params:
             if isinstance(param, oir.FieldDecl):

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 StorageDevice
 
 
 try:
@@ -51,15 +52,15 @@ def _compute_domain_origin_cache_key(
 
 @dataclass
 class ArgsInfo:
-    device: str
+    device: StorageDevice
     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]], device: StorageDevice
 ) -> Dict[str, Optional[ArgsInfo]]:
     array_infos: Dict[str, Optional[ArgsInfo]] = {}
     for name, arg in field_args.items():

src/gt4py/cartesian/testing/suites.py

@@ -24,6 +24,7 @@
 from gt4py.cartesian.gtc.definitions import Boundary, CartesianSpace, Index, Shape
 from gt4py.cartesian.stencil_object import StencilObject
 from gt4py.storage.cartesian import utils as storage_utils
+from gt4py.storage.cartesian.layout import StorageDevice
 
 from .input_strategies import (
     SymbolKind,
@@ -206,7 +207,10 @@ def hyp_wrapper(test_hyp, hypothesis_data):
                 )
 
                 marks = test["marks"].copy()
-                if gt4pyc.backend.from_name(test["backend"]).storage_info["device"] == "gpu":
+                if (
+                    gt4pyc.backend.from_name(test["backend"]).storage_info["device"]
+                    == StorageDevice.GPU
+                ):
                     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 +244,10 @@ def hyp_wrapper(test_hyp, hypothesis_data):
                 )
 
                 marks = test["marks"].copy()
-                if gt4pyc.backend.from_name(test["backend"]).storage_info["device"] == "gpu":
+                if (
+                    gt4pyc.backend.from_name(test["backend"]).storage_info["device"]
+                    == StorageDevice.GPU
+                ):
                     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/interface.py

@@ -81,7 +81,7 @@ 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"] == layout.StorageDevice.GPU:
         allocate_f = storage_utils.allocate_gpu
     else:
         allocate_f = storage_utils.allocate_cpu

src/gt4py/storage/cartesian/layout.py

@@ -6,13 +6,13 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
     Dict,
     Final,
-    Literal,
     Optional,
     Sequence,
     Tuple,
@@ -30,9 +30,14 @@
         cp = None
 
 
+class StorageDevice(Enum):
+    CPU = 1
+    GPU = 2
+
+
 class LayoutInfo(TypedDict):
     alignment: int  # measured in bytes
-    device: Literal["cpu", "gpu"]
+    device: StorageDevice
     layout_map: Callable[[Tuple[str, ...]], Tuple[Optional[int], ...]]
     is_optimal_layout: Callable[[Any, Tuple[str, ...]], bool]
 
@@ -136,15 +141,15 @@ def make_cuda_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ..
 
 NaiveCPULayout: Final[LayoutInfo] = {
     "alignment": 1,
-    "device": "cpu",
+    "device": StorageDevice.CPU,
     "layout_map": lambda axes: tuple(i for i in range(len(axes))),
     "is_optimal_layout": lambda *_: True,
 }
 register("naive_cpu", NaiveCPULayout)
 
 CPUIFirstLayout: Final[LayoutInfo] = {
     "alignment": 8,
-    "device": "cpu",
+    "device": StorageDevice.CPU,
     "layout_map": make_gtcpu_ifirst_layout_map,
     "is_optimal_layout": layout_checker_factory(make_gtcpu_ifirst_layout_map),
 }
@@ -153,7 +158,7 @@ def make_cuda_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ..
 
 CPUKFirstLayout: Final[LayoutInfo] = {
     "alignment": 1,
-    "device": "cpu",
+    "device": StorageDevice.CPU,
     "layout_map": make_gtcpu_kfirst_layout_map,
     "is_optimal_layout": layout_checker_factory(make_gtcpu_kfirst_layout_map),
 }
@@ -162,7 +167,7 @@ def make_cuda_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ..
 
 CUDALayout: Final[LayoutInfo] = {
     "alignment": 32,
-    "device": "gpu",
+    "device": StorageDevice.GPU,
     "layout_map": make_cuda_layout_map,
     "is_optimal_layout": layout_checker_factory(make_cuda_layout_map),
 }

src/gt4py/storage/cartesian/utils.py

@@ -22,6 +22,7 @@
 from gt4py.cartesian import config as gt_config
 from gt4py.eve.extended_typing import ArrayInterface, CUDAArrayInterface
 from gt4py.storage import allocators
+from gt4py.storage.cartesian.layout import StorageDevice
 
 
 try:
@@ -182,18 +183,16 @@ def cpu_copy(array: Union[np.ndarray, "cp.ndarray"]) -> np.ndarray:
         return np.array(array)
 
 
-def asarray(
-    array: FieldLike, *, device: Literal["cpu", "gpu", None] = None
-) -> np.ndarray | cp.ndarray:
+def asarray(array: FieldLike, *, device: Optional[StorageDevice] = None) -> np.ndarray | cp.ndarray:
     if hasattr(array, "ndarray"):
         # extract the buffer from a gt4py.next.Field
         # TODO(havogt): probably `Field` should provide the array interface methods when applicable
         array = array.ndarray
 
     xp = None
-    if device == "cpu":
+    if device == StorageDevice.CPU:
         xp = np
-    elif device == "gpu":
+    elif device == StorageDevice.GPU:
         assert cp is not None
         xp = cp
     elif not device:

tests/cartesian_tests/definitions.py

@@ -20,11 +20,12 @@
 
 from gt4py import cartesian as gt4pyc
 from gt4py.cartesian import utils as gt_utils
+from gt4py.storage.cartesian.layout import StorageDevice
 
 
 def _backend_name_as_param(name):
     marks = []
-    if gt4pyc.backend.from_name(name).storage_info["device"] == "gpu":
+    if gt4pyc.backend.from_name(name).storage_info["device"] == StorageDevice.GPU:
         marks.append(pytest.mark.requires_gpu)
     if "dace" in name:
         marks.append(pytest.mark.requires_dace)
@@ -34,18 +35,18 @@ def _backend_name_as_param(name):
 _ALL_BACKEND_NAMES = list(gt4pyc.backend.REGISTRY.keys())
 
 
-def _get_backends_with_storage_info(storage_info_kind: str):
+def _get_backends_with_storage_info(storage_device: StorageDevice):
     res = []
     for name in _ALL_BACKEND_NAMES:
         backend = gt4pyc.backend.from_name(name)
         if not getattr(backend, "disabled", False):
-            if backend.storage_info["device"] == storage_info_kind:
+            if backend.storage_info["device"] == storage_device:
                 res.append(_backend_name_as_param(name))
     return res
 
 
-CPU_BACKENDS = _get_backends_with_storage_info("cpu")
-GPU_BACKENDS = _get_backends_with_storage_info("gpu")
+CPU_BACKENDS = _get_backends_with_storage_info(StorageDevice.CPU)
+GPU_BACKENDS = _get_backends_with_storage_info(StorageDevice.GPU)
 ALL_BACKENDS = CPU_BACKENDS + GPU_BACKENDS
 
 _PERFORMANCE_BACKEND_NAMES = [name for name in _ALL_BACKEND_NAMES if name not in ("numpy", "cuda")]
@@ -61,7 +62,7 @@ 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":
+    if backend_cls.storage_info["device"] == StorageDevice.GPU:
         assert cp is not None
         return cp
     else:

tests/storage_tests/unit_tests/test_interface.py

@@ -11,6 +11,8 @@
 import numpy as np
 import pytest
 
+from gt4py.storage.cartesian.layout import StorageDevice
+
 
 try:
     import cupy as cp
@@ -23,12 +25,14 @@
 
 
 CPU_LAYOUTS = [
-    name for name, info in gt4py.storage.layout.REGISTRY.items() if info["device"] == "cpu"
+    name
+    for name, info in gt4py.storage.layout.REGISTRY.items()
+    if info["device"] == StorageDevice.CPU
 ]
 GPU_LAYOUTS = [
     pytest.param(name, marks=pytest.mark.requires_gpu)
     for name, info in gt4py.storage.layout.REGISTRY.items()
-    if info["device"] == "gpu"
+    if info["device"] == StorageDevice.GPU
 ]
 
 try: