diff --git a/ndsl/__init__.py b/ndsl/__init__.py index 18c96f3e..f6bb6117 100644 --- a/ndsl/__init__.py +++ b/ndsl/__init__.py @@ -29,6 +29,7 @@ from .performance.profiler import NullProfiler, Profiler from .performance.report import Experiment, Report, TimeReport from .quantity import Quantity +from .quantity.field_bundle import FieldBundle, FieldBundleType # Break circular import from .testing.dummy_comm import DummyComm from .types import Allocator from .utils import MetaEnumStr diff --git a/ndsl/dsl/stencil.py b/ndsl/dsl/stencil.py index daf78091..9260bd29 100644 --- a/ndsl/dsl/stencil.py +++ b/ndsl/dsl/stencil.py @@ -33,6 +33,7 @@ from ndsl.initialization.sizer import GridSizer, SubtileGridSizer from ndsl.logging import ndsl_log from ndsl.quantity import Quantity +from ndsl.quantity.field_bundle import FieldBundleType, MarkupFieldBundleType from ndsl.testing.comparison import LegacyMetric @@ -335,6 +336,14 @@ def __init__( ): block_waiting_for_compilation(MPI.COMM_WORLD, compilation_config) + # Field Bundle might have dropped a placeholder type that we now + # have to resolve to the proper type. + for name, types in func.__annotations__.items(): + if isinstance(types, MarkupFieldBundleType): + func.__annotations__[name] = FieldBundleType.T( + types.name, do_markup=False + ) + self.stencil_object = gtscript.stencil( definition=func, externals=externals, diff --git a/ndsl/quantity/__init__.py b/ndsl/quantity/__init__.py index fce0cf63..43751528 100644 --- a/ndsl/quantity/__init__.py +++ b/ndsl/quantity/__init__.py @@ -1,9 +1,11 @@ -from ndsl.quantity.metadata import QuantityHaloSpec, QuantityMetadata -from ndsl.quantity.quantity import Quantity +from .metadata import QuantityHaloSpec, QuantityMetadata +from .quantity import Quantity __all__ = [ "Quantity", "QuantityMetadata", "QuantityHaloSpec", + "FieldBundle", + "FieldBundleType", ] diff --git a/ndsl/quantity/field_bundle.py b/ndsl/quantity/field_bundle.py new file mode 100644 index 00000000..a2749361 --- /dev/null +++ b/ndsl/quantity/field_bundle.py @@ -0,0 +1,190 @@ +import copy +from dataclasses import dataclass +from typing import Any + +from ndsl.dsl.typing import Float +from ndsl.initialization.allocator import QuantityFactory +from ndsl.quantity.quantity import Quantity + + +import gt4py.cartesian.gtscript as gtscript # isort: skip + + +# ToDo: This is 4th dimensions restricted. We need a concept +# of data dimensions index here to be able to extend to N dimensions +_DataDimensionIndex = int +_FieldBundleIndexer = dict[str, _DataDimensionIndex] + + +class FieldBundle: + """Field Bundle wraps a nD array (3D + n Data Dimensions) into a complex + indexing scheme that allows a dual name-based and index-based + access to the underlying memory. It is paired with the `FieldBundleType` + which provides a way to type hint parameters for stencils in the `gtscript`. + + WARNING: The present implementation only allows for 4D array. + """ + + _quantity: Quantity + _indexer: _FieldBundleIndexer = {} + + def __init__( + self, + bundle_name: str, + quantity: Quantity, + mapping: _FieldBundleIndexer = {}, + register_type: bool = False, + ): + """ + Initialize a bundle from a nD quantity. + + Dev note: current implementation limits to 4D inputs. + + Args: + bundle_name: name of the bundle, accessible via `name`. + quantity: data inputs as a nD array. + mapping: sparse dict of [name, index] to be able to call tracers by name. + register_type: boolean to register the type as part of initilization. + """ + if len(quantity.shape) != 4: + raise NotImplementedError("FieldBundle implementation restricted to 4D") + + self.name = bundle_name + self._quantity = quantity + self._indexer = mapping + if register_type: + # ToDo: extend the dims below to work with more than 4 dims + assert len(quantity.shape) == 4 + FieldBundleType.register(bundle_name, quantity.shape[3:]) + + def map(self, index: _DataDimensionIndex, name: str): + """Map a single `index` to ` name`""" + self._indexer[name] = index + + @property + def quantity(self) -> Quantity: + return self._quantity + + @property + def shape(self) -> tuple[int, ...]: + return self._quantity.shape + + def groupby(self, name: str) -> Quantity: + """Not implemented""" + raise NotImplementedError + + def __getattr__(self, name: str) -> Quantity: + """Allow to reference sub-array using `field.a_name`""" + if name not in self._indexer.keys(): + # This replicates as close possible the default behavior of getattr + # without breaking orchestration + return None # type: ignore + # ToDo: extend the dims below to work with more than 4 dims + assert len(self._quantity.data.shape) == 4 + return Quantity( + data=self._quantity.data[:, :, :, self.index(name)], + dims=self._quantity.dims[:-1], + units=self._quantity.units, + origin=self._quantity.origin[:-1], + extent=self._quantity.extent[:-1], + ) + + def index(self, name: str) -> int: + """Get index from name.""" + return self._indexer[name] + + @property + def __array_interface__(self): + """Memory interface for CPU.""" + return self._quantity.__array_interface__ + + @property + def __cuda_array_interface__(self): + """Memory interface for GPU memory as defined by cupy.""" + return self._quantity.__cuda_array_interface__ + + def __descriptor__(self) -> Any: + """Data descriptor for DaCe.""" + return self._quantity.__descriptor__() + + @staticmethod + def extend_3D_quantity_factory( + quantity_factory: QuantityFactory, + extra_dims: dict[str, int], + ) -> QuantityFactory: + """Create a nD quantity factory from a cartesian 3D factory. + + Args: + quantity_factory: Cartesian 3D factory. + extra_dims: dict of [name, size] of the data dimensions to add. + """ + new_factory = copy.copy(quantity_factory) + new_factory.set_extra_dim_lengths( + **{ + **extra_dims, + } + ) + return new_factory + + +@dataclass +class MarkupFieldBundleType: + """Markup a field bundle to delay specialization. + + Properties: + name: name of the future type to look into the registrar. + """ + + name: str + + +class FieldBundleType: + """Field Bundle Types to help with static sizing of Data Dimensions. + + Methods: + register: Register a type by sizing it's data dimensions + T: access any registered types for type hinting. + """ + + _field_type_registrar: dict[str, gtscript._FieldDescriptor] = {} + + @classmethod + def register( + cls, name: str, data_dims: tuple[int], dtype=Float + ) -> gtscript._FieldDescriptor: + """Register a name type by name by giving the size of it's data dimensions. + + The same type cannot be registered twice and will error out. + + Args: + name: Type name, to be re-used with `T`. + data_dims: tuple of int giving size of each data dimensions. + dtype: Inner data type, defaults to Float. + """ + if name in cls._field_type_registrar.keys(): + raise RuntimeError(f"Registering {name} a second time!") + cls._field_type_registrar[name] = gtscript.Field[ + gtscript.IJK, (dtype, (data_dims)) + ] + return cls._field_type_registrar[name] + + @classmethod + def T( + cls, name: str, do_markup: bool = True + ) -> gtscript._FieldDescriptor | MarkupFieldBundleType: + """ + Get registered type. + + Dev note: The markup feature is to allow early parsing (at file import) + to go ahead - while we will resolve the full type when calling the stencil. + + Args: + name: name of the type as registered via `register` + do_markup: if name not registered, markup for a future specialization + at stencil call time + """ + if name not in cls._field_type_registrar: + if do_markup: + return MarkupFieldBundleType(name) + raise RuntimeError(f"FieldBundle type {name} as not been registered!") + return cls._field_type_registrar[name] diff --git a/tests/quantity/test_fieldbundle.py b/tests/quantity/test_fieldbundle.py new file mode 100644 index 00000000..f30d81c1 --- /dev/null +++ b/tests/quantity/test_fieldbundle.py @@ -0,0 +1,77 @@ +from ndsl.boilerplate import get_factories_single_tile +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import PARALLEL, computation, interval +from ndsl.dsl.typing import FloatField +from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType + + +def assign_4d_field_stcl(field_4d: FieldBundleType.T("Tracers")): # type: ignore # noqa + with computation(PARALLEL), interval(...): + field_4d[0, 0, 0][1] = 63.63 + field_4d[0, 0, 0][3] = 63.63 + + +def assign_3d_field_stcl(field_3d: FloatField): + with computation(PARALLEL), interval(...): + field_3d = 121.121 + + +def test_field_bundle(): + # Grid & Factories + NX = 2 + NY = 2 + NZ = 2 + N4th = 5 + stencil_factory, quantity_factory = get_factories_single_tile(NX, NY, NZ, 1) + + # Type register + FieldBundleType.register("Tracers", (N4th,)) + + # Make stencils + assign_4d_field = stencil_factory.from_dims_halo( + func=assign_4d_field_stcl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + assign_3d_field = stencil_factory.from_dims_halo( + func=assign_3d_field_stcl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # "Input" data + new_quantity_factory = FieldBundle.extend_3D_quantity_factory( + quantity_factory, {"tracers": N4th} + ) + data = new_quantity_factory.ones([X_DIM, Y_DIM, Z_DIM, "tracers"], units="kg/g") + + # Build Bundle + tracers = FieldBundle( + bundle_name="tracers", + quantity=data, + mapping={"vapor": 0, "cloud": 2}, + ) + + # Test + tracers.quantity.field[:, :, :, :] = 48.4 + tracers.quantity.field[:, :, :, 2] = 21.21 + + assign_4d_field(tracers.quantity) + + assert (tracers.quantity.field[:, :, :, 0] == 48.4).all() + assert (tracers.quantity.field[:, :, :, 1] == 63.63).all() + assert (tracers.quantity.field[:, :, :, 2] == 21.21).all() + assert (tracers.quantity.field[:, :, :, 3] == 63.63).all() + assert (tracers.quantity.field[:, :, :, 4] == 48.4).all() + + tracers.vapor.field[:] = 1000.1000 + + assert (tracers.quantity.field[:, :, :, 0] == 1000.1000).all() + assert (tracers.quantity.field[:, :, :, 1] == 63.63).all() + + assign_3d_field(tracers.cloud) + + assert (tracers.cloud.field[:] == 121.121).all() + assert (tracers.quantity.field[:, :, :, 2] == tracers.cloud.field[:]).all() + + +if __name__ == "__main__": + test_field_bundle()