cleanup

GridTools · Jan 15, 2024 · c670872 · c670872
1 parent 4281e70
commit c670872
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 34 deletions.
diff --git a/examples/lap_cartesian_vs_next.ipynb b/examples/lap_cartesian_vs_next.ipynb
@@ -24,7 +24,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Comparison gt4py.cartesian vs. gt4py.next"
+    "# Demonstrates gt4py.cartesian with gt4py.next compatibility"
    ]
   },
   {
@@ -37,7 +37,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -54,23 +54,41 @@
    "metadata": {},
    "source": [
     "Storages\n",
-    "--"
+    "--\n",
+    "\n",
+    "We create fields using the gt4py.next constructors. These fields are compatible with gt4py.cartesian when we use \"I\", \"J\", \"K\" as the dimension names."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<link href=\"https://fonts.googleapis.com/icon?family=Material+Icons\" rel=\"stylesheet\"><script src=\"https://spcl.github.io/dace/webclient2/dist/sdfv.js\"></script>\n",
+       "<link href=\"https://spcl.github.io/dace/webclient2/sdfv.css\" rel=\"stylesheet\">\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "import gt4py.next as gtx\n",
     "\n",
     "allocator = gtx.itir_embedded # should match the executor\n",
+    "# allocator = gtx.gtfn_cpu\n",
+    "# allocator = gtx.gtfn_gpu\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",
+    "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)\n",
     "\n",
     "domain = gtx.domain({I: nx, J: ny, K: nz})\n",
     "\n",
@@ -90,7 +108,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -108,7 +126,8 @@
     "import gt4py.cartesian.gtscript as gtscript\n",
     "\n",
     "cartesian_backend = \"numpy\"\n",
-    "cartesian_backend = \"gt:cpu_ifirst\"\n",
+    "# cartesian_backend = \"gt:cpu_ifirst\"\n",
+    "# cartesian_backend = \"gt:gpu\"\n",
     "\n",
     "@gtscript.stencil(backend=cartesian_backend)\n",
     "def lap_cartesian(\n",
@@ -123,26 +142,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 4,
    "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'"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from gt4py.next import Field\n",
     "\n",
-    "backend = gtx.itir_embedded\n",
-    "backend = gtx.gtfn_cpu\n",
+    "next_backend = gtx.itir_embedded\n",
+    "# next_backend = gtx.gtfn_cpu\n",
+    "# next_backend = gtx.gtfn_gpu\n",
     "\n",
     "Ioff = gtx.FieldOffset(\"I\", source=I, target=(I,))\n",
     "Joff = gtx.FieldOffset(\"J\", source=J, target=(J,))\n",
@@ -151,7 +159,7 @@
     "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",
+    "@gtx.program(backend=next_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",
@@ -160,7 +168,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [

diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py
@@ -574,18 +574,30 @@ def __call__(self, func: fbuiltins.BuiltInFunction[_R, _P], /) -> Callable[_P, _
             ...
 
 
-# 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`."""
+# TODO(havogt): we need to describe when this interface shouold be used instead of the `Field` protocol.
+class GTFieldInterface(core_defs.GTDimsInterface, core_defs.GTOriginInterface, Protocol):
+    """
+    Protocol for object providing the `__gt_domain__` property, specifying the :class:`Domain` of a :class:`Field`.
+
+    Note:
+    - A default impolementation of the `__gt_dims__` interface from `gt4py.cartesian` is provided.
+    - No implementation of `__gt_origin__` is provided because of infinite fields.
+    """
 
     @property
     def __gt_domain__(self) -> Domain:
+        # TODO probably should be changed to `DomainLike` (with a new concept `DimensionLike`)
+        # to allow implementations without having to import gtx.Domain.
         ...
 
+    @property
+    def __gt_dims__(self) -> tuple[str, ...]:
+        return tuple(d.value for d in self.__gt_domain__.dims)
+
 
 @extended_runtime_checkable
 class Field(
-    core_defs.GTDimsInterface, core_defs.GTOriginInterface, Protocol[DimsT, core_defs.ScalarT]
+    GTFieldInterface, Protocol[DimsT, core_defs.ScalarT]
 ):
     __gt_builtin_func__: ClassVar[GTBuiltInFuncDispatcher]
 
@@ -597,10 +609,6 @@ def domain(self) -> Domain:
     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:
         ...

diff --git a/src/gt4py/storage/cartesian/utils.py b/src/gt4py/storage/cartesian/utils.py
@@ -193,6 +193,8 @@ def asarray(
     array: FieldLike, *, device: Literal["cpu", "gpu", None] = 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
     if device == "gpu" or (not device and hasattr(array, "__cuda_array_interface__")):
         return cp.asarray(array)