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')

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 total_order(self):
+ iderivs = self.expr.find(IndexDerivative)
+ return 1 + max([i.total_order 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):