From 004ed3c0a25d3b0761cbfe14dd157c1a23549b7d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 16 May 2025 09:37:17 -0400 Subject: [PATCH 1/3] Field Bundle functional code for 4D array Deferred type hint for data dimensions Unit tests --- ndsl/__init__.py | 2 +- ndsl/dsl/stencil.py | 9 ++ ndsl/quantity/__init__.py | 7 +- ndsl/quantity/field_bundle.py | 141 +++++++++++++++++++++++++++++ tests/quantity/test_fieldbundle.py | 77 ++++++++++++++++ 5 files changed, 233 insertions(+), 3 deletions(-) create mode 100644 ndsl/quantity/field_bundle.py create mode 100644 tests/quantity/test_fieldbundle.py diff --git a/ndsl/__init__.py b/ndsl/__init__.py index 18c96f3e..1667f8a5 100644 --- a/ndsl/__init__.py +++ b/ndsl/__init__.py @@ -28,7 +28,7 @@ from .performance.collector import NullPerformanceCollector, PerformanceCollector from .performance.profiler import NullProfiler, Profiler from .performance.report import Experiment, Report, TimeReport -from .quantity import Quantity +from .quantity import FieldBundle, FieldBundleType, Quantity 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..ef850844 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 know + # 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..b498eafd 100644 --- a/ndsl/quantity/__init__.py +++ b/ndsl/quantity/__init__.py @@ -1,9 +1,12 @@ -from ndsl.quantity.metadata import QuantityHaloSpec, QuantityMetadata -from ndsl.quantity.quantity import Quantity +from .field_bundle import FieldBundle, FieldBundleType +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..005e01f6 --- /dev/null +++ b/ndsl/quantity/field_bundle.py @@ -0,0 +1,141 @@ +import copy +from typing import Any, Dict, Tuple +from ndsl.quantity.quantity import Quantity +from ndsl.initialization.allocator import QuantityFactory +from ndsl.dsl.typing import Float + +import gt4py.cartesian.gtscript as gtscript # isort: skip + +from dataclasses import dataclass + +# 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 + + WARNING: The present implementation only allows for 4D array. + """ + + _quantity: Quantity + _indexer: _FieldBundleIndexer = {} + + def __init__( + self, + bundle_name: str, + quantity: Quantity, + mapping: _FieldBundleIndexer = {}, + do_register_type: bool = False, + ): + if len(quantity.shape) != 4: + raise NotImplementedError("FieldBundle implementation restricted to 4D") + + self.name = bundle_name + self._quantity = quantity + self._indexer = mapping + if do_register_type: + FieldBundleType.register(bundle_name, quantity.shape[3:]) + + def map(self, index: _DataDimensionIndex, name: str): + 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: + raise NotImplementedError + + def __getattr__(self, name: str) -> Quantity: + 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._indexer[name]], + dims=self._quantity.dims[:-1], + units=self._quantity.units, + origin=self._quantity.origin[:-1], + extent=self._quantity.extent[:-1], + ) + + @property + def __array_interface__(self): + return self._quantity.__array_interface__ + + @property + def __cuda_array_interface__(self): + return self._quantity.__array_interface__ + + def __descriptor__(self) -> Any: + return self._quantity.__descriptor__() + + @staticmethod + def extend_3D_quantity_factory( + quantity_factory: QuantityFactory, + extra_dims: dict[str, int], + ) -> QuantityFactory: + new_factory = copy.copy(quantity_factory) + new_factory.set_extra_dim_lengths( + **{ + **extra_dims, + } + ) + return new_factory + + def index(self, name: str) -> int: + return self._indexer[name] + + +@dataclass +class MarkupFieldBundleType: + """Markup a field bundle to delay specialization + + Properties: + name: name of the future type to look into the registrat + """ + + 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]) -> gtscript._FieldDescriptor: + 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, (Float, (data_dims)) + ] + return cls._field_type_registrar[name] + + @classmethod + def T( + cls, name: str, do_markup: bool = True + ) -> gtscript._FieldDescriptor | MarkupFieldBundleType: + if name not in cls._field_type_registrar: + if do_markup: + return MarkupFieldBundleType(name) + else: + 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..8cc3abaf --- /dev/null +++ b/tests/quantity/test_fieldbundle.py @@ -0,0 +1,77 @@ +from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import FloatField +from ndsl.dsl.gt4py import computation, PARALLEL, interval +from ndsl.boilerplate import get_factories_single_tile + + +def assign_4d_field_stcl(field_4d: FieldBundleType.T("Tracers")): # type: ignore + 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() From eda8691ae226dddd0563b2cfd53855da00f11fff Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 16 May 2025 09:54:08 -0400 Subject: [PATCH 2/3] Fix imports --- ndsl/__init__.py | 3 ++- ndsl/quantity/__init__.py | 1 - ndsl/quantity/field_bundle.py | 8 +++++--- tests/quantity/test_fieldbundle.py | 8 ++++---- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ndsl/__init__.py b/ndsl/__init__.py index 1667f8a5..f6bb6117 100644 --- a/ndsl/__init__.py +++ b/ndsl/__init__.py @@ -28,7 +28,8 @@ from .performance.collector import NullPerformanceCollector, PerformanceCollector from .performance.profiler import NullProfiler, Profiler from .performance.report import Experiment, Report, TimeReport -from .quantity import FieldBundle, FieldBundleType, Quantity +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/quantity/__init__.py b/ndsl/quantity/__init__.py index b498eafd..43751528 100644 --- a/ndsl/quantity/__init__.py +++ b/ndsl/quantity/__init__.py @@ -1,4 +1,3 @@ -from .field_bundle import FieldBundle, FieldBundleType from .metadata import QuantityHaloSpec, QuantityMetadata from .quantity import Quantity diff --git a/ndsl/quantity/field_bundle.py b/ndsl/quantity/field_bundle.py index 005e01f6..05a0f3ad 100644 --- a/ndsl/quantity/field_bundle.py +++ b/ndsl/quantity/field_bundle.py @@ -1,12 +1,14 @@ import copy +from dataclasses import dataclass from typing import Any, Dict, Tuple -from ndsl.quantity.quantity import Quantity -from ndsl.initialization.allocator import QuantityFactory + 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 -from dataclasses import dataclass # ToDo: This is 4th dimensions restricted. We need a concept # of data dimensions index here to be able to extend to N dimensions diff --git a/tests/quantity/test_fieldbundle.py b/tests/quantity/test_fieldbundle.py index 8cc3abaf..f30d81c1 100644 --- a/tests/quantity/test_fieldbundle.py +++ b/tests/quantity/test_fieldbundle.py @@ -1,11 +1,11 @@ -from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType +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.dsl.gt4py import computation, PARALLEL, interval -from ndsl.boilerplate import get_factories_single_tile +from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType -def assign_4d_field_stcl(field_4d: FieldBundleType.T("Tracers")): # type: ignore +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 From 0c05cc63bb13f0767fae77d202db51c044ef1f8f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 16 May 2025 12:44:44 -0400 Subject: [PATCH 3/3] Fix `dtype` lacking in type registration Modernizing type hint Docstrings --- ndsl/dsl/stencil.py | 4 +- ndsl/quantity/field_bundle.py | 85 +++++++++++++++++++++++++++-------- 2 files changed, 68 insertions(+), 21 deletions(-) diff --git a/ndsl/dsl/stencil.py b/ndsl/dsl/stencil.py index ef850844..9260bd29 100644 --- a/ndsl/dsl/stencil.py +++ b/ndsl/dsl/stencil.py @@ -336,8 +336,8 @@ def __init__( ): block_waiting_for_compilation(MPI.COMM_WORLD, compilation_config) - # Field Bundle might have dropped a placeholder type that we know - # have to resolve to the proper type + # 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( diff --git a/ndsl/quantity/field_bundle.py b/ndsl/quantity/field_bundle.py index 05a0f3ad..a2749361 100644 --- a/ndsl/quantity/field_bundle.py +++ b/ndsl/quantity/field_bundle.py @@ -1,6 +1,6 @@ import copy from dataclasses import dataclass -from typing import Any, Dict, Tuple +from typing import Any from ndsl.dsl.typing import Float from ndsl.initialization.allocator import QuantityFactory @@ -13,14 +13,14 @@ # 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] +_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 + which provides a way to type hint parameters for stencils in the `gtscript`. WARNING: The present implementation only allows for 4D array. """ @@ -33,18 +33,32 @@ def __init__( bundle_name: str, quantity: Quantity, mapping: _FieldBundleIndexer = {}, - do_register_type: bool = False, + 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 do_register_type: + 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 @@ -52,13 +66,15 @@ def quantity(self) -> Quantity: return self._quantity @property - def shape(self) -> Tuple[int, ...]: + 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 @@ -66,22 +82,29 @@ def __getattr__(self, name: str) -> Quantity: # 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._indexer[name]], + 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): - return self._quantity.__array_interface__ + """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 @@ -89,6 +112,12 @@ 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( **{ @@ -97,37 +126,45 @@ def extend_3D_quantity_factory( ) return new_factory - def index(self, name: str) -> int: - return self._indexer[name] - @dataclass class MarkupFieldBundleType: - """Markup a field bundle to delay specialization + """Markup a field bundle to delay specialization. Properties: - name: name of the future type to look into the registrat + 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 + """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 + 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]) -> gtscript._FieldDescriptor: + 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, (Float, (data_dims)) + gtscript.IJK, (dtype, (data_dims)) ] return cls._field_type_registrar[name] @@ -135,9 +172,19 @@ def register(cls, name: str, data_dims: tuple[int]) -> gtscript._FieldDescriptor 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) - else: - raise RuntimeError(f"FieldBundle type {name} as not been registered!") + raise RuntimeError(f"FieldBundle type {name} as not been registered!") return cls._field_type_registrar[name]