Cartesian vs Next example

examples/lap_cartesian_vs_next.ipynb

@@ -0,0 +1,192 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "GT4Py - GridTools for Python\n",
+    "\n",
+    "Copyright (c) 2014-2023, ETH Zurich\n",
+    "All rights reserved.\n",
+    "\n",
+    "This file is part the GT4Py project and the GridTools framework.\n",
+    "GT4Py is free software: you can redistribute it and/or modify it under\n",
+    "the terms of the GNU General Public License as published by the\n",
+    "Free Software Foundation, either version 3 of the License, or any later\n",
+    "version. See the LICENSE.txt file at the top-level directory of this\n",
+    "distribution for a copy of the license or check <https://www.gnu.org/licenses/>.\n",
+    "\n",
+    "SPDX-License-Identifier: GPL-3.0-or-later"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Comparison gt4py.cartesian vs. gt4py.next"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "nx = 32\n",
+    "ny = 32\n",
+    "nz = 1\n",
+    "dtype = np.float64"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Storages\n",
+    "--"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gt4py.next as gtx\n",
+    "\n",
+    "allocator = gtx.itir_embedded # should match the executor\n",
+    "\n",
+    "# Note: for gt4py.next, names don't matter, for gt4py.cartesian they have to be \"I\", \"J\", \"K\"\n",
+    "I = gtx.Dimension(\"I\")\n",
+    "J = gtx.Dimension(\"J\")\n",
+    "K = gtx.Dimension(\"K\")\n",
+    "\n",
+    "domain = gtx.domain({I: nx, J: ny, K: nz})\n",
+    "\n",
+    "inp = gtx.as_field(domain, np.fromfunction(lambda x, y, z: x**2+y**2, shape=(nx, ny, nz)), dtype, allocator=allocator)\n",
+    "out_cartesian = gtx.zeros(domain, dtype, allocator=allocator)\n",
+    "out_next = gtx.zeros(domain, dtype, allocator=allocator)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "gt4py.cartesian\n",
+    "--"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/vogtha/git/gt4py/src/gt4py/cartesian/stencil_object.py:420: UserWarning: The layout of the field 'inp' is not recommended for this backend.This may lead to performance degradation. Please consider using theprovided allocators in `gt4py.storage`.\n",
+      "  warnings.warn(\n",
+      "/home/vogtha/git/gt4py/src/gt4py/cartesian/stencil_object.py:420: UserWarning: The layout of the field 'out' is not recommended for this backend.This may lead to performance degradation. Please consider using theprovided allocators in `gt4py.storage`.\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "import gt4py.cartesian.gtscript as gtscript\n",
+    "\n",
+    "cartesian_backend = \"numpy\"\n",
+    "cartesian_backend = \"gt:cpu_ifirst\"\n",
+    "\n",
+    "@gtscript.stencil(backend=cartesian_backend)\n",
+    "def lap_cartesian(\n",
+    "    inp: gtscript.Field[dtype],\n",
+    "    out: gtscript.Field[dtype],\n",
+    "):\n",
+    "    with computation(PARALLEL), interval(...):\n",
+    "        out = -4.0 * inp[0, 0, 0] + inp[-1, 0, 0] + inp[1, 0, 0] + inp[0, -1, 0] + inp[0, 1, 0]\n",
+    "\n",
+    "lap_cartesian(inp=inp, out=out_cartesian, origin=(1, 1, 0), domain=(nx-2, ny-2, nz))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "module 'gt4py.next' has no attribute 'cpu_backend'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[9], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mgt4py\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnext\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Field\n\u001b[1;32m      3\u001b[0m backend \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mitir_embedded\n\u001b[0;32m----> 4\u001b[0m backend \u001b[38;5;241m=\u001b[39m \u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcpu_backend\u001b[49m\n\u001b[1;32m      6\u001b[0m Ioff \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mFieldOffset(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mI\u001b[39m\u001b[38;5;124m\"\u001b[39m, source\u001b[38;5;241m=\u001b[39mI, target\u001b[38;5;241m=\u001b[39m(I,))\n\u001b[1;32m      7\u001b[0m Joff \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mFieldOffset(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mJ\u001b[39m\u001b[38;5;124m\"\u001b[39m, source\u001b[38;5;241m=\u001b[39mJ, target\u001b[38;5;241m=\u001b[39m(J,))\n",
+      "\u001b[0;31mAttributeError\u001b[0m: module 'gt4py.next' has no attribute 'cpu_backend'"
+     ]
+    }
+   ],
+   "source": [
+    "from gt4py.next import Field\n",
+    "\n",
+    "backend = gtx.itir_embedded\n",
+    "backend = gtx.gtfn_cpu\n",
+    "\n",
+    "Ioff = gtx.FieldOffset(\"I\", source=I, target=(I,))\n",
+    "Joff = gtx.FieldOffset(\"J\", source=J, target=(J,))\n",
+    "\n",
+    "@gtx.field_operator\n",
+    "def lap_next(inp: Field[[I, J, K], dtype]) -> Field[[I, J, K], dtype]:\n",
+    "    return -4.0 * inp + inp(Ioff[-1]) + inp(Ioff[1]) + inp(Joff[-1]) + inp(Joff[1])\n",
+    "\n",
+    "@gtx.program(backend=backend)\n",
+    "def lap_next_program(inp: Field[[I, J, K], dtype], out: Field[[I, J, K], dtype]):\n",
+    "    lap_next(inp, out=out[1:-1, 1:-1, :])\n",
+    "\n",
+    "lap_next_program(inp, out_next, offset_provider={\"Ioff\": I, \"Joff\": J})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert np.allclose(out_cartesian.asnumpy(), out_next.asnumpy())"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

src/gt4py/next/__init__.py

@@ -39,6 +39,11 @@
     index_field,
     np_as_located_field,
 )
+from .program_processors.runners.gtfn import (
+    run_gtfn_cached as gtfn_cpu,
+    run_gtfn_gpu_cached as gtfn_gpu,
+)
+from .program_processors.runners.roundtrip import backend as itir_embedded
 
 
 __all__ = [
@@ -74,5 +79,9 @@
     "field_operator",
     "program",
     "scan_operator",
+    # from program_processor
+    "gtfn_cpu",
+    "gtfn_gpu",
+    "itir_embedded",
     *fbuiltins.__all__,
 ]

src/gt4py/next/common.py

@@ -574,22 +574,7 @@ def __call__(self, func: fbuiltins.BuiltInFunction[_R, _P], /) -> Callable[_P, _
             ...
 
 
-# TODO(havogt): replace this protocol with the new `GTFieldInterface` protocol
-class NextGTDimsInterface(Protocol):
-    """
-    Protocol for objects providing the `__gt_dims__` property, naming :class:`Field` dimensions.
-
-    The dimension names are objects of type :class:`Dimension`, in contrast to
-    :mod:`gt4py.cartesian`, where the labels are `str` s with implied semantics,
-    see :class:`~gt4py._core.definitions.GTDimsInterface` .
-    """
-
-    @property
-    def __gt_dims__(self) -> tuple[Dimension, ...]:
-        ...
-
-
-# TODO(egparedes): add support for this new protocol in the cartesian module
+# TODO(havogt): we need to specify when we should use this interface vs the `Field` protocol.
 class GTFieldInterface(Protocol):
     """Protocol for object providing the `__gt_domain__` property, specifying the :class:`Domain` of a :class:`Field`."""
 
@@ -599,13 +584,23 @@ def __gt_domain__(self) -> Domain:
 
 
 @extended_runtime_checkable
-class Field(NextGTDimsInterface, core_defs.GTOriginInterface, Protocol[DimsT, core_defs.ScalarT]):
+class Field(
+    core_defs.GTDimsInterface, core_defs.GTOriginInterface, Protocol[DimsT, core_defs.ScalarT]
+):
     __gt_builtin_func__: ClassVar[GTBuiltInFuncDispatcher]
 
     @property
     def domain(self) -> Domain:
         ...
 
+    @property
+    def __gt_domain__(self) -> Domain:
+        return self.domain
+
+    @property
+    def __gt_dims__(self) -> tuple[str, ...]:
+        return tuple(d.value for d in self.domain.dims)
+
     @property
     def codomain(self) -> type[core_defs.ScalarT] | Dimension:
         ...
@@ -923,10 +918,6 @@ def asnumpy(self) -> Never:
     def domain(self) -> Domain:
         return Domain(dims=(self.dimension,), ranges=(UnitRange.infinite(),))
 
-    @property
-    def __gt_dims__(self) -> tuple[Dimension, ...]:
-        return self.domain.dims
-
     @property
     def __gt_origin__(self) -> Never:
         raise TypeError("'CartesianConnectivity' does not support this operation.")

src/gt4py/next/embedded/nd_array_field.py

@@ -107,10 +107,6 @@ def domain(self) -> common.Domain:
     def shape(self) -> tuple[int, ...]:
         return self._ndarray.shape
 
-    @property
-    def __gt_dims__(self) -> tuple[common.Dimension, ...]:
-        return self._domain.dims
-
     @property
     def __gt_origin__(self) -> tuple[int, ...]:
         assert common.Domain.is_finite(self._domain)

src/gt4py/next/iterator/embedded.py

@@ -172,7 +172,7 @@ class LocatedField(Protocol):
 
     @property
     @abc.abstractmethod
-    def __gt_dims__(self) -> tuple[common.Dimension, ...]:
+    def __gt_domain__(self) -> common.Domain:
         ...
 
     # TODO(havogt): define generic Protocol to provide a concrete return type
@@ -182,7 +182,7 @@ def field_getitem(self, indices: NamedFieldIndices) -> Any:
 
     @property
     def __gt_origin__(self) -> tuple[int, ...]:
-        return tuple([0] * len(self.__gt_dims__))
+        return tuple([0] * len(self.__gt_domain__.dims))
 
 
 @runtime_checkable
@@ -680,7 +680,18 @@ def _get_axes(
         assert all(first == _get_axes(f) for f in field_or_tuple)
         return first
     else:
-        return field_or_tuple.__gt_dims__
+        return field_or_tuple.__gt_domain__.dims
+
+
+def _get_domain(
+    field_or_tuple: LocatedField | tuple,
+) -> common.Domain:  # arbitrary nesting of tuples of LocatedField
+    if isinstance(field_or_tuple, tuple):
+        first = _get_domain(field_or_tuple[0])
+        assert all(first == _get_domain(f) for f in field_or_tuple)
+        return first
+    else:
+        return field_or_tuple.__gt_domain__
 
 
 def _single_vertical_idx(
@@ -894,14 +905,14 @@ class NDArrayLocatedFieldWrapper(MutableLocatedField):
     _ndarrayfield: common.Field
 
     @property
-    def __gt_dims__(self) -> tuple[common.Dimension, ...]:
-        return self._ndarrayfield.__gt_dims__
+    def __gt_domain__(self) -> common.Domain:
+        return self._ndarrayfield.__gt_domain__
 
     def _translate_named_indices(
         self, _named_indices: NamedFieldIndices
     ) -> common.AbsoluteIndexSequence:
         named_indices: Mapping[common.Dimension, FieldIndex | SparsePositionEntry] = {
-            d: _named_indices[d.value] for d in self._ndarrayfield.__gt_dims__
+            d: _named_indices[d.value] for d in self._ndarrayfield.__gt_domain__.dims
         }
         domain_slice: list[common.NamedRange | common.NamedIndex] = []
         for d, v in named_indices.items():
@@ -1046,8 +1057,8 @@ class IndexField(common.Field):
     _dimension: common.Dimension
 
     @property
-    def __gt_dims__(self) -> tuple[common.Dimension, ...]:
-        return (self._dimension,)
+    def __gt_domain__(self) -> common.Domain:
+        return self.domain
 
     @property
     def __gt_origin__(self) -> tuple[int, ...]:
@@ -1165,8 +1176,8 @@ class ConstantField(common.Field[Any, core_defs.ScalarT]):
     _value: core_defs.ScalarT
 
     @property
-    def __gt_dims__(self) -> tuple[common.Dimension, ...]:
-        return tuple()
+    def __gt_domain__(self) -> common.Domain:
+        return self.domain
 
     @property
     def __gt_origin__(self) -> tuple[int, ...]:
@@ -1452,7 +1463,7 @@ def _tuple_assign(field: tuple | MutableLocatedField, value: Any, named_indices:
 class TupleOfFields(TupleField):
     def __init__(self, data):
         self.data = data
-        self.__gt_dims__ = _get_axes(data)
+        self.__gt_domain__ = _get_domain(data)
 
     def field_getitem(self, named_indices: NamedFieldIndices) -> Any:
         return _build_tuple_result(self.data, named_indices)

src/gt4py/next/iterator/tracing.py

@@ -254,7 +254,7 @@ def _contains_tuple_dtype_field(arg):
     #  other `np.int32`). We just ignore the error here and postpone fixing this to when
     #  the new storages land (The implementation here works for LocatedFieldImpl).
 
-    return common.is_field(arg) and any(dim is None for dim in arg.__gt_dims__)
+    return common.is_field(arg) and any(dim is None for dim in arg.domain.dims)
 
 
 def _make_fencil_params(fun, args, *, use_arg_types: bool) -> list[Sym]:

src/gt4py/next/type_system/type_translation.py

@@ -184,7 +184,7 @@ def from_value(value: Any) -> ts.TypeSpec:
     elif isinstance(value, common.Dimension):
         symbol_type = ts.DimensionType(dim=value)
     elif common.is_field(value):
-        dims = list(value.__gt_dims__)
+        dims = list(value.domain.dims)
         dtype = from_type_hint(value.dtype.scalar_type)
         symbol_type = ts.FieldType(dims=dims, dtype=dtype)
     elif isinstance(value, tuple):

src/gt4py/storage/cartesian/utils.py

@@ -192,6 +192,8 @@ def cpu_copy(array: Union[np.ndarray, "cp.ndarray"]) -> np.ndarray:
 def asarray(
     array: FieldLike, *, device: Literal["cpu", "gpu", None] = None
 ) -> np.ndarray | cp.ndarray:
+    if hasattr(array, "ndarray"):
+        array = array.ndarray
     if device == "gpu" or (not device and hasattr(array, "__cuda_array_interface__")):
         return cp.asarray(array)
     if device == "cpu" or (