From 385860e8c8c1d542f41071d08bce603cccaf98d7 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 14 May 2024 10:50:50 -0400 Subject: [PATCH 1/4] Basic boilerplate + utset --- ndsl/boilerplate.py | 93 +++++++++++++++++++++++++++++++++++++++ tests/test_boilerplate.py | 61 +++++++++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 ndsl/boilerplate.py create mode 100644 tests/test_boilerplate.py diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py new file mode 100644 index 00000000..bf2ef7aa --- /dev/null +++ b/ndsl/boilerplate.py @@ -0,0 +1,93 @@ +import numpy as np +from ndsl import ( + StencilFactory, + DaceConfig, + DaCeOrchestration, + GridIndexing, + StencilConfig, + CompilationConfig, + RunMode, + SubtileGridSizer, + NullComm, + QuantityFactory, + TileCommunicator, + TilePartitioner, +) + +from typing import Tuple + + +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 + ) diff --git a/tests/test_boilerplate.py b/tests/test_boilerplate.py new file mode 100644 index 00000000..b677a7a2 --- /dev/null +++ b/tests/test_boilerplate.py @@ -0,0 +1,61 @@ +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import FloatField +from ndsl import StencilFactory, QuantityFactory +import numpy as np +from gt4py.cartesian.gtscript import ( + computation, PARALLEL, interval +) + +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) From 50b8d82692ea3ed12f1f27bc6a79e0d44f4d11ad Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 14 May 2024 11:01:16 -0400 Subject: [PATCH 2/4] lint --- ndsl/boilerplate.py | 26 ++++++++++++++------------ tests/test_boilerplate.py | 37 ++++++++++++++++++------------------- 2 files changed, 32 insertions(+), 31 deletions(-) diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py index bf2ef7aa..d2b0cb96 100644 --- a/ndsl/boilerplate.py +++ b/ndsl/boilerplate.py @@ -1,28 +1,29 @@ +from typing import Tuple + import numpy as np + from ndsl import ( - StencilFactory, + CompilationConfig, DaceConfig, DaCeOrchestration, GridIndexing, - StencilConfig, - CompilationConfig, - RunMode, - SubtileGridSizer, NullComm, QuantityFactory, + RunMode, + StencilConfig, + StencilFactory, + SubtileGridSizer, TileCommunicator, TilePartitioner, ) -from typing import Tuple - def _get_one_tile_factory( nx, ny, nz, nhalo, backend, orchestration -)-> Tuple[StencilFactory, QuantityFactory]: +) -> Tuple[StencilFactory, QuantityFactory]: """Build a Stencil & Quantity factory for: - - one tile - - no MPI communicator + - one tile + - no MPI communicator """ dace_config = DaceConfig( communicator=None, @@ -76,9 +77,10 @@ def get_one_tile_factory_orchestrated_cpu( nz=nz, nhalo=nhalo, backend="dace:cpu", - orchestration=DaCeOrchestration.BuildAndRun + orchestration=DaCeOrchestration.BuildAndRun, ) + def get_one_tile_factory_numpy( nx, ny, nz, nhalo ) -> Tuple[StencilFactory, QuantityFactory]: @@ -89,5 +91,5 @@ def get_one_tile_factory_numpy( nz=nz, nhalo=nhalo, backend="numpy", - orchestration=DaCeOrchestration.Python + orchestration=DaCeOrchestration.Python, ) diff --git a/tests/test_boilerplate.py b/tests/test_boilerplate.py index b677a7a2..983a8d07 100644 --- a/tests/test_boilerplate.py +++ b/tests/test_boilerplate.py @@ -1,10 +1,10 @@ +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 -from ndsl import StencilFactory, QuantityFactory -import numpy as np -from gt4py.cartesian.gtscript import ( - computation, PARALLEL, interval -) + def _copy_ops(stencil_factory: StencilFactory, quantity_factory: QuantityFactory): # Allocate data and fill input @@ -12,7 +12,8 @@ def _copy_ops(stencil_factory: StencilFactory, quantity_factory: QuantityFactory 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( + dtype=np.float64, + ).sum( axis=0 ) # Value of each entry is sum of the I and J index at each point @@ -22,22 +23,23 @@ def copy_stencil(input_field: FloatField, output_field: FloatField): output_field = input_field # Execute - copy = stencil_factory.from_dims_halo(func=copy_stencil, compute_dims=[X_DIM, Y_DIM, Z_DIM]) + 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. + + 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 + nx=5, ny=5, nz=2, nhalo=1 ) _copy_ops(stencil_factory, quantity_factory) @@ -45,17 +47,14 @@ def test_boilerplate_import_numpy(): 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. + + 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 + nx=5, ny=5, nz=2, nhalo=1 ) _copy_ops(stencil_factory, quantity_factory) From 0f8906b084bba1bb3774ae0ac1c9f9bb647aa385 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 May 2024 09:55:41 -0400 Subject: [PATCH 3/4] Adjust naming & verbose choice for private generic function --- ndsl/boilerplate.py | 65 ++++++++++++++++++++++++--------------- tests/test_boilerplate.py | 8 ++--- 2 files changed, 45 insertions(+), 28 deletions(-) diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py index d2b0cb96..a07a476e 100644 --- a/ndsl/boilerplate.py +++ b/ndsl/boilerplate.py @@ -4,9 +4,12 @@ from ndsl import ( CompilationConfig, + CubedSphereCommunicator, + CubedSpherePartitioner, DaceConfig, DaCeOrchestration, GridIndexing, + LocalComm, NullComm, QuantityFactory, RunMode, @@ -18,12 +21,22 @@ ) -def _get_one_tile_factory( - nx, ny, nz, nhalo, backend, orchestration +def _get_factories( + nx: int, + ny: int, + nz: int, + nhalo, + backend: str, + orchestration: DaCeOrchestration, + topology: str, ) -> Tuple[StencilFactory, QuantityFactory]: - """Build a Stencil & Quantity factory for: - - one tile - - no MPI communicator + """Build a Stencil & Quantity factory for a combination of options. + + Dev Note: We don't expose this function because we want the boilerplate to remain + as easy and self describing as possible. It should be a very easy call to make. + The other reason is that the orchestration requires two inputs instead of change + a backend name for now, making it confusing. Until refactor, we choose to hide this + pattern for boilerplate use. """ dace_config = DaceConfig( communicator=None, @@ -47,49 +60,53 @@ def _get_one_tile_factory( 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) + if topology == "tile": + 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, + ) + comm = TileCommunicator(comm=NullComm(0, 1, 42), partitioner=partitioner) + else: + raise NotImplementedError(f"Topology {topology} is not implemented.") - grid_indexing = GridIndexing.from_sizer_and_communicator(sizer, tile_comm) + grid_indexing = GridIndexing.from_sizer_and_communicator(sizer, 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( +def get_factories_single_tile_orchestrated_cpu( nx, ny, nz, nhalo ) -> Tuple[StencilFactory, QuantityFactory]: - """Build a Stencil & Quantity factory for orchestrated CPU""" - return _get_one_tile_factory( + """Build a Stencil & Quantity factory for orchestrated CPU, on a single tile toplogy.""" + return _get_factories( nx=nx, ny=ny, nz=nz, nhalo=nhalo, backend="dace:cpu", orchestration=DaCeOrchestration.BuildAndRun, + topology="tile", ) -def get_one_tile_factory_numpy( +def get_factories_single_tile_numpy( nx, ny, nz, nhalo ) -> Tuple[StencilFactory, QuantityFactory]: - """Build a Stencil & Quantity factory for Numpy""" - return _get_one_tile_factory( + """Build a Stencil & Quantity factory for Numpy, on a single tile toplogy.""" + return _get_factories( nx=nx, ny=ny, nz=nz, nhalo=nhalo, backend="numpy", orchestration=DaCeOrchestration.Python, + topology="tile", ) diff --git a/tests/test_boilerplate.py b/tests/test_boilerplate.py index 983a8d07..a8de4075 100644 --- a/tests/test_boilerplate.py +++ b/tests/test_boilerplate.py @@ -35,10 +35,10 @@ def test_boilerplate_import_numpy(): Dev Note: the import inside the function are part of the test. """ - from ndsl.boilerplate import get_one_tile_factory_numpy + from ndsl.boilerplate import get_factories_single_tile_numpy # Boilerplate - stencil_factory, quantity_factory = get_one_tile_factory_numpy( + stencil_factory, quantity_factory = get_factories_single_tile_numpy( nx=5, ny=5, nz=2, nhalo=1 ) @@ -50,10 +50,10 @@ def test_boilerplate_import_orchestrated_cpu(): Dev Note: the import inside the function are part of the test. """ - from ndsl.boilerplate import get_one_tile_factory_orchestrated_cpu + from ndsl.boilerplate import get_factories_single_tile_orchestrated_cpu # Boilerplate - stencil_factory, quantity_factory = get_one_tile_factory_orchestrated_cpu( + stencil_factory, quantity_factory = get_factories_single_tile_orchestrated_cpu( nx=5, ny=5, nz=2, nhalo=1 ) From 0789fc0151fc2239603560a31d1ea02dca91528f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 May 2024 10:01:21 -0400 Subject: [PATCH 4/4] Lint --- ndsl/boilerplate.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py index a07a476e..ced15115 100644 --- a/ndsl/boilerplate.py +++ b/ndsl/boilerplate.py @@ -4,12 +4,9 @@ from ndsl import ( CompilationConfig, - CubedSphereCommunicator, - CubedSpherePartitioner, DaceConfig, DaCeOrchestration, GridIndexing, - LocalComm, NullComm, QuantityFactory, RunMode,