Merge pull request #2782 from devitocodes/fd-eval-add

api: fix staggered evaluation for add/mul

Author: mloubout

devito/arch/archinfo.py

@@ -619,18 +619,12 @@ def get_m1_llvm_path(language):
 
 @memoized_func
 def check_cuda_runtime():
-    libnames = ('libcudart.so', 'libcudart.dylib', 'cudart.dll')
-    for libname in libnames:
-        try:
-            cuda = ctypes.CDLL(libname)
-        except OSError:
-            continue
-        else:
-            break
-    else:
+    libname = ctypes.util.find_library("cudart")
+    if not libname:
         warning("Unable to check compatibility of NVidia driver and runtime")
         return
 
+    cuda = ctypes.CDLL(libname)
     driver_version = ctypes.c_int()
     runtime_version = ctypes.c_int()
 
@@ -1069,6 +1063,32 @@ def march(self):
                 return 'tesla'
         return None
 
+    @cached_property
+    def max_shm_per_block(self):
+        """
+        Get the maximum amount of shared memory per thread block
+        """
+        # Load libcudart
+        libname = ctypes.util.find_library("cudart")
+        if not libname:
+            return 64 * 1024  # 64 KB default
+        lib = ctypes.CDLL(libname)
+
+        cudaDevAttrMaxSharedMemoryPerBlockOptin = 97
+        # get current device
+        dev = ctypes.c_int()
+        lib.cudaGetDevice(ctypes.byref(dev))
+
+        # query attribute
+        value = ctypes.c_int()
+        lib.cudaDeviceGetAttribute(
+            ctypes.byref(value),
+            ctypes.c_int(cudaDevAttrMaxSharedMemoryPerBlockOptin),
+            dev
+        )
+
+        return value.value
+
     def supports(self, query, language=None):
         if language != 'cuda':
             return False
@@ -1125,6 +1145,8 @@ class AmdDevice(Device):
 
     max_mem_trans_nbytes = 256
 
+    max_shm_per_block = 64*1024  # 64 KB
+
     @cached_property
     def march(cls):
         # TODO: this corresponds to Vega, which acts as the fallback `march`

devito/finite_differences/derivative.py

@@ -335,20 +335,6 @@ def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs):
             except AttributeError:
                 raise TypeError("fd_order incompatible with dimensions") from None
 
-        if isinstance(self.expr, Derivative):
-            # In case this was called on a perfect cross-derivative `u.dxdy`
-            # we need to propagate the call to the nested derivative
-            rkwe = dict(rkw)
-            rkwe.pop('weights', None)
-            if 'x0' in rkwe:
-                rkwe['x0'] = self._filter_dims(self.expr._filter_dims(rkw['x0']),
-                                               neg=True)
-            if fd_order is not None:
-                fdo = self.expr._filter_dims(_fd_order)
-                if fdo:
-                    rkwe['fd_order'] = fdo
-            rkw['expr'] = self.expr(**rkwe)
-
         if fd_order is not None:
             rkw['fd_order'] = self._filter_dims(_fd_order, as_tuple=True)
 
@@ -530,9 +516,12 @@ def _eval_at(self, func):
             # it into `u(x + h_x/2).dx` and `v(x).dx`, since they require
             # different FD indices
             mapper = as_mapper(self.expr._args_diff, lambda i: i.staggered)
+            if len(mapper) == 1:
+                # All terms have the same staggering, we can use expr as is
+                return self._rebuild(self.expr, **rkw)
             args = [self.expr.func(*v) for v in mapper.values()]
             args.extend([a for a in self.expr.args if a not in self.expr._args_diff])
-            args = [self._rebuild(a, **rkw) for a in args]
+            args = [self._rebuild(a)._eval_at(func) for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
             # For Mul, We treat the basic case `u(x + h_x/2) * v(x) which is what appear

devito/finite_differences/differentiable.py

@@ -259,6 +259,12 @@ def __rfloordiv__(self, other):
         from .elementary import floor
         return floor(other / self)
 
+    def _inv(self, ref, safe=False):
+        if safe:
+            return SafeInv(self, ref or self)
+        else:
+            return 1 / self
+
     def __mod__(self, other):
         return Mod(self, other)
 
@@ -964,7 +970,11 @@ def _subs(self, old, new, **hints):
 
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
-    pass
+
+    def _eval_at(self, func):
+        # Like EvalDerivative, a DiffDerivative must have already been evaluated
+        # at a valid x0 and should not be re-evaluated at a different location
+        return self
 
 
 # SymPy args ordering is the same for Derivatives and IndexDerivatives
@@ -1012,6 +1022,11 @@ def _new_rawargs(self, *args, **kwargs):
         kwargs.pop('is_commutative', None)
         return self.func(*args, **kwargs)
 
+    def _eval_at(self, func):
+        # An EvalDerivative must have already been evaluated at a valid x0
+        # and should not be re-evaluated at a different location
+        return self
+
 
 class diffify:
 
@@ -1145,13 +1160,8 @@ def _(expr, x0, **kwargs):
 
 @interp_for_fd.register(AbstractFunction)
 def _(expr, x0, **kwargs):
-    from devito.finite_differences.derivative import Derivative
-    x0_expr = {d: v for d, v in x0.items() if v is not expr.indices_ref[d]}
-    if expr.is_parameter:
-        return expr
-    elif x0_expr:
-        dims = tuple((d, 0) for d in x0_expr)
-        fd_o = tuple([expr.interp_order]*len(dims))
-        return Derivative(expr, *dims, fd_order=fd_o, x0=x0_expr)
+    x0_expr = {d: v for d, v in x0.items() if v.has(d)}
+    if x0_expr:
+        return expr.subs({expr.indices[d]: v for d, v in x0_expr.items()})
     else:
         return expr

devito/passes/clusters/implicit.py

@@ -175,13 +175,17 @@ def callback(self, clusters, prefix):
 
             # Make sure the "implicit expressions" are scheduled in
             # the innermost loop such that the thicknesses can be computed
-            edims = set(retrieve_dimensions(mapper.values(), deep=True))
-            if dim not in edims or not edims.issubset(prefix.dimensions):
+            def key(tkn):
+                edims = set(retrieve_dimensions(tkn, deep=True))
+                return dim._defines & edims and edims.issubset(prefix.dimensions)
+
+            mapper = {k: v for k, v in mapper.items() if key(v)}
+            if not mapper:
                 continue
 
             found[d.functions].clusters.append(c)
             found[d.functions].mapper = reduce(found[d.functions].mapper,
-                                               mapper, edims, prefix)
+                                               mapper, {dim}, prefix)
 
         # Turn the reduced mapper into a list of equations
         processed = []
@@ -259,7 +263,11 @@ def reduce(m0, m1, edims, prefix):
     else:
         func = min
 
-    key = lambda i: i.indices[d]
+    def key(i):
+        try:
+            return i.indices[d]
+        except (KeyError, AttributeError):
+            return i
 
     mapper = {}
     for k, e in m1.items():

devito/passes/iet/misc.py

@@ -6,7 +6,6 @@
 
 from devito.finite_differences import Max, Min
 from devito.finite_differences.differentiable import SafeInv
-from devito.logger import warning
 from devito.ir import (Any, Forward, DummyExpr, Iteration, EmptyList, Prodder,
                        FindApplications, FindNodes, FindSymbols, Transformer,
                        Uxreplace, filter_iterations, retrieve_iteration_tree,
@@ -155,7 +154,7 @@ def _generate_macros(iet, tracker=None, langbb=None, printer=CPrinter, **kwargs)
                      for define, expr in headers)
 
     # Generate Macros from higher-level SymPy objects
-    mheaders, includes = _generate_macros_math(iet, langbb=langbb)
+    mheaders, includes = _generate_macros_math(iet, langbb=langbb, printer=printer)
     includes = sorted(includes, key=str)
     headers.extend(sorted(mheaders, key=str))
 
@@ -199,25 +198,25 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs):
     return iet
 
 
-def _generate_macros_math(iet, langbb=None):
+def _generate_macros_math(iet, langbb=None, printer=CPrinter):
     headers = []
     includes = []
     for i in FindApplications().visit(iet):
-        header, include = _lower_macro_math(i, langbb)
+        header, include = _lower_macro_math(i, langbb, printer)
         headers.extend(header)
         includes.extend(include)
 
     return headers, set(includes) - {None}
 
 
 @singledispatch
-def _lower_macro_math(expr, langbb):
+def _lower_macro_math(expr, langbb, printer):
     return (), {}
 
 
 @_lower_macro_math.register(Min)
 @_lower_macro_math.register(sympy.Min)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     if has_integer_args(*expr.args):
         return (('MIN(a,b)', ('(((a) < (b)) ? (a) : (b))')),), {}
     else:
@@ -226,23 +225,28 @@ def _(expr, langbb):
 
 @_lower_macro_math.register(Max)
 @_lower_macro_math.register(sympy.Max)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     if has_integer_args(*expr.args):
         return (('MAX(a,b)', ('(((a) > (b)) ? (a) : (b))')),), {}
     else:
         return (), as_tuple(langbb.get('header-math'))
 
 
 @_lower_macro_math.register(SafeInv)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     try:
-        eps = np.finfo(expr.base.dtype).resolution**2
-    except ValueError:
-        warning(f"dtype not recognized in SafeInv for {expr.base}, assuming float32")
-        eps = np.finfo(np.float32).resolution**2
-    b = Cast('b', dtype=np.float32)
+        dtype = expr.base.dtype
+        eps = np.finfo(dtype).resolution**2
+    except (AttributeError, ValueError):
+        dtype = np.float32
+        eps = np.finfo(dtype).resolution**2
+
+    b = printer()._print(Cast('b', dtype=dtype))
+    ext = 'F' if dtype is np.float32 else ''
+
     return (('SAFEINV(a, b)',
-             f'(((a) < {eps}F || ({b}) < {eps}F) ? (0.0F) : ((1.0F) / (a)))'),), {}
+             f'(((a) < {eps}{ext} || ({b}) < {eps}{ext}) ? '
+             f'(0.0{ext}) : ((1.0{ext}) / (a)))'),), {}
 
 
 @iet_pass

devito/passes/iet/parpragma.py

@@ -317,6 +317,9 @@ def _make_parregion(self, partree, parrays):
             i = n.write
             if not (i.is_Array or i.is_TempFunction):
                 continue
+            elif partree.dim in i.dimensions:
+                # Non-local Array (full iteration space): no need to vector-expand
+                continue
             elif i in parrays:
                 pi = parrays[i]
             else:

devito/symbolics/inspection.py

@@ -4,6 +4,7 @@
 from sympy import (Function, Indexed, Integer, Mul, Number,
                    Pow, S, Symbol, Tuple)
 from sympy.core.numbers import ImaginaryUnit
+from sympy.core.function import Application
 
 from devito.finite_differences import Derivative
 from devito.finite_differences.differentiable import IndexDerivative
@@ -12,7 +13,7 @@
 from devito.symbolics.extended_sympy import (CallFromPointer, Cast,
                                              DefFunction, ReservedWord)
 from devito.symbolics.queries import q_routine
-from devito.tools import as_tuple, prod
+from devito.tools import as_tuple, prod, is_integer
 from devito.tools.dtypes_lowering import infer_dtype
 
 __all__ = ['compare_ops', 'estimate_cost', 'has_integer_args', 'sympy_dtype']
@@ -116,7 +117,7 @@ def estimate_cost(exprs, estimate=False):
 estimate_values = {
     'elementary': 100,
     'pow': 50,
-    'SafeInv': 10,
+    'SafeInv': 50,
     'div': 5,
     'Abs': 5,
     'floor': 1,
@@ -211,6 +212,7 @@ def _(expr, estimate, seen):
 
 
 @_estimate_cost.register(Function)
+@_estimate_cost.register(Application)
 def _(expr, estimate, seen):
     if q_routine(expr):
         flops, _ = zip(*[_estimate_cost(a, estimate, seen) for a in expr.args])
@@ -227,6 +229,7 @@ def _(expr, estimate, seen):
             flops += 1
     else:
         flops = 0
+
     return flops, False
 
 
@@ -284,12 +287,14 @@ def has_integer_args(*args):
         try:
             return np.issubdtype(args[0].dtype, np.integer)
         except AttributeError:
-            return args[0].is_integer
+            return is_integer(args[0])
 
     res = True
     for a in args:
         try:
-            if isinstance(a, INT):
+            if isinstance(a, INT) or \
+               is_integer(a) or \
+               has_integer_args(a):
                 res = res and True
             elif len(a.args) > 0:
                 res = res and has_integer_args(*a.args)