compiler: Support StencilDimension in estimate_cost

devito/operator/operator.py

@@ -355,7 +355,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
  clusters = cls._specialize_clusters(clusters, **kwargs)
 
  # Operation count after specialization
- #TODO: extend estimate_cost to parse StencilDimension correctly
  final_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense)
  try:
  profiler.record_ops_variation(init_ops, final_ops)

devito/symbolics/inspection.py

@@ -4,12 +4,14 @@
 import numpy as np
 from sympy import Function, Indexed, Integer, Mul, Number, Pow, S, Symbol, Tuple
 
+from devito.finite_differences import Derivative
+from devito.finite_differences.differentiable import IndexDerivative
 from devito.logger import warning
 from devito.symbolics.extended_sympy import (INT, CallFromPointer, Cast,
  DefFunction, ReservedWord)
 from devito.symbolics.queries import q_routine
 from devito.symbolics.search import search
-from devito.tools import as_tuple
+from devito.tools import as_tuple, prod
 
 __all__ = ['compare_ops', 'count', 'estimate_cost']
 
@@ -220,3 +222,22 @@ def _(expr, estimate):
  else:
  flops += 1
  return flops, False
+
+
+@_estimate_cost.register(Derivative)
+def _(expr, estimate):
+ return _estimate_cost(expr._evaluate(expand=False), estimate)
+
+
+@_estimate_cost.register(IndexDerivative)
+def _(expr, estimate):
+ flops, _ = _estimate_cost(expr.expr, estimate)
+
+ # It's an increment
+ flops += 1
+
+ # To be multiplied by the number of points this index sum implicitly
+ # iterates over
+ flops *= prod(i._size for i in expr.dimensions)
+
+ return flops, False

tests/test_dse.py

@@ -18,7 +18,7 @@
 from devito.symbolics import (INT, FLOAT, DefFunction, FieldFromPointer, # noqa
  Keyword, SizeOf, estimate_cost, pow_to_mul, indexify)
 from devito.tools import as_tuple, generator
-from devito.types import Scalar, Array, Symbol
+from devito.types import Array, Scalar, Symbol
 
 from examples.seismic.acoustic import AcousticWaveSolver
 from examples.seismic import demo_model, AcquisitionGeometry
@@ -225,6 +225,11 @@ def test_pow_to_mul(expr, expected):
  ('Eq(t0, ffp)', 0, True),
  ('Eq(t0, ffp + 1.)', 1, True),
  ('Eq(t0, ffp + ffp)', 1, True),
+ # W/ StencilDimensions
+ ('Eq(fb, fd.dx)', 10, False),
+ ('Eq(fb, fd.dx)', 10, True),
+ ('Eq(fb, fd.dx._evaluate(expand=False))', 10, False),
+ ('Eq(fb, fd.dx.dy + fa.dx)', 66, False),
 ])
 def test_estimate_cost(expr, expected, estimate):
  # Note: integer arithmetic isn't counted
@@ -241,6 +246,7 @@ def test_estimate_cost(expr, expected, estimate):
  fa = Function(name='fa', grid=grid, shape=(4,), dimensions=(x,)) # noqa
  fb = Function(name='fb', grid=grid, shape=(4,), dimensions=(x,)) # noqa
  fc = Function(name='fc', grid=grid) # noqa
+ fd = Function(name='fd', grid=grid, space_order=4) # noqa
  foo = lambda *args: DefFunction('foo', tuple(args)) # noqa
  k = Keyword('k') # noqa
  ffp = FieldFromPointer('size', fa._C_symbol) # noqa
@@ -2730,7 +2736,7 @@ def test_fullopt(self):
  bns, _ = assert_blocking(op1, {'x0_blk0'}) # due to loop blocking
 
  assert summary0[('section0', None)].ops == 50
- assert summary0[('section1', None)].ops == 140
+ assert summary0[('section1', None)].ops == 148
  assert np.isclose(summary0[('section0', None)].oi, 2.851, atol=0.001)
 
  assert summary1[('section0', None)].ops == 31