compiler: Misc tweaks to fully support FP16 storage/compute across all backends

devito/core/operator.py

@@ -202,10 +202,9 @@ class BasicOperator(Operator):
     See `examples/userapi/08_staggered_interp.ipynb` for a worked example.
     """
 
-    HALF_ARITH = False
+    ARITHM_MODE = 'default'
     """
-    Whether an Operator working in half precision carries out the arithmetic there
-    too, rounding its literals and its FD weights to half. Off by default.
+    Whether to use default arithmetic or FP16 arithmetic for half-precision fields.
     """
 
     @classmethod
@@ -237,13 +236,17 @@ def _normalize_sym_kwargs(cls, **kwargs):
         """
         so = dict(kwargs.get('sym_options', {}))
         out = {'interp-mode': so.pop('interp-mode', cls.INTERP_MODE),
-               'half-arith': so.pop('half-arith', cls.HALF_ARITH)}
+               'arithm-mode': so.pop('arithm-mode', cls.ARITHM_MODE)}
 
         if so:
             raise InvalidOperator(
                 f'Unrecognized symbolic options: [{", ".join(list(so))}]'
             )
 
+        mode = out['arithm-mode']
+        if not isinstance(mode, str) or mode not in ('default', 'fp16'):
+            raise InvalidOperator("`arithm-mode` must be 'default' or 'fp16'")
+
         return out
 
     @classmethod

devito/ir/iet/visitors.py

@@ -12,6 +12,7 @@
 from typing import Any, Generic, TypeVar
 
 import cgen as c
+import numpy as np
 from sympy import IndexedBase
 from sympy.core.function import Application
 
@@ -24,12 +25,12 @@
 from devito.ir.support.space import Backward
 from devito.ir.support.utils import IMask
 from devito.symbolics import (
-    FieldFromComposite, FieldFromPointer, ListInitializer, uxreplace
+    FieldFromComposite, FieldFromPointer, ListInitializer, sympy_dtype, uxreplace
 )
 from devito.symbolics.extended_dtypes import NoDeclStruct
 from devito.tools import (
-    GenericVisitor, as_tuple, c_restrict_void_p, filter_ordered, filter_sorted, flatten,
-    is_external_ctype, natural_sort_key
+    GenericVisitor, as_tuple, c_restrict_void_p, dtypes_vector_mapper, filter_ordered,
+    filter_sorted, flatten, is_external_ctype, natural_sort_key
 )
 from devito.types import (
     ArrayObject, CompositeObject, DeviceMap, Dimension, IndexedData, Pointer
@@ -261,8 +262,16 @@ def __init__(self, *args, printer=None, **kwargs):
     def ccode(self, expr, dtype=None, exact_prec=None):
         defaults = self.printer._default_settings
         dtype = defaults['dtype'] if dtype is None else dtype
+
         if exact_prec is None:
             exact_prec = defaults['exact_prec']
+
+        base = dtypes_vector_mapper.get_base_dtype(dtype, dtype)
+        if exact_prec and base == np.float16:
+            # E.g., `half_out = 0.9998*half_in + float_in` needs FP32
+            # literals despite its FP16 destination
+            dtype = sympy_dtype(expr, base=base)
+
         return get_printer(self.printer, dtype, exact_prec).doprint(expr, None)
 
     @property

devito/mpi/routines.py

@@ -1286,7 +1286,8 @@ def _arg_defaults(self, alias=None, args=None):
         super()._arg_defaults(alias, args=args)
 
         f = alias or self.target.c0
-        neighborhood = f.grid.distributor.neighborhood
+        # Internal Arrays may retain a serialized, non-distributed Grid
+        neighborhood = args.grid.distributor.neighborhood
 
         for i, halo in enumerate(self.halos):
             entry = self.value[i]

devito/operator/operator.py

@@ -822,6 +822,10 @@ def arguments(self, **kwargs):
 
     # Code generation and JIT compilation
 
+    def _signature_items(self):
+        # Retain literal precision and the complete toolchain ABI in JIT identity
+        return (str(self.ccode), repr(self._compiler.abi_id()))
+
     @cached_property
     def _soname(self):
         """A unique name for the shared object resulting from JIT compilation."""
@@ -831,7 +835,7 @@ def _soname(self):
     def _printer(self):
         # A Target may offer a second printer e.g. for Operators
         # using half-precision arithmetic
-        if self._sym_options.get('half-arith'):
+        if self._sym_options.get('arithm-mode') == 'fp16':
             with suppress(AttributeError):
                 return self._Target.HalfArithPrinter
         return self._Target.Printer

devito/passes/clusters/fusion.py

@@ -248,9 +248,14 @@ def _toposort(self, cgroups, prefix):
 
         # Are there any ClusterGroups that could potentially be topologically
         # reordered? If not, do not waste time
-        counter = Counter(self._key(cg).strict for cg in cgroups)
-        if len(counter.most_common()) == 1 or \
-           not any(v > 1 for it, v in counter.most_common()):
+        keys = [self._key(cg) for cg in cgroups]
+        if len({k.full for k in keys}) == 1:
+            return ClusterGroup(cgroups, prefix)
+
+        # Same strict keys can still hide different scheduling preferences.
+        # E.g., `r0, shared0, r1, shared1` may reorder to fuse `r0, r1`.
+        counter = Counter(k.strict for k in keys)
+        if not any(v > 1 for v in counter.values()):
             return ClusterGroup(cgroups, prefix)
 
         dag = self._build_dag(cgroups, prefix)

devito/passes/iet/definitions.py

@@ -201,7 +201,7 @@ def _alloc_host_array_on_high_bw_mem(self, site, obj, storage, *args):
 
         storage.update(obj, site, allocs=tuple(allocs), frees=free, efuncs=efuncs)
 
-    def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, *args):
+    def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, *args, **kwargs):
         """
         Allocate a local Array in the host high bandwidth memory.
         """
@@ -471,7 +471,7 @@ def _inject_definitions(self, iet, storage):
         return processed, flatten(efuncs)
 
     @iet_pass
-    def place_definitions(self, iet, globs=None, **kwargs):
+    def place_definitions(self, iet, globs=None, alignments=None, **kwargs):
         """
         Create a new IET where all symbols have been declared, allocated, and
         deallocated in one or more memory spaces.
@@ -507,7 +507,9 @@ def place_definitions(self, iet, globs=None, **kwargs):
                     if i.is_transient:
                         self._alloc_bundle_struct_on_high_bw_mem(iet, i, storage)
                     elif i._mem_local:
-                        self._alloc_local_array_on_high_bw_mem(iet, i, storage)
+                        self._alloc_local_array_on_high_bw_mem(
+                            iet, i, storage, alignments=alignments
+                        )
                     elif i._mem_mapped:
                         self._alloc_mapped_array_on_high_bw_mem(iet, i, storage)
                 elif i._mem_stack:
@@ -518,7 +520,9 @@ def place_definitions(self, iet, globs=None, **kwargs):
                     if i._mem_host:
                         self._alloc_host_array_on_high_bw_mem(iet, i, storage)
                     elif i._mem_local:
-                        self._alloc_local_array_on_high_bw_mem(iet, i, storage)
+                        self._alloc_local_array_on_high_bw_mem(
+                            iet, i, storage, alignments=alignments
+                        )
                     else:
                         self._alloc_mapped_array_on_high_bw_mem(iet, i, storage)
                 elif i._mem_stack:
@@ -579,11 +583,22 @@ def place_casts(self, iet, **kwargs):
 
         return iet, {}
 
+    def derive_alignments(self, graph):
+        """
+        Choose allocation padding from alignment requirements across `graph`.
+
+        Backends may inspect all kernels and helpers before deciding how many
+        elements to insert before each field's logical allocation origin.
+        An empty mapper retains the backend's ordinary allocation behavior.
+        """
+        return {}
+
     def process(self, graph):
         """
         Apply the `place_definitions` and `place_casts` passes.
         """
-        self.place_definitions(graph, globs=set())
+        alignments = self.derive_alignments(graph)
+        self.place_definitions(graph, globs=set(), alignments=alignments)
         self.place_casts(graph)
 
 
@@ -596,7 +611,7 @@ def __init__(self, options=None, **kwargs):
 
         super().__init__(options=options, **kwargs)
 
-    def _alloc_local_array_on_high_bw_mem(self, site, obj, storage):
+    def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, **kwargs):
         """
         Allocate a local Array in the device high bandwidth memory.
         """
@@ -673,7 +688,7 @@ def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm,
         storage.update(obj, site, maps=mmap, unmaps=unmap, efuncs=efuncs)
 
     @iet_pass
-    def place_transfers(self, iet, data_movs=None, ctx=None, **kwargs):
+    def place_transfers(self, iet, data_movs=None, alignments=None, **kwargs):
         """
         Create a new IET with host-device data transfers. This requires mapping
         symbols to the suitable memory spaces.
@@ -695,14 +710,14 @@ def place_transfers(self, iet, data_movs=None, ctx=None, **kwargs):
                 self._map_array_on_high_bw_mem(iet, i, storage)
             else:
                 self._map_function_on_high_bw_mem(
-                    iet, i, storage, devicerm, ctx=ctx
+                    iet, i, storage, devicerm, alignments=alignments
                 )
         for i in filter_sorted(reads - writes):
             if i.is_Array:
                 self._map_array_on_high_bw_mem(iet, i, storage)
             else:
                 self._map_function_on_high_bw_mem(
-                    iet, i, storage, devicerm, read_only=True, ctx=ctx
+                    iet, i, storage, devicerm, read_only=True, alignments=alignments
                 )
 
         iet, efuncs = self._inject_definitions(iet, storage)
@@ -735,8 +750,9 @@ def process(self, graph):
         """
         Apply the `place_transfers`, `place_definitions` and `place_casts` passes.
         """
-        self.place_transfers(graph, data_movs=graph.data_movs)
-        self.place_definitions(graph, globs=set())
+        alignments = self.derive_alignments(graph)
+        self.place_transfers(graph, data_movs=graph.data_movs, alignments=alignments)
+        self.place_definitions(graph, globs=set(), alignments=alignments)
         self.place_devptr(graph)
         self.place_bundling(graph, writes_input=graph.writes_input)
         self.place_casts(graph)

devito/passes/iet/dtypes.py

@@ -12,7 +12,7 @@
 
 @iet_pass
 def _complex_includes(iet: Callable, langbb: type[LangBB], compiler: Compiler,
-                      sregistry: SymbolRegistry) -> tuple[Callable, dict]:
+                      sregistry: SymbolRegistry, **kwargs) -> tuple[Callable, dict]:
     """
     Includes complex arithmetic headers for the given language, if needed.
     """
@@ -48,9 +48,9 @@ def lower_dtypes(graph: Callable,
                  compiler: Compiler = None,
                  sregistry: SymbolRegistry = None, **kwargs) -> tuple[Callable, dict]:
     """
-    Lowers float16 scalar types to pointers since we can't directly pass their
-    value. Also includes headers for complex arithmetic if needed.
+    Lower dtype-specific parameters and arithmetic, and include required headers.
     """
 
     for dtype_pass in dtype_passes:
-        dtype_pass(graph, langbb=langbb, compiler=compiler, sregistry=sregistry)
+        dtype_pass(graph, langbb=langbb, compiler=compiler, sregistry=sregistry,
+                   **kwargs)

devito/symbolics/extended_sympy.py

@@ -15,7 +15,8 @@
 from devito.finite_differences.elementary import Max, Min
 from devito.tools import (  # noqa
     Bunch, Pickable, as_tuple, ctypes_to_cstr, ctypes_vector_mapper, double2, double3,
-    double4, dtype_to_ctype, float2, float3, float4, int2, int3, int4, is_integer
+    double4, dtype_to_ctype, extract_dtype, float2, float3, float4, int2, int3, int4,
+    is_integer
 )
 from devito.types import Symbol
 from devito.types.basic import Basic
@@ -543,6 +544,10 @@ def __new__(cls, base, index, **kwargs):
     def base(self):
         return self._base
 
+    @property
+    def dtype(self):
+        return self.base.dtype
+
     @property
     def index(self):
         return self._index
@@ -660,6 +665,12 @@ def __new__(cls, value, step, **kwargs):
 
         return super().__new__(cls, value, step, **kwargs)
 
+    def _eval_is_integer(self):
+        dtype = extract_dtype(self.value)
+        if (self.value.is_integer or
+                isinstance(dtype, type) and issubclass(dtype, np.integer)):
+            return True
+
     @property
     def value(self):
         return self.args[0]

devito/symbolics/inspection.py

@@ -290,7 +290,7 @@ def has_integer_args(*args):
     if len(args) == 1:
         try:
             return np.issubdtype(args[0].dtype, np.integer)
-        except AttributeError:
+        except (AttributeError, TypeError):
             return is_integer(args[0])
 
     res = True

devito/types/array.py

@@ -700,7 +700,8 @@ def indices(self):
     @property
     def dtype(self):
         try:
-            return self.function.c0.dtype
+            dtype = self.function.c0.dtype
+            return dtypes_vector_mapper.get_base_dtype(dtype, dtype)
         except AttributeError:
             # Vector-component access over a scalar symbol, e.g. a float4 register.
             if self.function.is_Symbol:

devito/types/basic.py

@@ -1799,6 +1799,12 @@ def __new__(cls, label, shape, function=None):
 
     func = Pickable._rebuild
 
+    def __getstate__(self):
+        state = dict(self.__dict__)
+        # Recreate process-local ctypes classes from the Function's dtype
+        state.pop('_C_ctype', None)
+        return state
+
     @sympy.cacheit
     def sort_key(self, order=None):
         class_key, args, exp, coeff = super().sort_key(order=order)

devito/types/misc.py

@@ -1,4 +1,4 @@
-from ctypes import c_double, c_void_p
+from ctypes import POINTER, c_double
 
 import numpy as np
 import sympy
@@ -201,15 +201,36 @@ def _defines(self):
 
 class Pointer(LocalObject, sympy.Expr):
 
-    __rkwargs__ = LocalObject.__rkwargs__ + ('dtype',)
+    """
+    A Pointer is defined by the pointee base type and its indirections.
+
+    Unlike plain LocalObjects, Pointers can be used within arbitrary expressions.
 
-    def __init__(self, *args, dtype=c_void_p, **kwargs):
+    Examples
+    --------
+    `Pointer('p', dtype_base=c_float, stars='**')` declares `float **p`.
+    """
+
+    __rkwargs__ = LocalObject.__rkwargs__ + ('dtype_base', 'stars')
+
+    def __init__(self, *args, dtype_base=None, stars='*', **kwargs):
         super().__init__(*args, **kwargs)
-        self._dtype = dtype
+        if not isinstance(stars, str) or any(i != '*' for i in stars):
+            raise ValueError('`stars` must be a string of zero or more `*` characters')
+        self.dtype_base = dtype_base
+        self.stars = stars
+
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.stars,)
 
     @property
     def dtype(self):
-        return self._dtype
+        ctype = self.dtype_base
+        if isinstance(ctype, CustomDtype):
+            return CustomDtype(ctype.name, ctype.template, ctype.modifier + self.stars)
+        for _ in self.stars:
+            ctype = POINTER(ctype)
+        return ctype
 
 
 class Indirection(Symbol):