compiler: add scalar type option

Author: mloubout

devito/core/cpu.py

@@ -32,6 +32,9 @@ def _normalize_kwargs(cls, **kwargs):
         o['mpi'] = oo.pop('mpi')
         o['parallel'] = o['openmp']  # Backwards compatibility
 
+        # Minimum scalar type
+        o['scalar-min-type'] = oo.pop('scalar-min-type', cls.SCALAR_MIN_TYPE)
+
         # Buffering
         o['buf-async-degree'] = oo.pop('buf-async-degree', None)
 

devito/core/gpu.py

@@ -40,6 +40,9 @@ def _normalize_kwargs(cls, **kwargs):
         o['mpi'] = oo.pop('mpi')
         o['parallel'] = True
 
+        # Minimum scalar type
+        o['scalar-min-type'] = oo.pop('scalar-min-type', cls.SCALAR_MIN_TYPE)
+
         # Buffering
         o['buf-async-degree'] = oo.pop('buf-async-degree', None)
 

devito/core/operator.py

@@ -1,6 +1,8 @@
 from collections.abc import Iterable
 from functools import cached_property
 
+import numpy as np
+
 from devito.core.autotuning import autotune
 from devito.exceptions import InvalidOperator
 from devito.ir import FindSymbols
@@ -67,6 +69,11 @@ class BasicOperator(Operator):
     intensity of the generated kernel.
     """
 
+    SCALAR_MIN_TYPE = np.float16
+    """
+    Minimum datatype for a scalar alias for a common sub-expression or cire temp.
+    """
+
     PAR_COLLAPSE_NCORES = 4
     """
     Use a collapse clause if the number of available physical cores is greater

devito/passes/clusters/aliases.py

@@ -109,11 +109,11 @@ class CireTransformer:
     def __init__(self, sregistry, options, platform):
         self.sregistry = sregistry
         self.platform = platform
-
         self.opt_minstorage = options['min-storage']
         self.opt_rotate = options['cire-rotate']
         self.opt_ftemps = options['cire-ftemps']
         self.opt_mingain = options['cire-mingain']
+        self.opt_mindtype = options['scalar-min-type']
         self.opt_multisubdomain = True
 
     def _aliases_from_clusters(self, clusters, exclude, meta):
@@ -143,7 +143,7 @@ def _aliases_from_clusters(self, clusters, exclude, meta):
 
         # Schedule -> [Clusters]_k
         processed, subs = lower_schedule(schedule, meta, self.sregistry,
-                                         self.opt_ftemps)
+                                         self.opt_ftemps, self.opt_mindtype)
 
         # [Clusters]_k -> [Clusters]_k (optimization)
         if self.opt_multisubdomain:
@@ -831,7 +831,7 @@ def optimize_schedule_rotations(schedule, sregistry):
     return schedule.rebuild(*processed, rmapper=rmapper)
 
 
-def lower_schedule(schedule, meta, sregistry, ftemps):
+def lower_schedule(schedule, meta, sregistry, ftemps, mindtype):
     """
     Turn a Schedule into a sequence of Clusters.
     """
@@ -849,7 +849,8 @@ def lower_schedule(schedule, meta, sregistry, ftemps):
         # This prevents cases such as `floor(a*b)` with `a` and `b` floats
         # that would creat a temporary `int r = b` leading to erronous
         # numerical results
-        dtype = sympy_dtype(pivot, base=meta.dtype)
+        mindtype = None if writeto else mindtype
+        dtype = sympy_dtype(pivot, base=meta.dtype, smin=mindtype)
 
         if writeto:
             # The Dimensions defining the shape of Array

devito/passes/clusters/cse.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
 from functools import cached_property, singledispatch
 
+import numpy as np
 import sympy
 from sympy import Add, Function, Indexed, Mul, Pow
 try:
@@ -69,11 +70,12 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     """
     min_cost = options['cse-min-cost']
     mode = options['cse-algo']
+    mindtype = np.promote_types(options['scalar-min-type'], cluster.dtype).type
 
     if cluster.is_fence:
         return cluster
 
-    make = lambda: CTemp(name=sregistry.make_name(), dtype=cluster.dtype)
+    make = lambda: CTemp(name=sregistry.make_name(), dtype=mindtype)
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)
 

devito/symbolics/inspection.py

@@ -298,7 +298,7 @@ def has_integer_args(*args):
     return res
 
 
-def sympy_dtype(expr, base=None, default=None):
+def sympy_dtype(expr, base=None, default=None, smin=None):
     """
     Infer the dtype of the expression.
     """
@@ -322,4 +322,7 @@ def sympy_dtype(expr, base=None, default=None):
         else:
             dtype = np.promote_types(dtype, np.complex64).type
 
+    if smin is not None and not np.issubdtype(dtype, np.integer):
+        dtype = np.promote_types(dtype, smin).type
+
     return dtype