compiler: Patch IterationSpace.intersection

Author: FabioLuporini

devito/ir/support/space.py

@@ -378,10 +378,16 @@ def generate(cls, op, *interval_groups, relations=None):
 
         >>> ig = IntervalGroup.generate('intersection', ig0, ig1, ig2)
         """
+        if op == 'intersection':
+            dims = set.intersection(*[set(ig.dimensions) for ig in interval_groups])
+        else:
+            dims = set().union(*[ig.dimensions for ig in interval_groups])
+
         mapper = {}
         for ig in interval_groups:
             for i in ig:
-                mapper.setdefault(i.dim, []).append(i)
+                if i.dim in dims:
+                    mapper.setdefault(i.dim, []).append(i)
 
         intervals = []
         for v in mapper.values():

tests/test_ir.py

@@ -13,7 +13,7 @@
                                      Vector, AFFINE, REGULAR, IRREGULAR, mocksym0,
                                      mocksym1)
 from devito.ir.support.space import (NullInterval, Interval, Forward, Backward,
-                                     IterationSpace)
+                                     IntervalGroup, IterationSpace)
 from devito.ir.support.guards import GuardOverflow
 from devito.symbolics import DefFunction, FieldFromPointer
 from devito.tools import prod
@@ -508,6 +508,24 @@ def test_intervals_switch(self, x, y):
         assert iy.switch(x) == ix
         assert ix.switch(y).switch(x) == ix
 
+    def test_space_intersection(self, x, y):
+        ig0 = IntervalGroup([Interval(x, 1, -1)])
+        ig1 = IntervalGroup([Interval(x, 2, -2), Interval(y, 3, -3)])
+
+        ig = IntervalGroup.generate('intersection', ig0, ig1)
+
+        assert len(ig) == 1
+        assert ig[0] == Interval(x, 2, -2)
+
+        # Now the same but with IterationSpaces
+        ispace0 = IterationSpace(ig0)
+        ispace1 = IterationSpace(ig1)
+
+        ispace = IterationSpace.intersection(ispace0, ispace1)
+
+        assert len(ispace) == 1
+        assert ispace.intervals == ig
+
 
 class TestDependenceAnalysis: