Boilerplate code for fast deploy of NDSL

ndsl/boilerplate.py

@@ -0,0 +1,95 @@
+from typing import Tuple
+
+import numpy as np
+
+from ndsl import (
+    CompilationConfig,
+    DaceConfig,
+    DaCeOrchestration,
+    GridIndexing,
+    NullComm,
+    QuantityFactory,
+    RunMode,
+    StencilConfig,
+    StencilFactory,
+    SubtileGridSizer,
+    TileCommunicator,
+    TilePartitioner,
+)
+
+
+def _get_one_tile_factory(
+    nx, ny, nz, nhalo, backend, orchestration
+) -> Tuple[StencilFactory, QuantityFactory]:
+    """Build a Stencil & Quantity factory for:
+    - one tile
+    - no MPI communicator
+    """
+    dace_config = DaceConfig(
+        communicator=None,
+        backend=backend,
+        orchestration=orchestration,
+    )
+
+    compilation_config = CompilationConfig(
+        backend=backend,
+        rebuild=True,
+        validate_args=True,
+        format_source=False,
+        device_sync=False,
+        run_mode=RunMode.BuildAndRun,
+        use_minimal_caching=False,
+    )
+
+    stencil_config = StencilConfig(
+        compare_to_numpy=False,
+        compilation_config=compilation_config,
+        dace_config=dace_config,
+    )
+
+    partitioner = TilePartitioner((1, 1))
+    sizer = SubtileGridSizer.from_tile_params(
+        nx_tile=nx,
+        ny_tile=ny,
+        nz=nz,
+        n_halo=nhalo,
+        extra_dim_lengths={},
+        layout=partitioner.layout,
+        tile_partitioner=partitioner,
+    )
+
+    tile_comm = TileCommunicator(comm=NullComm(0, 1, 42), partitioner=partitioner)
+
+    grid_indexing = GridIndexing.from_sizer_and_communicator(sizer, tile_comm)
+    stencil_factory = StencilFactory(config=stencil_config, grid_indexing=grid_indexing)
+    quantity_factory = QuantityFactory(sizer, np)
+
+    return stencil_factory, quantity_factory
+
+
+def get_one_tile_factory_orchestrated_cpu(
+    nx, ny, nz, nhalo
+) -> Tuple[StencilFactory, QuantityFactory]:
+    """Build a Stencil & Quantity factory for orchestrated CPU"""
+    return _get_one_tile_factory(
+        nx=nx,
+        ny=ny,
+        nz=nz,
+        nhalo=nhalo,
+        backend="dace:cpu",
+        orchestration=DaCeOrchestration.BuildAndRun,
+    )
+
+
+def get_one_tile_factory_numpy(
+    nx, ny, nz, nhalo
+) -> Tuple[StencilFactory, QuantityFactory]:
+    """Build a Stencil & Quantity factory for Numpy"""
+    return _get_one_tile_factory(
+        nx=nx,
+        ny=ny,
+        nz=nz,
+        nhalo=nhalo,
+        backend="numpy",
+        orchestration=DaCeOrchestration.Python,
+    )

tests/test_boilerplate.py

@@ -0,0 +1,60 @@
+import numpy as np
+from gt4py.cartesian.gtscript import PARALLEL, computation, interval
+
+from ndsl import QuantityFactory, StencilFactory
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.typing import FloatField
+
+
+def _copy_ops(stencil_factory: StencilFactory, quantity_factory: QuantityFactory):
+    # Allocate data and fill input
+    qty_out = quantity_factory.zeros(dims=[X_DIM, Y_DIM, Z_DIM], units="n/a")
+    qty_in = quantity_factory.zeros(dims=[X_DIM, Y_DIM, Z_DIM], units="n/a")
+    qty_in.view[:] = np.indices(
+        dimensions=quantity_factory.sizer.get_extent([X_DIM, Y_DIM, Z_DIM]),
+        dtype=np.float64,
+    ).sum(
+        axis=0
+    )  # Value of each entry is sum of the I and J index at each point
+
+    # Define a stencil
+    def copy_stencil(input_field: FloatField, output_field: FloatField):
+        with computation(PARALLEL), interval(...):
+            output_field = input_field
+
+    # Execute
+    copy = stencil_factory.from_dims_halo(
+        func=copy_stencil, compute_dims=[X_DIM, Y_DIM, Z_DIM]
+    )
+    copy(qty_in, qty_out)
+    assert (qty_in.view[:] == qty_out.view[:]).all()
+
+
+def test_boilerplate_import_numpy():
+    """Test make sure the basic numpy boilerplate works as expected.
+
+    Dev Note: the import inside the function are part of the test.
+    """
+    from ndsl.boilerplate import get_one_tile_factory_numpy
+
+    # Boilerplate
+    stencil_factory, quantity_factory = get_one_tile_factory_numpy(
+        nx=5, ny=5, nz=2, nhalo=1
+    )
+
+    _copy_ops(stencil_factory, quantity_factory)
+
+
+def test_boilerplate_import_orchestrated_cpu():
+    """Test make sure the basic orchestrate boilerplate works as expected.
+
+    Dev Note: the import inside the function are part of the test.
+    """
+    from ndsl.boilerplate import get_one_tile_factory_orchestrated_cpu
+
+    # Boilerplate
+    stencil_factory, quantity_factory = get_one_tile_factory_orchestrated_cpu(
+        nx=5, ny=5, nz=2, nhalo=1
+    )
+
+    _copy_ops(stencil_factory, quantity_factory)