compiler: Revamp FIndexed for correct reconstruction

devito/passes/iet/linearization.py

@@ -256,7 +256,7 @@ def _(f, indexeds, tracker, strides, sregistry):
 
         if len(i.indices) == i.function.ndim:
             v = tuple(strides.values())[-n:]
-            subs[i] = FIndexed(i, pname, strides=v)
+            subs[i] = FIndexed.from_indexed(i, pname, strides=v)
         else:
             # Honour custom indexing
             subs[i] = i.base[sum(i.indices)]

devito/types/misc.py

@@ -68,20 +68,21 @@ class FIndexed(Indexed, Pickable):
     `uX[x*ny + y]`, where `X` is a string provided by the caller.
     """
 
-    __rargs__ = ('indexed', 'pname')
+    __rargs__ = ('base', '*indices')
     __rkwargs__ = ('strides',)
 
-    def __new__(cls, indexed, pname, strides=None):
-        plabel = Symbol(name=pname, dtype=indexed.dtype)
-        base = IndexedData(plabel, None, function=indexed.function)
-        obj = super().__new__(cls, base, *indexed.indices)
-
-        obj.indexed = indexed
-        obj.pname = pname
+    def __new__(cls, base, *args, strides=None):
+        obj = super().__new__(cls, base, *args)
         obj.strides = as_tuple(strides)
 
         return obj
 
+    @classmethod
+    def from_indexed(cls, indexed, pname, strides=None):
+        label = Symbol(name=pname, dtype=indexed.dtype)
+        base = IndexedData(label, None, function=indexed.function)
+        return FIndexed(base, *indexed.indices, strides=strides)
+
     def __repr__(self):
         return "%s(%s)" % (self.name, ", ".join(str(i) for i in self.indices))
 
@@ -90,10 +91,16 @@ def __repr__(self):
     def _hashable_content(self):
         return super()._hashable_content() + (self.strides,)
 
+    func = Pickable._rebuild
+
     @property
     def name(self):
         return self.function.name
 
+    @property
+    def pname(self):
+        return self.base.name
+
     @property
     def free_symbols(self):
         # The functional representation of the FIndexed "hides" the strides, which

tests/test_symbolics.py

@@ -11,7 +11,7 @@
 from devito.symbolics import (retrieve_functions, retrieve_indexed, evalrel,  # noqa
                               CallFromPointer, Cast, FieldFromPointer,
                               FieldFromComposite, IntDiv, ccode, uxreplace)
-from devito.types import Array, LocalObject, Object, Symbol as dSymbol
+from devito.types import Array, FIndexed, LocalObject, Object, Symbol as dSymbol  # noqa
 
 
 def test_float_indices():
@@ -333,6 +333,7 @@ def test_solve_time():
     ('f[x, y+1]', '{f.indexed: g.indexed}', 'g[x, y+1]'),
     ('cos(f)', '{cos: sin}', 'sin(f)'),
     ('cos(f + sin(g))', '{cos: sin, sin: cos}', 'sin(f + cos(g))'),
+    ('FIndexed(f.indexed, x, y)', '{x: 0}', 'FIndexed(f.indexed, 0, y)'),
 ])
 def test_uxreplace(expr, subs, expected):
     grid = Grid(shape=(4, 4))