dsl: Look a DimensionTuple up by the Dimension asked for

`__getitem_hook__` matched on `_defines` overlap alone. A derived Dimension
carries its parent in `_defines`, so for a Bundle indexed by
`(p_rec, rp_recx)` -- `rp_recx` being a `CustomDimension` whose parent is
`p_rec` -- the lookup for `rp_recx` matched the `p_rec` entry first and
returned the number of sparse points where the number of interpolation
weights was meant.

That size becomes the innermost stride in `_generate_fsz`, so the receiver
kernels of a vectorized Operator read `w[p*npoint + rp]` instead of
`w[p*2 + rp]` and run off the end of the array. Observed as an out-of-bounds
`__global__` read under compute-sanitizer and a run-to-run varying, sometimes
NaN, elastic TTI gradient on CUDA.

Try an exact hit before falling back to the overlap, in both
`__getitem_hook__` and `dindex`.

Author: mloubout

devito/types/utils.py

@@ -45,12 +45,21 @@ class Stagger(Tag):
 class DimensionTuple(EnrichedTuple):
 
     def __getitem_hook__(self, dim):
+        # An exact hit wins over a `_defines` overlap.  A derived Dimension
+        # carries its parent in `_defines`, so a plain overlap test would
+        # match the parent's entry first and silently return the wrong
+        # component (e.g. a Bundle's `rp_rec` sub-dimension picking up the
+        # `p_rec` size).
+        if dim in self.getters:
+            return self.getters[dim]
         for d in self.getters:
             if d._defines & dim._defines:
                 return self.getters[d]
         raise KeyError
 
     def dindex(self, dim):
+        if dim in self.getters:
+            return list(self.getters).index(dim)
         for d in self.getters:
             if d._defines & dim._defines:
                 return list(self.getters).index(d)

tests/test_linearize.py

@@ -8,7 +8,7 @@
 )
 from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry
 from devito.passes import Graph, generate_macros, linearize
-from devito.types import Array, Bundle, DefaultDimension
+from devito.types import Array, Bundle, CustomDimension, DefaultDimension
 
 
 def test_basic():
@@ -716,3 +716,28 @@ def test_cire_n_strides():
     # NOTE: not exact equality because `op2` slightly changes the order of
     # arithmetic operations, which in turn causes some rounding differences
     assert np.allclose(u.data, u1.data, rtol=1e-4)
+
+
+def test_bundle_derived_dim_stride():
+    """
+    A Bundle's stride must come from the Dimension that is asked for, not from
+    a Dimension that merely appears in its `_defines`.
+
+    `MatrixSparseTimeFunction`-style interpolation introduces a `CustomDimension`
+    whose parent is the sparse Dimension, so `rp._defines` contains `p`.  Looking
+    the shape up by `_defines` overlap alone matched `p` first and gave the
+    Bundle the number of points as its innermost stride instead of the number of
+    interpolation weights, which reads out of bounds on the device.
+    """
+    grid = Grid(shape=(4, 4))
+    p = DefaultDimension(name='p', default_value=5)
+    rp = CustomDimension(name='rp', parent=p, symbolic_size=2)
+
+    w0 = Function(name='w0', dimensions=(p, rp), shape=(5, 2))
+    w1 = Function(name='w1', dimensions=(p, rp), shape=(5, 2))
+    bundle = Bundle(name='w0w1', components=(w0, w1), grid=grid)
+
+    assert rp._defines & p._defines  # the overlap that used to mislead
+    assert bundle.symbolic_shape[p] is not bundle.symbolic_shape[rp]
+    assert bundle.symbolic_shape[rp] == 2
+    assert bundle.symbolic_shape.dindex(rp) == 1