diff --git a/src/gt4py/cartesian/backend/base.py b/src/gt4py/cartesian/backend/base.py index 5bab0453a9..b25095ffca 100644 --- a/src/gt4py/cartesian/backend/base.py +++ b/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] diff --git a/src/gt4py/cartesian/backend/dace_backend.py b/src/gt4py/cartesian/backend/dace_backend.py index c84e1061c4..b1b66890fb 100644 --- a/src/gt4py/cartesian/backend/dace_backend.py +++ b/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 {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)) diff --git a/src/gt4py/cartesian/backend/dace_stencil_object.py b/src/gt4py/cartesian/backend/dace_stencil_object.py index 21006475a0..41a41980fc 100644 --- a/src/gt4py/cartesian/backend/dace_stencil_object.py +++ b/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, diff --git a/src/gt4py/cartesian/backend/gtc_common.py b/src/gt4py/cartesian/backend/gtc_common.py index 348e85de92..2b1dd86ab1 100644 --- a/src/gt4py/cartesian/backend/gtc_common.py +++ b/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>({name})""".format( diff --git a/src/gt4py/cartesian/backend/gtcpp_backend.py b/src/gt4py/cartesian/backend/gtcpp_backend.py index 96f5672ae4..bf0eb41caf 100644 --- a/src/gt4py/cartesian/backend/gtcpp_backend.py +++ b/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 {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, diff --git a/src/gt4py/cartesian/gtc/dace/oir_to_dace.py b/src/gt4py/cartesian/gtc/dace/oir_to_dace.py index 8bbca3b590..60ffcae93d 100644 --- a/src/gt4py/cartesian/gtc/dace/oir_to_dace.py +++ b/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): diff --git a/src/gt4py/cartesian/stencil_object.py b/src/gt4py/cartesian/stencil_object.py index 5e5976e3e5..cbfcabdc9a 100644 --- a/src/gt4py/cartesian/stencil_object.py +++ b/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,7 +52,7 @@ 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 @@ -59,7 +60,7 @@ class ArgsInfo: 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(): diff --git a/src/gt4py/cartesian/testing/suites.py b/src/gt4py/cartesian/testing/suites.py index 423f834f51..47e26444f3 100644 --- a/src/gt4py/cartesian/testing/suites.py +++ b/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. diff --git a/src/gt4py/storage/cartesian/interface.py b/src/gt4py/storage/cartesian/interface.py index 8b38bcdd42..43034f2117 100644 --- a/src/gt4py/storage/cartesian/interface.py +++ b/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 diff --git a/src/gt4py/storage/cartesian/layout.py b/src/gt4py/storage/cartesian/layout.py index 868f7af7c4..521fc2fa73 100644 --- a/src/gt4py/storage/cartesian/layout.py +++ b/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,7 +141,7 @@ 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, } @@ -144,7 +149,7 @@ def make_cuda_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], .. 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), } diff --git a/src/gt4py/storage/cartesian/utils.py b/src/gt4py/storage/cartesian/utils.py index bd89c85052..a59f1543cb 100644 --- a/src/gt4py/storage/cartesian/utils.py +++ b/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: diff --git a/tests/cartesian_tests/definitions.py b/tests/cartesian_tests/definitions.py index 7499ad4a95..7986065902 100644 --- a/tests/cartesian_tests/definitions.py +++ b/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: diff --git a/tests/storage_tests/unit_tests/test_interface.py b/tests/storage_tests/unit_tests/test_interface.py index ba7bc2aaef..6091ef0805 100644 --- a/tests/storage_tests/unit_tests/test_interface.py +++ b/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: