Merge pull request #2128 from devitocodes/revamp-prec-sparse-func

api: Revamp interpolation/injection

.github/workflows/examples-mpi.yml

@@ -57,6 +57,7 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install -e .[extras,mpi,tests]
+        python3 scripts/clear_devito_cache.py
 
     - name: Test mpi notebooks
       continue-on-error: true

.github/workflows/examples.yml

@@ -49,7 +49,8 @@ jobs:
         pip install -e .[tests,extras]
 
     - name: Tests in examples
-      run: py.test --cov --cov-config=.coveragerc --cov-report=xml examples/
+      run: |
+        py.test --cov --cov-config=.coveragerc --cov-report=xml examples/
 
     - name: Seismic acoustic examples
       run: |

devito/ir/equations/algorithms.py

@@ -1,10 +1,10 @@
 from collections.abc import Iterable
-from operator import attrgetter
 
 from sympy import sympify
 
-from devito.symbolics import retrieve_indexed, uxreplace
-from devito.tools import PartialOrderTuple, as_tuple, filter_sorted, flatten
+from devito.symbolics import retrieve_indexed, uxreplace, retrieve_dimensions
+from devito.tools import (PartialOrderTuple, as_tuple, flatten,
+                          filter_sorted, filter_ordered)
 from devito.types import Dimension, IgnoreDimSort
 from devito.types.basic import AbstractFunction
 
@@ -33,8 +33,7 @@ def handle_indexed(indexed):
 
                 # Fallback: Just insert all the Dimensions we find, regardless of
                 # what the user is attempting to do
-                relation.extend([d for d in filter_sorted(i.free_symbols)
-                                 if isinstance(d, Dimension)])
+                relation.extend(filter_sorted(i.atoms(Dimension)))
 
         # StencilDimensions are lowered subsequently through special compiler
         # passes, so they can be ignored here
@@ -51,30 +50,42 @@ def handle_indexed(indexed):
     relations.add(expr.implicit_dims)
 
     # Add in leftover free dimensions (not an Indexed' index)
-    extra = set([i for i in expr.free_symbols if isinstance(i, Dimension)])
+    extra = set(retrieve_dimensions(expr, deep=True))
 
     # Add in pure data dimensions (e.g., those accessed only via explicit values,
     # such as A[3])
     indexeds = retrieve_indexed(expr, deep=True)
-    extra.update(set().union(*[set(i.function.dimensions) for i in indexeds]))
+    for i in indexeds:
+        extra.update({d for d in i.function.dimensions if i.indices[d].is_integer})
 
     # Enforce determinism
-    extra = filter_sorted(extra, key=attrgetter('name'))
+    extra = filter_sorted(extra)
 
     # Add in implicit relations for parent dimensions
     # -----------------------------------------------
     # 1) Note that (d.parent, d) is what we want, while (d, d.parent) would be
     # wrong; for example, in `((t, time), (t, x, y), (x, y))`, `x` could now
     # preceed `time`, while `t`, and therefore `time`, *must* appear before `x`,
     # as indicated by the second relation
-    implicit_relations = {(d.parent, d) for d in extra if d.is_Derived}
+    implicit_relations = {(d.parent, d) for d in extra if d.is_Derived and not d.indirect}
+
     # 2) To handle cases such as `((time, xi), (x,))`, where `xi` a SubDimension
     # of `x`, besides `(x, xi)`, we also have to add `(time, x)` so that we
     # obtain the desired ordering `(time, x, xi)`. W/o `(time, x)`, the ordering
     # `(x, time, xi)` might be returned instead, which would be non-sense
-    implicit_relations.update({tuple(d.root for d in i) for i in relations})
-
-    ordering = PartialOrderTuple(extra, relations=(relations | implicit_relations))
+    for i in relations:
+        dims = []
+        for d in i:
+            # Only add index if a different Dimension name to avoid dropping conditionals
+            # with the same name as the parent
+            if d.index.name == d.name:
+                dims.append(d)
+            else:
+                dims.extend([d.index, d])
+
+        implicit_relations.update({tuple(filter_ordered(dims))})
+
+    ordering = PartialOrderTuple(extra, relations=implicit_relations)
 
     return ordering
 

devito/ir/support/basic.py

@@ -123,7 +123,7 @@ def index_mode(self):
     def aindices(self):
         retval = []
         for i, fi in zip(self, self.findices):
-            dims = {j for j in i.free_symbols if isinstance(j, Dimension)}
+            dims = set(d.root if d.indirect else d for d in i.atoms(Dimension))
             sdims = {d for d in dims if d.is_Stencil}
             candidates = dims - sdims
 
@@ -660,9 +660,9 @@ def is_const(self, dim):
         """
         True if a constant dependence, that is no Dimensions involved, False otherwise.
         """
-        return (self.source.aindices[dim] is None and
-                self.sink.aindices[dim] is None and
-                self.distance_mapper[dim] == 0)
+        return (self.source.aindices.get(dim) is None and
+                self.sink.aindices.get(dim) is None and
+                self.distance_mapper.get(dim, 0) == 0)
 
     @memoized_meth
     def is_carried(self, dim=None):

devito/ir/support/utils.py

@@ -183,6 +183,7 @@ def detect_accesses(exprs):
     for e in as_tuple(exprs):
         other_dims.update(i for i in e.free_symbols if isinstance(i, Dimension))
         other_dims.update(e.implicit_dims)
+    other_dims = filter_sorted(other_dims)
     mapper[None] = Stencil([(i, 0) for i in other_dims])
 
     return mapper