refactor backend selection

Author: havogt

docs/development/ADRs/0015-Test_Exclusion_Matrices.md

@@ -1,13 +1,13 @@
 ---
-tags: []
+tags: [testing]
 ---
 
 # Test-Exclusion Matrices
 
 - **Status**: valid
 - **Authors**: Edoardo Paone (@edopao), Enrique G. Paredes (@egparedes)
 - **Created**: 2023-09-21
-- **Updated**: 2023-09-21
+- **Updated**: 2024-01-25
 
 In the context of Field View testing, lacking support for specific ITIR features while a certain backend
 is being developed, we decided to use `pytest` fixtures to exclude unsupported tests.
@@ -33,7 +33,7 @@ In the example below, `test_offset_field` requires the backend to support dynami
 def test_offset_field(cartesian_case):
 ```
 
-In order to selectively enable the backends, the dictionary `next_tests.exclusion_matrices.BACKEND_SKIP_TEST_MATRIX`
+In order to selectively enable the backends, the dictionary `next_tests.definitions.BACKEND_SKIP_TEST_MATRIX`
 lists for each backend the features that are not supported.
 The fixture will check if the annotated feature is present in the exclusion-matrix for the selected backend.
 If so, the exclusion matrix will also specify the action `pytest` should take (e.g. `SKIP` or `XFAIL`).

src/gt4py/next/allocators.py

@@ -231,6 +231,7 @@ def __init__(self) -> None:
 
 device_allocators[core_defs.DeviceType.CPU] = StandardCPUFieldBufferAllocator()
 
+
 assert is_field_allocator(device_allocators[core_defs.DeviceType.CPU])
 
 
@@ -349,3 +350,15 @@ def allocate(
         device_id=device.device_id,
         aligned_index=aligned_index,
     )
+
+
+def __getattr__(name: str) -> Any:
+    if name == "default_cpu_allocator":
+        return device_allocators[core_defs.DeviceType.CPU]
+    if name == "default_gpu_allocator":
+        if CUPY_DEVICE is core_defs.DeviceType.CUDA:
+            return device_allocators[core_defs.DeviceType.CUDA]
+        elif CUPY_DEVICE is core_defs.DeviceType.ROCM:
+            return device_allocators[core_defs.DeviceType.ROCM]
+        # else: fall through
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tests/next_tests/definitions.py

@@ -19,8 +19,6 @@
 
 import pytest
 
-from gt4py.next import allocators as next_allocators
-
 
 # Skip definitions
 XFAIL = pytest.xfail
@@ -46,11 +44,6 @@ def short_id(self, num_components: int = 2) -> str:
         return ".".join(self.value.split(".")[-num_components:])
 
 
-class _PythonObjectIdMixinForAllocator(_PythonObjectIdMixin):
-    def short_id(self, num_components: int = 1) -> str:
-        return "None-" + super().short_id(num_components)
-
-
 class ProgramBackendId(_PythonObjectIdMixin, str, enum.Enum):
     GTFN_CPU = "gt4py.next.program_processors.runners.gtfn.run_gtfn"
     GTFN_CPU_IMPERATIVE = "gt4py.next.program_processors.runners.gtfn.run_gtfn_imperative"
@@ -62,13 +55,9 @@ class ProgramBackendId(_PythonObjectIdMixin, str, enum.Enum):
     DOUBLE_ROUNDTRIP = "gt4py.next.program_processors.runners.double_roundtrip.backend"
 
 
-cpu_allocator = next_allocators.StandardCPUFieldBufferAllocator()
-gpu_allocator = next_allocators.StandardGPUFieldBufferAllocator()
-
-
-class AllocatorId(_PythonObjectIdMixinForAllocator, str, enum.Enum):
-    CPU_ALLOCATOR = "next_tests.definitions.cpu_allocator"
-    GPU_ALLOCATOR = "next_tests.definitions.gpu_allocator"
+class AllocatorId(_PythonObjectIdMixin, str, enum.Enum):
+    CPU_ALLOCATOR = "gt4py.next.allocators.default_cpu_allocator"
+    GPU_ALLOCATOR = "gt4py.next.allocators.default_gpu_allocator"
 
 
 class OptionalProgramBackendId(_PythonObjectIdMixin, str, enum.Enum):

tests/next_tests/integration_tests/cases.py

@@ -482,18 +482,18 @@ def verify_with_default_data(
 @pytest.fixture
 def cartesian_case(fieldview_backend):  # noqa: F811 # fixtures
     yield Case(
-        fieldview_backend if isinstance(fieldview_backend, ppi.ProgramExecutor) else None,
+        fieldview_backend.executor,
         offset_provider={"Ioff": IDim, "Joff": JDim, "Koff": KDim},
         default_sizes={IDim: 10, JDim: 10, KDim: 10},
         grid_type=common.GridType.CARTESIAN,
-        allocator=fieldview_backend,
+        allocator=fieldview_backend.allocator,
     )
 
 
 @pytest.fixture
 def unstructured_case(reduction_setup, fieldview_backend):  # noqa: F811 # fixtures
     yield Case(
-        fieldview_backend if isinstance(fieldview_backend, ppi.ProgramExecutor) else None,
+        fieldview_backend.executor,
         offset_provider=reduction_setup.offset_provider,
         default_sizes={
             Vertex: reduction_setup.num_vertices,
@@ -502,7 +502,7 @@ def unstructured_case(reduction_setup, fieldview_backend):  # noqa: F811 # fixtu
             KDim: reduction_setup.k_levels,
         },
         grid_type=common.GridType.UNSTRUCTURED,
-        allocator=fieldview_backend,
+        allocator=fieldview_backend.allocator,
     )
 
 

tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py

@@ -14,7 +14,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from collections import namedtuple
-from typing import Any, TypeVar
+from typing import Any, Optional, TypeVar
 
 import numpy as np
 import pytest
@@ -52,6 +52,23 @@ def no_backend(program: itir.FencilDefinition, *args: Any, **kwargs: Any) -> Non
         )
     ),
 
+_ExecutorAndAllocator = namedtuple("_ExecutorAndAllocator", ["executor", "allocator"])
+
+
+def _backend_name(
+    p: next_tests.definitions._PythonObjectIdMixin
+    | tuple[
+        Optional[next_tests.definitions._PythonObjectIdMixin],
+        next_tests.definitions._PythonObjectIdMixin,
+    ]
+):
+    if isinstance(p, tuple):
+        executor = p[0].short_id() if p[0] is not None else "None"
+        allocator = p[1].short_id()
+        return f"{executor}-{allocator}"
+    else:
+        return p.short_id()
+
 
 @pytest.fixture(
     params=[
@@ -63,13 +80,13 @@ def no_backend(program: itir.FencilDefinition, *args: Any, **kwargs: Any) -> Non
             next_tests.definitions.ProgramBackendId.GTFN_GPU, marks=pytest.mark.requires_gpu
         ),
         # will use the default (embedded) execution, but input/output allocated with the provided allocator
-        next_tests.definitions.AllocatorId.CPU_ALLOCATOR,
+        (None, next_tests.definitions.AllocatorId.CPU_ALLOCATOR),
         pytest.param(
-            next_tests.definitions.AllocatorId.GPU_ALLOCATOR, marks=pytest.mark.requires_gpu
+            (None, next_tests.definitions.AllocatorId.GPU_ALLOCATOR), marks=pytest.mark.requires_gpu
         ),
     ]
     + OPTIONAL_PROCESSORS,
-    ids=lambda p: p.short_id() if p is not None else "None",
+    ids=_backend_name,
 )
 def fieldview_backend(request):
     """
@@ -79,17 +96,27 @@ def fieldview_backend(request):
         Check ADR 15 for details on the test-exclusion matrices.
     """
     backend_or_allocator_id = request.param
-    backend_or_allocator = backend_or_allocator_id.load()
+    if isinstance(backend_or_allocator_id, tuple):
+        backend_id, allocator_id = backend_or_allocator_id
+        assert backend_id is None  # revisit if we have need for backend_id != None
+        executor = None
+        allocator = allocator_id.load()
+        exclusion_matrix_id = backend_or_allocator_id[1]
+    else:
+        backend = backend_or_allocator_id.load()
+        executor = backend.executor
+        allocator = backend
+        exclusion_matrix_id = backend_or_allocator_id
 
     for marker, skip_mark, msg in next_tests.definitions.BACKEND_SKIP_TEST_MATRIX.get(
-        backend_or_allocator_id, []
+        exclusion_matrix_id, []
     ):
         if request.node.get_closest_marker(marker):
-            skip_mark(msg.format(marker=marker, backend=backend_or_allocator_id))
+            skip_mark(msg.format(marker=marker, backend=_backend_name(backend_or_allocator_id)))
 
     backup_backend = decorator.DEFAULT_BACKEND
     decorator.DEFAULT_BACKEND = no_backend
-    yield backend_or_allocator
+    yield _ExecutorAndAllocator(executor, allocator)
     decorator.DEFAULT_BACKEND = backup_backend
 
 

tests/next_tests/integration_tests/feature_tests/test_util_cases.py

@@ -24,6 +24,7 @@
     cartesian_case,
     fieldview_backend,
 )
+from next_tests.integration_tests.feature_tests.ffront_tests import ffront_test_utils
 
 
 @gtx.field_operator
@@ -70,7 +71,13 @@ def test_allocate_const(cartesian_case):  # noqa: F811 # fixtures
     assert b == 42.0
 
 
-@pytest.mark.parametrize("fieldview_backend", [~definitions.ProgramBackendId.ROUNDTRIP])
+_roundtrip_executor_and_allocator = ffront_test_utils._ExecutorAndAllocator(
+    definitions.ProgramBackendId.ROUNDTRIP.load().executor,
+    definitions.ProgramBackendId.ROUNDTRIP.load(),
+)
+
+
+@pytest.mark.parametrize("fieldview_backend", [_roundtrip_executor_and_allocator])
 def test_verify_fails_with_wrong_reference(cartesian_case):  # noqa: F811 # fixtures
     a = cases.allocate(cartesian_case, addition, "a")()
     b = cases.allocate(cartesian_case, addition, "b")()
@@ -81,7 +88,7 @@ def test_verify_fails_with_wrong_reference(cartesian_case):  # noqa: F811 # fixt
         cases.verify(cartesian_case, addition, a, b, out=out, ref=wrong_ref)
 
 
-@pytest.mark.parametrize("fieldview_backend", [~definitions.ProgramBackendId.ROUNDTRIP])
+@pytest.mark.parametrize("fieldview_backend", [_roundtrip_executor_and_allocator])
 def test_verify_fails_with_wrong_type(cartesian_case):  # noqa: F811 # fixtures
     a = cases.allocate(cartesian_case, addition, "a").dtype(np.float32)()
     b = cases.allocate(cartesian_case, addition, "b")()
@@ -91,7 +98,7 @@ def test_verify_fails_with_wrong_type(cartesian_case):  # noqa: F811 # fixtures
         cases.verify(cartesian_case, addition, a, b, out=out, ref=a + b)
 
 
-@pytest.mark.parametrize("fieldview_backend", [~definitions.ProgramBackendId.ROUNDTRIP])
+@pytest.mark.parametrize("fieldview_backend", [_roundtrip_executor_and_allocator])
 def test_verify_with_default_data_fails_with_wrong_reference(
     cartesian_case,  # noqa: F811 # fixtures
 ):

tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py

@@ -196,11 +196,11 @@ def reference(
 @pytest.fixture
 def test_setup(fieldview_backend):
     test_case = cases.Case(
-        fieldview_backend if isinstance(fieldview_backend, ppi.ProgramExecutor) else None,
+        fieldview_backend.executor,
         offset_provider={"Koff": KDim},
         default_sizes={Cell: 14, KDim: 10},
         grid_type=common.GridType.UNSTRUCTURED,
-        allocator=fieldview_backend,
+        allocator=fieldview_backend.allocator,
     )
 
     @dataclasses.dataclass(frozen=True)