Fp16 constant weights

devito/core/operator.py

@@ -186,6 +186,14 @@ class BasicOperator(Operator):
     # ------------------------------------------------------------------
 
     INTERP_MODE = 'direct'
+
+    HALF_ARITH = False
+    """
+    Whether an Operator working in half precision carries the arithmetic there
+    too, rounding its literals and its FD weights to half. Off by default: half
+    is a storage format, and giving up the accuracy of the coefficients as well
+    is a mathematical choice rather than a consequence of it.
+    """
     """
     Default for the `sym_opt={'interp-mode': ...}` option. Controls how
     a product of fields living at different staggered locations is mapped
@@ -230,7 +238,8 @@ def _normalize_sym_kwargs(cls, **kwargs):
         the Operator. Returns the normalized `sym_options` dict.
         """
         so = dict(kwargs.get('sym_options', {}))
-        out = {'interp-mode': so.pop('interp-mode', cls.INTERP_MODE)}
+        out = {'interp-mode': so.pop('interp-mode', cls.INTERP_MODE),
+               'half-arith': so.pop('half-arith', cls.HALF_ARITH)}
 
         if so:
             raise InvalidOperator(

devito/finite_differences/differentiable.py

@@ -979,12 +979,17 @@ def __eq__(self, other):
                 self.name == other.name and
                 self.dimension == other.dimension and
                 self.indices == other.indices and
+                self.dtype is other.dtype and
                 self.weights == other.weights)
 
     __hash__ = sympy.Basic.__hash__
 
     def _hashable_content(self):
-        return (self.name, self.dimension, str(self.weights), self.scope)
+        # NOTE: `dtype` belongs here. The same coefficients at two precisions
+        # are two different arrays, and leaving it out has one of them fetched
+        # from the cache in place of the other
+        return (self.name, self.dimension, str(self.weights), self.scope,
+                np.dtype(self.dtype).name)
 
     @property
     def dimension(self):

devito/finite_differences/finite_difference.py

@@ -223,7 +223,8 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
         expand = expand(dim)
 
     if not expand and indices.expr is not None:
-        weights = Weights(name='w', dimensions=indices.free_dim, initvalue=weights)
+        weights = Weights(name='w', dimensions=indices.free_dim,
+                          initvalue=weights, dtype=expr.dtype)
 
         # Inject the StencilDimension
         # E.g. `x + i*h_x` into `f(x)` s.t. `f(x + i*h_x)`

devito/ir/cgen/printer.py

@@ -43,6 +43,12 @@ class BasePrinter(CodePrinter):
     _func_literals = {}
     _prec_literals = {np.float32: 'F', np.complex64: 'F'}
 
+    # Whether the arithmetic is carried at the Operator's own precision, even
+    # where that is narrower than `float32`. Off by default: a narrow dtype is
+    # a storage choice, and it takes a deliberate one to also give up the
+    # accuracy of the literals
+    _half_arith = False
+
     _qualifiers_mapper = {
         'is_extern': 'extern',
         'is_const': 'const',
@@ -79,6 +85,13 @@ def _prec(self, expr):
         dtype = sympy_dtype(expr, default=self.dtype)
         if dtype is None or np.issubdtype(dtype, np.integer):
             if any(isinstance(i, Float) for i in expr.atoms()):
+                # A real literal in an otherwise integer (or untyped)
+                # expression is emitted at the Operator's precision, floored at
+                # `float32` so that an integer default doesn't degrade it.
+                # A printer that has opted into narrow arithmetic keeps its own
+                # precision instead, rather than have the literal widen it
+                if self._half_arith and np.issubdtype(self.dtype, np.floating):
+                    return self.dtype
                 try:
                     return np.promote_types(self.dtype, np.float32).type
                 except np.exceptions.DTypePromotionError:
@@ -373,6 +386,17 @@ def _print_FieldFromComposite(self, expr):
     def _print_ListInitializer(self, expr):
         return f"{{{', '.join(self._print(i) for i in expr.params)}}}"
 
+    def initvalue(self, init, dtype):
+        """
+        Print the aggregate initializer `init` of an Array of type `dtype`.
+
+        Kept separate from `_print_ListInitializer` because a static
+        initializer, unlike an expression, cannot rely on implicit conversions:
+        some types (e.g. CUDA's `__half`) are only constructible from a literal
+        via a runtime call, which is illegal in that position.
+        """
+        return self._print(init)
+
     def _print_IndexedPointer(self, expr):
         base = self._print(expr.base)
         return f"{base}{''.join(f'[{self._print(i)}]' for i in expr.index)}"

devito/ir/iet/visitors.py

@@ -360,12 +360,21 @@ def _gen_value(self, obj, mode=1, masked=()):
         if obj.is_Array and obj.initvalue is not None and mode == 1:
             init = ListInitializer(obj.initvalue)
             if not obj._mem_constant or init.is_numeric:
-                value = c.Initializer(value, self.ccode(init))
+                value = c.Initializer(value, self._gen_initvalue(obj, init))
         elif obj.is_LocalObject and obj.initvalue is not None and mode == 1:
             value = c.Initializer(value, self.ccode(obj.initvalue))
 
         return value
 
+    def _gen_initvalue(self, obj, init):
+        """
+        Convert the aggregate initializer `init` of the Array `obj` into a C
+        string, delegating to the printer so that languages whose types cannot
+        be built from plain literals (e.g. CUDA's `__half`) can specialize it.
+        """
+        printer = get_printer(self.printer, obj.dtype)
+        return printer.initvalue(init, obj.dtype)
+
     def _gen_rettype(self, obj):
         try:
             return self._gen_value(obj, 0).typename

devito/operator/operator.py

@@ -811,6 +811,11 @@ def _soname(self):
 
     @cached_property
     def _printer(self):
+        # A Target may offer a second printer for Operators that have opted
+        # into carrying their precision into the arithmetic
+        if self._sym_options.get('half-arith'):
+            with suppress(AttributeError):
+                return self._Target.HalfArithPrinter
         return self._Target.Printer
 
     @cached_property

devito/passes/iet/errors.py

@@ -56,7 +56,12 @@ def _check_stability(iet, wmovs=(), rcompile=None, sregistry=None):
         else:
             continue
 
-        accumulator = Symbol(name='accumulator', dtype=f.dtype)
+        # The accumulator sums the whole field, so it is given at least single
+        # precision: in half precision it would overflow within a few thousand
+        # points and report an instability that isn't there
+        dtype = np.promote_types(f.dtype, np.float32).type
+
+        accumulator = Symbol(name='accumulator', dtype=dtype)
         eqns = [Eq(accumulator, 0.0),
                 Inc(accumulator, f.subs(f.time_dim, 0))]
         irs, byproduct = rcompile(eqns)