Merge pull request #2598 from devitocodes/streaming-loads-3

compiler: Add VectorAccess as an internal type

Author: FabioLuporini

devito/arch/compiler.py

@@ -705,10 +705,14 @@ def __init_finalize__(self, **kwargs):
         if not configuration['safe-math']:
             self.cflags.append('--use_fast_math')
 
-        # Optionally print out per-kernel shared memory and register usage
         if configuration['profiling'] == 'advanced2':
+            # Optionally print out per-kernel shared memory and register usage
             self.cflags.append('--ptxas-options=-v')
 
+            # Useful for Nsight Compute to associate the source code with
+            # profiling information
+            self.cflags.append('-lineinfo')
+
         self.src_ext = 'cu'
 
         # NOTE: not sure where we should place this. It definitely needs

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):
@@ -453,10 +454,7 @@ def _C_ctype(self):
 
     @property
     def _op(self):
-        cstr = ctypes_to_cstr(self._C_ctype)
-        if self.stars:
-            cstr = f"{cstr}{self.stars}"
-        return f'({cstr})'
+        return f'({ctypes_to_cstr(self._C_ctype)}{self.stars})'
 
     def __str__(self):
         return f"{self._op}{self.base}"
@@ -796,6 +794,36 @@ def __str__(self):
     __repr__ = __str__
 
 
+class VectorAccess(Expr, Pickable, BasicWrapperMixin):
+
+    """
+    Represent a vector access operation at high-level.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        return Expr.__new__(cls, *args)
+
+    def __str__(self):
+        return f"VL<{self.base}>"
+
+    __repr__ = __str__
+
+    func = Pickable._rebuild
+
+    @property
+    def base(self):
+        return self.args[0]
+
+    @property
+    def indices(self):
+        return self.base.indices
+
+    @cacheit
+    def sort_key(self, order=None):
+        # Ensure that the VectorAccess is sorted as the base
+        return self.base.sort_key(order=order)
+
+
 # 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) == 'VL<f[x, y, z]>'
+
+    # Reconstruction
+    v1 = v.func(g.indexify())
+    assert ccode(v1) == 'VL<g[x, y, z]>'
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions