Try #868:

src/gt4py/backend/base.py

@@ -72,7 +72,6 @@ class Backend(abc.ABC):
     #:  - "device": "cpu" | "gpu"
     #:  - "layout_map": callback converting a mask to a layout
     #:  - "is_compatible_layout": callback checking if a storage has compatible layout
-    #:  - "is_compatible_type": callback checking if storage has compatible type
     storage_info: ClassVar[Dict[str, Any]]
 
     #: Language support:

src/gt4py/backend/cuda_backend.py

@@ -30,9 +30,8 @@
 
 from .gtc_common import (
     BaseGTBackend,
-    GTCUDAPyModuleGenerator,
+    PyExtModuleGenerator,
     cuda_is_compatible_layout,
-    cuda_is_compatible_type,
     make_cuda_layout_map,
 )
 
@@ -145,10 +144,9 @@ class CudaBackend(BaseGTBackend, CLIBackendMixin):
         "device": "gpu",
         "layout_map": make_cuda_layout_map,
         "is_compatible_layout": cuda_is_compatible_layout,
-        "is_compatible_type": cuda_is_compatible_type,
     }
     PYEXT_GENERATOR_CLASS = CudaExtGenerator  # type: ignore
-    MODULE_GENERATOR_CLASS = GTCUDAPyModuleGenerator
+    MODULE_GENERATOR_CLASS = PyExtModuleGenerator
     GT_BACKEND_T = "gpu"
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:

src/gt4py/backend/dace_backend.py

@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type
 
 import dace
-import numpy as np
+import dace.data
 from dace.serialize import dumps
 
 from eve import codegen
@@ -27,10 +27,8 @@
 from gt4py.backend.gtc_common import (
     BackendCodegen,
     BaseGTBackend,
-    GTCUDAPyModuleGenerator,
     PyExtModuleGenerator,
     bindings_main_template,
-    cuda_is_compatible_type,
     pybuffer_to_sid,
 )
 from gt4py.backend.module_generator import make_args_data_from_gtir
@@ -366,7 +364,6 @@ def unique_index(self) -> int:
 
     def generate_entry_params(self, stencil_ir: gtir.Stencil, sdfg: dace.SDFG) -> List[str]:
         res: Dict[str, str] = {}
-        import dace.data
 
         for name in sdfg.signature_arglist(with_types=False, for_call=True):
             if name in sdfg.arrays:
@@ -383,22 +380,12 @@ def generate_entry_params(self, stencil_ir: gtir.Stencil, sdfg: dace.SDFG) -> Li
 
     def generate_sid_params(self, sdfg: dace.SDFG) -> List[str]:
         res: List[str] = []
-        import dace.data
 
         for name, array in sdfg.arrays.items():
             if array.transient:
                 continue
-            domain_dim_flags = tuple(
-                True
-                if any(
-                    dace.symbolic.pystr_to_symbolic(f"__{dim.upper()}") in s.free_symbols
-                    for s in array.shape
-                    if hasattr(s, "free_symbols")
-                )
-                else False
-                for dim in "ijk"
-            )
-            data_ndim = len(array.shape) - sum(array_dimensions(array))
+            domain_dim_flags = array_dimensions(array)
+            data_ndim = len(array.shape) - sum(domain_dim_flags)
             sid_def = pybuffer_to_sid(
                 name=name,
                 ctype=array.dtype.ctype,
@@ -457,10 +444,6 @@ def generate_class_members(self):
         return res
 
 
-class DaCeCUDAPyExtModuleGenerator(DaCePyExtModuleGenerator, GTCUDAPyModuleGenerator):
-    pass
-
-
 class BaseDaceBackend(BaseGTBackend, CLIBackendMixin):
 
     GT_BACKEND_T = "dace"
@@ -497,8 +480,7 @@ class DaceCPUBackend(BaseDaceBackend):
         "alignment": 1,
         "device": "cpu",
         "layout_map": layout_maker_factory((1, 0, 2)),
-        "is_compatible_layout": lambda x: True,
-        "is_compatible_type": lambda x: isinstance(x, np.ndarray),
+        "is_compatible_layout": lambda x, m: True,
     }
     MODULE_GENERATOR_CLASS = DaCePyExtModuleGenerator
 
@@ -519,9 +501,8 @@ class DaceGPUBackend(BaseDaceBackend):
         "device": "gpu",
         "layout_map": layout_maker_factory((2, 1, 0)),
         "is_compatible_layout": lambda x: True,
-        "is_compatible_type": cuda_is_compatible_type,
     }
-    MODULE_GENERATOR_CLASS = DaCeCUDAPyExtModuleGenerator
+    MODULE_GENERATOR_CLASS = DaCePyExtModuleGenerator
     options = {
         **BaseGTBackend.GT_BACKEND_OPTS,
         "device_sync": {"versioning": True, "type": bool},

src/gt4py/backend/dace_stencil_object.py

@@ -30,6 +30,30 @@
 from gt4py.utils import shash
 
 
+@dataclass
+class _ArgsInfo:
+    device: str
+    array: dace.data.Array
+    origin: Optional[Tuple[int]] = None
+    dimensions: Optional[Tuple[str]] = None
+
+
+def _extract_array_infos(field_args, device) -> Dict[str, _ArgsInfo]:
+    return {
+        name: _ArgsInfo(
+            array=arg,
+            dimensions=getattr(arg, "dimensions", None),
+            device=device,
+            origin=getattr(arg, "origin", None),
+        )
+        for name, arg in field_args.items()
+    }
+
+
+def _extract_stencil_arrays(array_infos: Dict[str, _ArgsInfo]):
+    return {name: info.array for name, info in array_infos.items()}
+
+
 @dataclass(frozen=True)
 class DaCeFrozenStencil(FrozenStencil, SDFGConvertible):
 
@@ -298,10 +322,14 @@ def _normalize_args(self, *args, domain=None, origin=None, **kwargs):
         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=gt_backend.from_name(self.backend).storage_info["device"],
+        )
 
-        origin = self._normalize_origins(args_as_kwargs, origin)
+        origin = self._normalize_origins(arg_infos, origin)
         if domain is None:
-            domain = self._get_max_domain(args_as_kwargs, origin)
+            domain = self._get_max_domain(arg_infos, origin)
         for key, value in kwargs.items():
             args_as_kwargs.setdefault(key, value)
         args_as_kwargs["domain"] = domain

src/gt4py/backend/gtc_common.py

@@ -18,17 +18,17 @@
 import pathlib
 import textwrap
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 import numpy as np
 
+import gtc.utils
 import gtc.utils as gtc_utils
 from eve.codegen import MakoTemplate as as_mako
 from gt4py import backend as gt_backend
 from gt4py import utils as gt_utils
 from gt4py.backend import Backend
 from gt4py.backend.module_generator import BaseModuleGenerator, ModuleData
-from gt4py.definitions import AccessKind
 from gtc import gtir
 from gtc.passes.gtir_pipeline import GtirPipeline
 from gtc.passes.oir_pipeline import OirPipeline
@@ -42,9 +42,10 @@
 
 def _get_unit_stride_dim(backend, domain_dim_flags, data_ndim):
     make_layout_map = backend.storage_info["layout_map"]
-    layout_map = [
-        x for x in make_layout_map(domain_dim_flags + (True,) * data_ndim) if x is not None
+    dimensions = list(gtc.utils.dimension_flags_to_names(domain_dim_flags).upper()) + [
+        str(d) for d in range(data_ndim)
     ]
+    layout_map = [x for x in make_layout_map(dimensions) if x is not None]
     return layout_map.index(max(layout_map))
 
 
@@ -370,56 +371,52 @@ def generate_imports(self) -> str:
         return source
 
 
-class GTCUDAPyModuleGenerator(CUDAPyExtModuleGenerator):
-    def generate_pre_run(self) -> str:
-        field_names = [
-            key
-            for key in self.args_data.field_info
-            if self.args_data.field_info[key].access != AccessKind.NONE
-        ]
-
-        return "\n".join([f + ".host_to_device()" for f in field_names])
+def _dimensions_to_mask(dimensions: Tuple[str, ...]) -> Tuple[bool, ...]:
+    ndata_dims = sum(d.isdigit() for d in dimensions)
+    mask = [(d in dimensions) for d in "IJK"] + [True for _ in range(ndata_dims)]
+    return tuple(mask)
 
-    def generate_post_run(self) -> str:
-        output_field_names = [
-            name
-            for name, info in self.args_data.field_info.items()
-            if info is not None and bool(info.access & AccessKind.WRITE)
-        ]
 
-        return "\n".join([f + "._set_device_modified()" for f in output_field_names])
+def _permute_layout_to_dimensions(
+    layout: Sequence[int], dimensions: Tuple[str, ...]
+) -> Tuple[int, ...]:
+    data_dims = [int(d) for d in dimensions if d.isdigit()]
+    canonical_dimensions = [d for d in "IJK" if d in dimensions] + [
+        str(d) for d in sorted(data_dims)
+    ]
+    res_layout = []
+    for d in dimensions:
+        res_layout.append(layout[canonical_dimensions.index(d)])
+    return tuple(res_layout)
 
 
-def debug_layout(mask):
+def debug_layout(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ...]:
+    mask = _dimensions_to_mask(dimensions)
     ctr = iter(range(sum(mask)))
-    layout = [next(ctr) if m else None for m in mask]
-    return tuple(layout)
+    layout = [next(ctr) for m in mask if m]
+    return _permute_layout_to_dimensions(layout, dimensions)
 
 
-def debug_is_compatible_layout(field):
+def debug_is_compatible_layout(field, dimensions: Tuple[str, ...]):
     return sum(field.shape) > 0
 
 
-def debug_is_compatible_type(field):
-    return isinstance(field, np.ndarray)
-
-
-def make_x86_layout_map(mask: Tuple[int, ...]) -> Tuple[Optional[int], ...]:
+def make_x86_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ...]:
+    mask = _dimensions_to_mask(dimensions)
     ctr = iter(range(sum(mask)))
     if len(mask) < 3:
-        layout: List[Optional[int]] = [next(ctr) if m else None for m in mask]
+        layout: List[Optional[int]] = [next(ctr) for m in mask if m]
     else:
         swapped_mask: List[Optional[int]] = [*mask[3:], *mask[:3]]
         layout = [next(ctr) if m else None for m in swapped_mask]
 
         layout = [*layout[-3:], *layout[:-3]]
-
-    return tuple(layout)
+    return _permute_layout_to_dimensions([lt for lt in layout if lt is not None], dimensions)
 
 
-def x86_is_compatible_layout(field: "Storage") -> bool:
+def x86_is_compatible_layout(field: "Storage", dimensions: Tuple[str, ...]) -> bool:
     stride = 0
-    layout_map = make_x86_layout_map(field.mask)
+    layout_map = make_x86_layout_map(dimensions)
     flattened_layout = [index for index in layout_map if index is not None]
     if len(field.strides) < len(flattened_layout):
         return False
@@ -430,14 +427,11 @@ def x86_is_compatible_layout(field: "Storage") -> bool:
     return True
 
 
-def gtcpu_is_compatible_type(field: "Storage") -> bool:
-    return isinstance(field, np.ndarray)
-
-
-def make_mc_layout_map(mask: Tuple[int, ...]) -> Tuple[Optional[int], ...]:
-    ctr = reversed(range(sum(mask)))
+def make_mc_layout_map(dimensions: Tuple[str, ...]) -> Tuple[int, ...]:
+    mask = _dimensions_to_mask(dimensions)
+    ctr = reversed(range(len(dimensions)))
     if len(mask) < 3:
-        layout: List[Optional[int]] = [next(ctr) if m else None for m in mask]
+        layout: List[Optional[int]] = [next(ctr) for m in mask if m]
     else:
         swapped_mask: List[Optional[int]] = list(mask)
         tmp = swapped_mask[1]
@@ -449,13 +443,12 @@ def make_mc_layout_map(mask: Tuple[int, ...]) -> Tuple[Optional[int], ...]:
         tmp = layout[1]
         layout[1] = layout[2]
         layout[2] = tmp
-
-    return tuple(layout)
+    return _permute_layout_to_dimensions([lt for lt in layout if lt is not None], dimensions)
 
 
-def mc_is_compatible_layout(field: "Storage") -> bool:
+def mc_is_compatible_layout(field: "Storage", dimensions: Tuple[str, ...]) -> bool:
     stride = 0
-    layout_map = make_mc_layout_map(field.mask)
+    layout_map = make_mc_layout_map(dimensions)
     flattened_layout = [index for index in layout_map if index is not None]
     if len(field.strides) < len(flattened_layout):
         return False
@@ -466,14 +459,14 @@ def mc_is_compatible_layout(field: "Storage") -> bool:
     return True
 
 
-def make_cuda_layout_map(mask: Tuple[int, ...]) -> Tuple[Optional[int], ...]:
-    ctr = reversed(range(sum(mask)))
-    return tuple([next(ctr) if m else None for m in mask])
+def make_cuda_layout_map(dimensions: Tuple[str, ...]) -> Tuple[Optional[int], ...]:
+    layout = tuple(reversed(range(len(dimensions))))
+    return _permute_layout_to_dimensions(layout, dimensions)
 
 
-def cuda_is_compatible_layout(field: "Storage") -> bool:
+def cuda_is_compatible_layout(field: "Storage", dimensions: Tuple[str, ...]) -> bool:
     stride = 0
-    layout_map = make_cuda_layout_map(field.mask)
+    layout_map = make_cuda_layout_map(dimensions)
     flattened_layout = [index for index in layout_map if index is not None]
     if len(field.strides) < len(flattened_layout):
         return False
@@ -482,9 +475,3 @@ def cuda_is_compatible_layout(field: "Storage") -> bool:
             return False
         stride = field.strides[dim]
     return True
-
-
-def cuda_is_compatible_type(field: Any) -> bool:
-    from gt4py.storage.storage import ExplicitlySyncedGPUStorage, GPUStorage
-
-    return isinstance(field, (GPUStorage, ExplicitlySyncedGPUStorage))

src/gt4py/backend/gtcpp_backend.py

@@ -28,10 +28,8 @@
 
 from .gtc_common import (
     BaseGTBackend,
-    GTCUDAPyModuleGenerator,
+    PyExtModuleGenerator,
     cuda_is_compatible_layout,
-    cuda_is_compatible_type,
-    gtcpu_is_compatible_type,
     make_cuda_layout_map,
     make_mc_layout_map,
     make_x86_layout_map,
@@ -175,7 +173,6 @@ class GTCpuIfirstBackend(GTBaseBackend):
         "device": "cpu",
         "layout_map": make_mc_layout_map,
         "is_compatible_layout": mc_is_compatible_layout,
-        "is_compatible_type": gtcpu_is_compatible_type,
     }
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
@@ -194,7 +191,6 @@ class GTCpuKfirstBackend(GTBaseBackend):
         "device": "cpu",
         "layout_map": make_x86_layout_map,
         "is_compatible_layout": x86_is_compatible_layout,
-        "is_compatible_type": gtcpu_is_compatible_type,
     }
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
@@ -205,7 +201,7 @@ def generate_extension(self, **kwargs: Any) -> Tuple[str, str]:
 class GTGpuBackend(GTBaseBackend):
     """GridTools python backend using gtc."""
 
-    MODULE_GENERATOR_CLASS = GTCUDAPyModuleGenerator
+    MODULE_GENERATOR_CLASS = PyExtModuleGenerator
     name = "gt:gpu"
     GT_BACKEND_T = "gpu"
     languages = {"computation": "cuda", "bindings": ["python"]}
@@ -215,7 +211,6 @@ class GTGpuBackend(GTBaseBackend):
         "device": "gpu",
         "layout_map": make_cuda_layout_map,
         "is_compatible_layout": cuda_is_compatible_layout,
-        "is_compatible_type": cuda_is_compatible_type,
     }
 
     def generate_extension(self, **kwargs: Any) -> Tuple[str, str]: