Merge pull request #2060 from devitocodes/unexpansion-final

compiler: Implement graceful lowering of derivatives (aka "unexpansion")

.github/workflows/pytest-core-nompi.yml

@@ -26,7 +26,7 @@ jobs:
 
       matrix:
         name: [
-           pytest-ubuntu-py37-gcc5-omp,
+           pytest-ubuntu-py39-gcc11-noomp,
            pytest-ubuntu-py38-gcc12-omp,
            pytest-ubuntu-py36-gcc7-omp,
            pytest-ubuntu-py310-gcc10-noomp,
@@ -38,19 +38,19 @@ jobs:
         ]
         set: [base, adjoint]
         include:
-        - name: pytest-ubuntu-py37-gcc5-omp
-          python-version: '3.7'
-          os: ubuntu-18.04
-          arch: "gcc-5"
-          language: "openmp"
-          sympy: "1.7"
+        - name: pytest-ubuntu-py39-gcc11-noomp
+          python-version: '3.9'
+          os: ubuntu-22.04
+          arch: "gcc-11"
+          language: "C"
+          sympy: "1.11"
 
         - name: pytest-ubuntu-py38-gcc12-omp
           python-version: '3.8'
           os: ubuntu-22.04
           arch: "gcc-12"
           language: "openmp"
-          sympy: "1.8"
+          sympy: "1.10"
 
         - name: pytest-ubuntu-py36-gcc7-omp
           python-version: '3.6'
@@ -78,7 +78,7 @@ jobs:
           os: ubuntu-20.04
           arch: "gcc-9"
           language: "openmp"
-          sympy: "1.8"
+          sympy: "1.9"
 
         - name: pytest-osx-py37-clang-omp
           python-version: '3.7'

devito/arch/compiler.py

@@ -12,8 +12,8 @@
 from codepy.jit import compile_from_string
 from codepy.toolchain import GCCToolchain
 
-from devito.arch import (AMDGPUX, NVIDIAX, M1, SKX, POWER8, POWER9, get_nvidia_cc,
-                         check_cuda_runtime, get_m1_llvm_path)
+from devito.arch import (AMDGPUX, Cpu64, M1, NVIDIAX, SKX, POWER8, POWER9,
+                         get_nvidia_cc, check_cuda_runtime, get_m1_llvm_path)
 from devito.exceptions import CompilationError
 from devito.logger import debug, warning, error
 from devito.parameters import configuration
@@ -535,6 +535,9 @@ def __init__(self, *args, **kwargs):
                 self.cflags.extend(['-mp', '-acc:gpu'])
             elif language == 'openmp':
                 self.cflags.extend(['-mp=gpu'])
+        elif isinstance(platform, Cpu64):
+            if language == 'openmp':
+                self.cflags.append('-mp')
 
         if not configuration['safe-math']:
             self.cflags.append('-fast')
@@ -764,6 +767,7 @@ def __lookup_cmds__(self):
     'clang': ClangCompiler,
     'cray': CrayCompiler,
     'aomp': AOMPCompiler,
+    'amdclang': AOMPCompiler,
     'hip': HipCompiler,
     'pgcc': PGICompiler,
     'pgi': PGICompiler,

devito/core/cpu.py

@@ -61,6 +61,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['par-nested'] = oo.pop('par-nested', cls.PAR_NESTED)
 
         # Misc
+        o['expand'] = oo.pop('expand', cls.EXPAND)
         o['optcomms'] = oo.pop('optcomms', True)
         o['linearize'] = oo.pop('linearize', False)
         o['mapify-reduce'] = oo.pop('mapify-reduce', cls.MAPIFY_REDUCE)

devito/core/gpu.py

@@ -75,6 +75,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['gpu-fit'] = as_tuple(oo.pop('gpu-fit', cls._normalize_gpu_fit(**kwargs)))
 
         # Misc
+        o['expand'] = oo.pop('expand', cls.EXPAND)
         o['optcomms'] = oo.pop('optcomms', True)
         o['linearize'] = oo.pop('linearize', False)
         o['mapify-reduce'] = oo.pop('mapify-reduce', cls.MAPIFY_REDUCE)

devito/core/operator.py

@@ -17,6 +17,7 @@
 class BasicOperator(Operator):
 
     # Default values for various optimization options
+
     CSE_MIN_COST = 1
     """
     Minimum computational cost of an operation to be eliminated as a
@@ -88,6 +89,12 @@ class BasicOperator(Operator):
     which may be easier to parallelize for certain backends.
     """
 
+    EXPAND = True
+    """
+    Unroll all loops with short, numeric trip count, such as loops created by
+    finite-difference derivatives.
+    """
+
     MPI_MODES = tuple(mpi_registry)
     """
     The supported MPI modes.

devito/finite_differences/derivative.py

@@ -83,9 +83,11 @@ class Derivative(sympy.Derivative, Differentiable):
     evaluation are `x0`, `fd_order` and `side`.
     """
 
-    _state = ('expr', 'dims', 'side', 'fd_order', 'transpose', '_ppsubs', 'x0')
     _fd_priority = 3
 
+    __rargs__ = ('expr', 'dims')
+    __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs', 'x0')
+
     def __new__(cls, expr, *dims, **kwargs):
         if type(expr) == sympy.Derivative:
             raise ValueError("Cannot nest sympy.Derivative with devito.Derivative")
@@ -104,7 +106,8 @@ def __new__(cls, expr, *dims, **kwargs):
         obj._deriv_order = orders if skip else DimensionTuple(*orders, getters=obj._dims)
         obj._side = kwargs.get("side")
         obj._transpose = kwargs.get("transpose", direct)
-        obj._ppsubs = as_tuple(frozendict(i) for i in kwargs.get("subs", []))
+        obj._ppsubs = as_tuple(frozendict(i) for i in
+                               kwargs.get("subs", kwargs.get("_ppsubs", [])))
         obj._x0 = frozendict(kwargs.get('x0', {}))
         return obj
 
@@ -237,7 +240,7 @@ def _xreplace(self, subs):
 
     @cached_property
     def _metadata(self):
-        state = list(self._state)
+        state = list(self.__rargs__ + self.__rkwargs__)
         state.remove('expr')
         ret = [getattr(self, i) for i in state]
         ret.append(self.expr.staggered or (None,))

devito/finite_differences/differentiable.py

@@ -2,20 +2,22 @@
 from itertools import product
 from functools import singledispatch
 
+from cached_property import cached_property
 import numpy as np
 import sympy
 from sympy.core.add import _addsort
 from sympy.core.mul import _keep_coeff, _mulsort
 from sympy.core.decorators import call_highest_priority
 from sympy.core.evalf import evalf_table
 
-from cached_property import cached_property
 from devito.finite_differences.tools import make_shift_x0
 from devito.logger import warning
-from devito.tools import as_tuple, filter_ordered, flatten, is_integer, split
-from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
+from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
+                          infer_dtype, is_integer, split)
+from devito.types import (Array, DimensionTuple, Evaluable, Indexed, Spacing,
+                          StencilDimension)
 
-__all__ = ['Differentiable', 'IndexDerivative', 'EvalDerivative']
+__all__ = ['Differentiable', 'IndexDerivative', 'EvalDerivative', 'Weights']
 
 
 class Differentiable(sympy.Expr, Evaluable):
@@ -29,7 +31,7 @@ class Differentiable(sympy.Expr, Evaluable):
     # operators to be used
     _op_priority = sympy.Expr._op_priority + 1.
 
-    _state = ('space_order', 'time_order', 'indices')
+    __rkwargs__ = ('space_order', 'time_order', 'indices')
 
     @cached_property
     def _functions(self):
@@ -63,6 +65,11 @@ def grid(self):
         except KeyError:
             return None
 
+    @cached_property
+    def dtype(self):
+        dtypes = {f.dtype for f in self.find(Indexed)} - {None}
+        return infer_dtype(dtypes)
+
     @cached_property
     def indices(self):
         return tuple(filter_ordered(flatten(getattr(i, 'indices', ())
@@ -222,7 +229,8 @@ def __neg__(self):
 
     def __eq__(self, other):
         return super(Differentiable, self).__eq__(other) and\
-            all(getattr(self, i, None) == getattr(other, i, None) for i in self._state)
+            all(getattr(self, i, None) == getattr(other, i, None)
+                for i in self.__rkwargs__)
 
     @property
     def name(self):
@@ -503,6 +511,8 @@ class IndexSum(DifferentiableOp):
     Dimensions, of an indexed expression.
     """
 
+    __rargs__ = ('expr', 'dimensions')
+
     is_commutative = True
 
     def __new__(cls, expr, dimensions, **kwargs):
@@ -517,16 +527,21 @@ def __new__(cls, expr, dimensions, **kwargs):
                 pass
             raise ValueError("Expected Dimension with numeric size, "
                              "got `%s` instead" % d)
-        if not expr.has_free(*dimensions):
+
+        # TODO: `has_free` only available with SymPy v>=1.10
+        # We should start using `not expr.has_free(*dimensions)` once we drop
+        # support for SymPy 1.8<=v<1.0
+        if not all(d in expr.free_symbols for d in dimensions):
             raise ValueError("All Dimensions `%s` must appear in `expr` "
                              "as free variables" % str(dimensions))
+
         for i in expr.find(IndexSum):
             for d in dimensions:
                 if d in i.dimensions:
                     raise ValueError("Dimension `%s` already appears in a "
                                      "nested tensor contraction" % d)
 
-        obj = sympy.Expr.__new__(cls, expr, *dimensions)
+        obj = sympy.Expr.__new__(cls, expr)
         obj._expr = expr
         obj._dimensions = dimensions
 
@@ -538,6 +553,9 @@ def __repr__(self):
 
     __str__ = __repr__
 
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.dimensions,)
+
     @property
     def expr(self):
         return self._expr
@@ -563,6 +581,8 @@ def _evaluate(self, **kwargs):
     def free_symbols(self):
         return super().free_symbols - set(self.dimensions)
 
+    func = DifferentiableOp._rebuild
+
 
 class Weights(Array):
 
@@ -579,47 +599,92 @@ def __init_finalize__(self, *args, **kwargs):
         assert isinstance(d, StencilDimension) and d.symbolic_size == len(weights)
         assert isinstance(weights, (list, tuple, np.ndarray))
 
-        kwargs['scope'] = 'static'
+        try:
+            self._spacings = set().union(*[i.find(Spacing) for i in weights])
+        except AttributeError:
+            self._spacing = set()
+
+        kwargs['scope'] = 'constant'
 
         super().__init_finalize__(*args, **kwargs)
 
+    def __eq__(self, other):
+        return (isinstance(other, Weights) and
+                self.dimension is other.dimension and
+                self.name == other.name and
+                self.indices == other.indices and
+                self.weights == other.weights)
+
+    __hash__ = sympy.Basic.__hash__
+
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.name,) + tuple(self.weights)
+
     @property
     def dimension(self):
         return self.dimensions[0]
 
+    @property
+    def spacings(self):
+        return self._spacings
+
     weights = Array.initvalue
 
 
 class IndexDerivative(IndexSum):
 
-    def __new__(cls, expr, dimensions, **kwargs):
-        if not (expr.is_Mul and len(expr.args) == 2):
-            raise ValueError("Expect expr*weights, got `%s` instead" % str(expr))
-        _, weights = expr.args
-        if not isinstance(weights, Weights):
-            # All of the SymPy versions we support end up placing the Weights
-            # array here, so if something changes we'll get an alarm through
-            # this exception
-            raise ValueError("Couldn't find weights array")
+    __rargs__ = ('expr', 'mapper')
+
+    def __new__(cls, expr, mapper, **kwargs):
+        dimensions = as_tuple(mapper.values())
+
+        # Detect the Weights among the arguments
+        weightss = []
+        for a in expr.args:
+            try:
+                f = a.function
+            except AttributeError:
+                continue
+            if isinstance(f, Weights):
+                weightss.append(a)
+
+        # Sanity check
+        if not (expr.is_Mul and len(weightss) == 1):
+            raise ValueError("Expect `expr*weights`, got `%s` instead" % str(expr))
+        weights = weightss.pop()
 
         obj = super().__new__(cls, expr, dimensions)
         obj._weights = weights
+        obj._mapper = frozendict(mapper)
 
         return obj
 
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.mapper,)
+
     @property
     def weights(self):
         return self._weights
 
+    @property
+    def mapper(self):
+        return self._mapper
+
+    @property
+    def depth(self):
+        iderivs = self.expr.find(IndexDerivative)
+        return 1 + max([i.depth for i in iderivs], default=0)
+
     def _evaluate(self, **kwargs):
         expr = super()._evaluate(**kwargs)
 
         if not kwargs.get('expand', True):
             return expr
 
         w = self.weights
+        f = w.function
         d = w.dimension
-        mapper = {w.subs(d, i): w.weights[n] for n, i in enumerate(d.range)}
+        mapper = {w.subs(d, i): f.weights[n] for n, i in enumerate(d.range)}
         expr = expr.xreplace(mapper)
 
         return expr

devito/finite_differences/finite_difference.py

@@ -94,7 +94,7 @@ def first_derivative(expr, dim, fd_order=None, side=centered, matvec=direct, x0=
 
 @check_input
 @check_symbolic
-def cross_derivative(expr, dims, fd_order, deriv_order, **kwargs):
+def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, **kwargs):
     """
     Arbitrary-order cross derivative of a given expression.
 
@@ -155,9 +155,10 @@ def cross_derivative(expr, dims, fd_order, deriv_order, **kwargs):
     (-1/h_y)*(-f(1, 2)*g(1, 2)/h_x + f(h_x + 1, 2)*g(h_x + 1, 2)/h_x) + (-f(1, h_y + 2)*\
 g(1, h_y + 2)/h_x + f(h_x + 1, h_y + 2)*g(h_x + 1, h_y + 2)/h_x)/h_y
     """
-    x0 = kwargs.get('x0', {})
+    x0 = x0 or {}
     for d, fd, dim in zip(deriv_order, fd_order, dims):
-        expr = generic_derivative(expr, dim=dim, fd_order=fd, deriv_order=d, x0=x0)
+        expr = generic_derivative(expr, dim=dim, fd_order=fd, deriv_order=d, x0=x0,
+                                  **kwargs)
 
     return expr
 
@@ -240,7 +241,7 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, symbolic
             # Pure number
             pass
 
-        deriv = IndexDerivative(expr*weights, weights.dimension)
+        deriv = IndexDerivative(expr*weights, {dim: indices.free_dim})
     else:
         terms = []
         for i, c in zip(indices, weights):