compiler: Add VectorAccess

Author: FabioLuporini

devito/ir/cgen/printer.py

@@ -401,6 +401,7 @@ def _print_Fallback(self, expr):
     _print_IndexSum = _print_Fallback
     _print_ReservedWord = _print_Fallback
     _print_Basic = _print_Fallback
+    _print_VectorAccess = _print_Fallback
 
 
 # Lifted from SymPy so that we go through our own `_print_math_func`

devito/symbolics/extended_sympy.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import sympy
-from sympy import Expr, Function, Number, Tuple, sympify
+from sympy import Expr, Function, Number, Tuple, cacheit, sympify
 from sympy.core.decorators import call_highest_priority
 
 from devito.finite_differences.elementary import Min, Max
@@ -21,7 +21,8 @@
            'ListInitializer', 'Byref', 'IndexedPointer', 'Cast', 'DefFunction',
            'MathFunction', 'InlineIf', 'ReservedWord', 'Keyword', 'String',
            'Macro', 'Class', 'MacroArgument', 'Deref', 'Namespace',
-           'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin', 'ValueLimit']
+           'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin', 'ValueLimit',
+           'VectorAccess']
 
 
 class CondEq(sympy.Eq):
@@ -793,6 +794,51 @@ def __str__(self):
     __repr__ = __str__
 
 
+class VectorAccess(Expr, Pickable):
+
+    """
+    Represent a vector access operation at high-level.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        return Expr.__new__(cls, *args)
+
+    def __str__(self):
+        return f"VL4<{self.base}>"
+
+    __repr__ = __str__
+
+    func = Pickable._rebuild
+
+    def _sympystr(self, printer):
+        return str(self)
+
+    @property
+    def base(self):
+        return self.args[0]
+
+    @property
+    def function(self):
+        return self.base.function
+
+    @property
+    def indices(self):
+        return self.base.indices
+
+    @property
+    def dtype(self):
+        return self.function.dtype
+
+    @cacheit
+    def sort_key(self, order=None):
+        # Ensure that the VectorAccess is sorted as the base
+        return self.base.sort_key(order=order)
+
+    # Default assumptions correspond to those of the `base`
+    for i in ('is_real', 'is_imaginary', 'is_commutative'):
+        locals()[i] = property(lambda self, v=i: getattr(self.base, v))
+
+
 # Some other utility objects
 Null = Macro('NULL')
 

tests/test_symbolics.py

@@ -14,7 +14,7 @@
                               CallFromPointer, Cast, DefFunction, FieldFromPointer,
                               INT, FieldFromComposite, IntDiv, Namespace, Rvalue,
                               ReservedWord, ListInitializer, uxreplace, pow_to_mul,
-                              retrieve_derivatives, BaseCast, SizeOf)
+                              retrieve_derivatives, BaseCast, SizeOf, VectorAccess)
 from devito.tools import as_tuple, CustomDtype
 from devito.types import (Array, Bundle, FIndexed, LocalObject, Object,
                           ComponentAccess, StencilDimension, Symbol as dSymbol)
@@ -501,6 +501,25 @@ def test_component_access():
     assert cf2 == cf1
 
 
+def test_vector_access():
+    grid = Grid(shape=(3, 3, 3))
+
+    f = Function(name='f', grid=grid)
+    g = Function(name='g', grid=grid)
+
+    v = VectorAccess(f.indexify())
+
+    assert v.base == f.indexify()
+    assert v.function is f
+
+    # Code generation
+    assert ccode(v) == 'VL4<f[x, y, z]>'
+
+    # Reconstruction
+    v1 = v.func(g.indexify())
+    assert ccode(v1) == 'VL4<g[x, y, z]>'
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions