Extended corner usage

pyfv3/stencils/copy_corners.py

@@ -0,0 +1,169 @@
+from ndsl import StencilFactory, orchestrate
+from ndsl.dsl.typing import FloatField
+
+
+def corner_copy_x(field_to_copy):
+    """Equivalent to the copy_corners_x functions in fortran.
+
+    This is written to operate on plain ndarrarys and not use the GT4Py framework.
+    This choice was made because we've seen a lot of performance left on the table using
+    orchestration without explicitly describing the operations but rather have full 3d-
+    sweeps with conditionals.
+    Since DaCe can handle (simple) operations on ndarrays directly this gives us a more
+    explicit entrypoint to the language and more optimization-potential.
+
+    Args:
+        field_to_copy (ndarray): field to apply the corner copy on.
+            This is explicitly not type-hinted for orchestration
+    """
+    field_to_copy[0, 0] = field_to_copy[0, 5]
+    field_to_copy[0, 1] = field_to_copy[1, 5]
+    field_to_copy[0, 2] = field_to_copy[2, 5]
+
+    field_to_copy[1, 0] = field_to_copy[0, 4]
+    field_to_copy[1, 1] = field_to_copy[1, 4]
+    field_to_copy[1, 2] = field_to_copy[2, 4]
+
+    field_to_copy[2, 0] = field_to_copy[0, 3]
+    field_to_copy[2, 1] = field_to_copy[1, 3]
+    field_to_copy[2, 2] = field_to_copy[2, 3]
+
+    field_to_copy[0, -4] = field_to_copy[2, -7]
+    field_to_copy[0, -3] = field_to_copy[1, -7]
+    field_to_copy[0, -2] = field_to_copy[0, -7]
+
+    field_to_copy[1, -4] = field_to_copy[2, -6]
+    field_to_copy[1, -3] = field_to_copy[1, -6]
+    field_to_copy[1, -2] = field_to_copy[0, -6]
+
+    field_to_copy[2, -4] = field_to_copy[2, -5]
+    field_to_copy[2, -3] = field_to_copy[1, -5]
+    field_to_copy[2, -2] = field_to_copy[0, -5]
+
+    field_to_copy[-4, 0] = field_to_copy[-2, 3]
+    field_to_copy[-4, 1] = field_to_copy[-3, 3]
+    field_to_copy[-4, 2] = field_to_copy[-4, 3]
+
+    field_to_copy[-3, 0] = field_to_copy[-2, 4]
+    field_to_copy[-3, 1] = field_to_copy[-3, 4]
+    field_to_copy[-3, 2] = field_to_copy[-4, 4]
+
+    field_to_copy[-2, 0] = field_to_copy[-2, 5]
+    field_to_copy[-2, 1] = field_to_copy[-3, 5]
+    field_to_copy[-2, 2] = field_to_copy[-4, 5]
+
+    field_to_copy[-4, -2] = field_to_copy[-2, -5]
+    field_to_copy[-4, -3] = field_to_copy[-3, -5]
+    field_to_copy[-4, -4] = field_to_copy[-4, -5]
+
+    field_to_copy[-3, -2] = field_to_copy[-2, -6]
+    field_to_copy[-3, -3] = field_to_copy[-3, -6]
+    field_to_copy[-3, -4] = field_to_copy[-4, -6]
+
+    field_to_copy[-2, -2] = field_to_copy[-2, -7]
+    field_to_copy[-2, -3] = field_to_copy[-3, -7]
+    field_to_copy[-2, -4] = field_to_copy[-4, -7]
+
+
+def corner_copy_y(field_to_copy):
+    """Equivalent to the copy_corners_y functions in fortran.
+
+    This is written to operate on plain ndarrarys and not use the GT4Py framework.
+    This choice was made because we've seen a lot of performance left on the table using
+    orchestration without explicitly describing the operations but rather have full 3d-
+    sweeps with conditionals.
+    Since DaCe can handle (simple) operations on ndarrays directly this gives us a more
+    explicit entrypoint to the language and more optimization-potential.
+
+    Args:
+        field_to_copy (ndarray): field to apply the corner copy on.
+            This is explicitly not type-hinted for orchestration
+    """
+    field_to_copy[0, 0] = field_to_copy[5, 0]
+    field_to_copy[1, 0] = field_to_copy[5, 1]
+    field_to_copy[2, 0] = field_to_copy[5, 2]
+
+    field_to_copy[0, 1] = field_to_copy[4, 0]
+    field_to_copy[1, 1] = field_to_copy[4, 1]
+    field_to_copy[2, 1] = field_to_copy[4, 2]
+
+    field_to_copy[0, 2] = field_to_copy[3, 0]
+    field_to_copy[1, 2] = field_to_copy[3, 1]
+    field_to_copy[2, 2] = field_to_copy[3, 2]
+
+    field_to_copy[-4, 0] = field_to_copy[-7, 2]
+    field_to_copy[-3, 0] = field_to_copy[-7, 1]
+    field_to_copy[-2, 0] = field_to_copy[-7, 0]
+
+    field_to_copy[-4, 1] = field_to_copy[-6, 2]
+    field_to_copy[-3, 1] = field_to_copy[-6, 1]
+    field_to_copy[-2, 1] = field_to_copy[-6, 0]
+
+    field_to_copy[-4, 2] = field_to_copy[-5, 2]
+    field_to_copy[-3, 2] = field_to_copy[-5, 1]
+    field_to_copy[-2, 2] = field_to_copy[-5, 0]
+
+    field_to_copy[0, -2] = field_to_copy[5, -2]
+    field_to_copy[0, -3] = field_to_copy[4, -2]
+    field_to_copy[0, -4] = field_to_copy[3, -2]
+
+    field_to_copy[1, -2] = field_to_copy[5, -3]
+    field_to_copy[1, -3] = field_to_copy[4, -3]
+    field_to_copy[1, -4] = field_to_copy[3, -3]
+
+    field_to_copy[2, -2] = field_to_copy[5, -4]
+    field_to_copy[2, -3] = field_to_copy[4, -4]
+    field_to_copy[2, -4] = field_to_copy[3, -4]
+
+    field_to_copy[-2, -4] = field_to_copy[-5, -2]
+    field_to_copy[-2, -3] = field_to_copy[-6, -2]
+    field_to_copy[-2, -2] = field_to_copy[-7, -2]
+
+    field_to_copy[-3, -4] = field_to_copy[-5, -3]
+    field_to_copy[-3, -3] = field_to_copy[-6, -3]
+    field_to_copy[-3, -2] = field_to_copy[-7, -3]
+
+    field_to_copy[-4, -4] = field_to_copy[-5, -4]
+    field_to_copy[-4, -3] = field_to_copy[-6, -4]
+    field_to_copy[-4, -2] = field_to_copy[-7, -4]
+
+
+class CopyCornersX:
+    """
+    Helper-class to copy corners corresponding to the fortran function copy_corners_x
+    """
+
+    def __init__(self, stencil_factory: StencilFactory) -> None:
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+        )
+
+        if stencil_factory.grid_indexing.n_halo != 3:
+            raise NotImplementedError(
+                "Corner-Copy only implemented for exactly 3 Halo-Points"
+            )
+
+    def __call__(self, field: FloatField):
+        corner_copy_x(field)
+
+
+class CopyCornersY:
+    """
+    Helper-class to copy corners corresponding to the fortran function
+    copy_corners_y
+    """
+
+    def __init__(self, stencil_factory: StencilFactory) -> None:
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+        )
+
+        if stencil_factory.grid_indexing.n_halo != 3:
+            raise NotImplementedError(
+                "Corner-Copy only implemented for exactly 3 Halo-Points"
+            )
+
+    def __call__(self, field: FloatField):
+        corner_copy_y(field)