Merge pull request #2728 from devitocodes/simplify-queue-dim-reuse

compiler: Enhance Array codegen

Author: mloubout

devito/arch/compiler.py

@@ -291,6 +291,22 @@ def load(self, soname):
         """
         return npct.load_library(str(self.get_jit_dir().joinpath(soname)), '.')
 
+    def save_header(self, filename, code):
+        """
+        Store some source code into a header file within the same temporary directory
+        used for JIT compilation.
+
+        Parameters
+        ----------
+        filename : str
+            The name of the header file (w/o the suffix).
+        code : str
+            The source code to be stored.
+        """
+        hfile = self.get_jit_dir().joinpath(filename).with_suffix('.h')
+        with open(str(hfile), 'w') as f:
+            f.write(code)
+
     def save(self, soname, binary):
         """
         Store a binary into a file within a temporary directory.

devito/ir/cgen/printer.py

@@ -41,6 +41,7 @@ class BasePrinter(CodePrinter):
     _prec_literals = {np.float32: 'F', np.complex64: 'F'}
 
     _qualifiers_mapper = {
+        'is_extern': 'extern',
         'is_const': 'const',
         'is_volatile': 'volatile',
         '_mem_constant': 'static',

devito/ir/iet/nodes.py

@@ -195,7 +195,7 @@ class List(Node):
 
     _traversable = ['body']
 
-    def __init__(self, header=None, body=None, footer=None):
+    def __init__(self, header=None, body=None, footer=None, inline=False):
         body = as_tuple(body)
         if len(body) == 1 and all(type(i) is List for i in [self, body[0]]):
             # De-nest Lists
@@ -213,6 +213,8 @@ def __init__(self, header=None, body=None, footer=None):
             self.body = as_tuple(body)
             self.footer = as_tuple(footer)
 
+        self.inline = inline
+
     def __repr__(self):
         return "<%s (%d, %d, %d)>" % (self.__class__.__name__, len(self.header),
                                       len(self.body), len(self.footer))
@@ -1047,9 +1049,8 @@ class Dereference(ExprStmt, Node):
     A node encapsulating a dereference from a `pointer` to a `pointee`.
     The following cases are supported:
 
-        * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array.
-        * `pointer` is an ArrayObject representing a pointer to a C struct, and
-          `pointee` is a field in `pointer`.
+        * `pointer` is an AbstractFunction, and `pointee` is an Array.
+        * `pointer` is an AbstractObject, and `pointee` is an Array.
         * `pointer` is a Symbol with its _C_ctype deriving from ct._Pointer, and
           `pointee` is a Symbol representing the dereferenced value.
     """
@@ -1075,20 +1076,23 @@ def expr_symbols(self):
             assert issubclass(self.pointer._C_ctype, ctypes._Pointer), \
                    "Scalar dereference must have a pointer ctype"
             ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
-        elif self.pointer.is_PointerArray or self.pointer.is_TempFunction:
+        elif self.pointer.is_AbstractFunction:
             ret.extend([self.pointer.indexed, self.pointee.indexed])
             ret.extend(flatten(i.free_symbols
                                for i in self.pointee.symbolic_shape[1:]))
             ret.extend(self.pointer.free_symbols)
+        elif self.pointer.is_AbstractObject:
+            ret.extend([self.pointer, self.pointee.indexed])
+            ret.extend(flatten(i.free_symbols
+                               for i in self.pointee.symbolic_shape[1:]))
         else:
-            ret.extend([self.pointer.indexed, self.pointee._C_symbol])
+            assert False, f"Unexpected pointer type {type(self.pointer)}"
+
         return tuple(filter_ordered(ret))
 
     @property
     def defines(self):
-        if self.pointer.is_PointerArray or \
-           self.pointer.is_TempFunction or \
-           self.pointee._mem_stack:
+        if self.pointer.is_AbstractFunction or self.pointee._mem_stack:
             return (self.pointee.indexed, self.pointee)
         else:
             return (self.pointee,)

devito/ir/iet/visitors.py

@@ -506,24 +506,30 @@ def visit_PointerCast(self, o):
 
     def visit_Dereference(self, o):
         a0, a1 = o.functions
-        if a1.is_PointerArray or a1.is_TempFunction:
-            i = a1.indexed
-            cstr = self.ccode(i._C_typedata)
+        if a0.is_AbstractFunction:
+            cstr = self.ccode(a0.indexed._C_typedata)
+
+            try:
+                # Special AbstractFunctions such as PointerArray or TempFunction
+                cdim = f'[{a1.dim.name}]'
+            except AttributeError:
+                cdim = ''
+
             if o.flat is None:
                 shape = ''.join(f"[{self.ccode(i)}]" for i in a0.symbolic_shape[1:])
-                rvalue = f'({cstr} (*){shape}) {a1.name}[{a1.dim.name}]'
+                rvalue = f'({cstr} (*){shape}) {a1.name}{cdim}'
                 lvalue = c.Value(cstr, f'(*{self._restrict_keyword} {a0.name}){shape}')
             else:
-                rvalue = f'({cstr} *) {a1.name}[{a1.dim.name}]'
+                rvalue = f'({cstr} *) {a1.name}{cdim}'
                 lvalue = c.Value(cstr, f'*{self._restrict_keyword} {a0.name}')
-            if a0._data_alignment:
-                lvalue = c.AlignedAttribute(a0._data_alignment, lvalue)
+
         else:
             if a1.is_Symbol:
                 rvalue = f'*{a1.name}'
             else:
                 rvalue = f'{a1.name}->{a0._C_name}'
             lvalue = self._gen_value(a0, 0)
+
         return c.Initializer(lvalue, rvalue)
 
     def visit_Block(self, o):
@@ -532,7 +538,11 @@ def visit_Block(self, o):
 
     def visit_List(self, o):
         body = flatten(self._visit(i) for i in self._blankline_logic(o.children))
-        return c.Module(o.header + (c.Collection(body),) + o.footer)
+        if o.inline:
+            body = c.Line(' '.join(str(i) for i in body))
+        else:
+            body = c.Collection(body)
+        return c.Module(o.header + (body,) + o.footer)
 
     def visit_Section(self, o):
         body = flatten(self._visit(i) for i in o.children)

devito/types/array.py

@@ -503,6 +503,10 @@ def symbolic_shape(self):
             # shape, i.e. something along the lines of  `(x_size, y_size, z_size)`
             return self.c0.symbolic_shape
 
+    @property
+    def nbytes(self):
+        return self.size*self.dtype.itemsize
+
     @property
     def _mem_heap(self):
         return not any([self._mem_stack, self._mem_shared, self._mem_shared_remote])

devito/types/misc.py

@@ -176,7 +176,7 @@ def _defines(self):
         return frozenset().union(*[i._defines for i in self])
 
 
-class Pointer(LocalObject):
+class Pointer(LocalObject, sympy.Expr):
 
     __rkwargs__ = LocalObject.__rkwargs__ + ('dtype',)
 

examples/performance/00_overview.ipynb

@@ -701,7 +701,7 @@
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    float (*restrict r0)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr0[tid];\n",
+      "    float (*restrict r0)[z_size] = (float (*)[z_size]) pr0[tid];\n",
       "\n",
       "    #pragma omp for schedule(dynamic,1)\n",
       "    for (int x = x_m; x <= x_M; x += 1)\n",
@@ -836,7 +836,7 @@
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    float (*restrict r1)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr1[tid];\n",
+      "    float (*restrict r1)[z_size] = (float (*)[z_size]) pr1[tid];\n",
       "\n",
       "    #pragma omp for schedule(dynamic,1)\n",
       "    for (int x = x_m; x <= x_M; x += 1)\n",
@@ -968,7 +968,7 @@
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    float (*restrict r0)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr0[tid];\n",
+      "    float (*restrict r0)[z_size] = (float (*)[z_size]) pr0[tid];\n",
       "\n",
       "    #pragma omp for schedule(dynamic,1)\n",
       "    for (int x = x_m; x <= x_M; x += 1)\n",
@@ -1236,7 +1236,7 @@
       "    #pragma omp parallel num_threads(nthreads)\n",
       "    {\n",
       "      const int tid = omp_get_thread_num();\n",
-      "      float (*restrict r2)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr2[tid];\n",
+      "      float (*restrict r2)[z_size] = (float (*)[z_size]) pr2[tid];\n",
       "\n",
       "      #pragma omp for collapse(2) schedule(dynamic,1)\n",
       "      for (int x0_blk0 = x_m; x0_blk0 <= x_M; x0_blk0 += x0_blk0_size)\n",
@@ -1348,7 +1348,7 @@
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    float (*restrict r2)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr2[tid];\n",
+      "    float (*restrict r2)[z_size] = (float (*)[z_size]) pr2[tid];\n",
       "\n",
       "    #pragma omp for collapse(2) schedule(dynamic,1)\n",
       "    for (int x0_blk0 = x_m; x0_blk0 <= x_M; x0_blk0 += x0_blk0_size)\n",
@@ -1527,7 +1527,7 @@
       "    #pragma omp parallel num_threads(nthreads)\n",
       "    {\n",
       "      const int tid = omp_get_thread_num();\n",
-      "      float (*restrict r2)[z_size] __attribute__ ((aligned (64))) = (float (*)[z_size]) pr2[tid];\n",
+      "      float (*restrict r2)[z_size] = (float (*)[z_size]) pr2[tid];\n",
       "\n",
       "      #pragma omp for schedule(dynamic,1)\n",
       "      for (int x = x_m; x <= x_M; x += 1)\n",

tests/test_iet.py

@@ -7,17 +7,18 @@
 
 from devito import (Eq, Grid, Function, TimeFunction, Operator, Dimension,  # noqa
                     switchconfig)
-from devito.ir.iet import (Call, Callable, Conditional, DeviceCall, DummyExpr,
-                           Iteration, List, KernelLaunch, Lambda, ElementalFunction,
-                           CGen, FindSymbols, filter_iterations, make_efunc,
-                           retrieve_iteration_tree, Transformer)
+from devito.ir.iet import (
+    Call, Callable, Conditional, Definition, DeviceCall, DummyExpr, Iteration, List,
+    KernelLaunch, Lambda, ElementalFunction, CGen, FindSymbols, filter_iterations,
+    make_efunc, retrieve_iteration_tree, Transformer
+)
 from devito.ir import SymbolRegistry
 from devito.passes.iet.engine import Graph
 from devito.passes.iet.languages.C import CDataManager
 from devito.symbolics import (Byref, FieldFromComposite, InlineIf, Macro, Class,
-                              FLOAT)
+                              String, FLOAT)
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
-from devito.types import Array, LocalObject, Symbol
+from devito.types import CustomDimension, Array, LocalObject, Symbol
 
 
 @pytest.fixture
@@ -475,3 +476,23 @@ def test_codegen_quality0():
     assert len(foo.parameters) == 3
     assert len(foo1.parameters) == 1
     assert foo1.parameters[0] is a
+
+
+def test_special_array_definition():
+
+    class MyArray(Array):
+        is_extern = True
+        _data_alignment = False
+
+    dim = CustomDimension(name='d', symbolic_size=String(''))
+    a = MyArray(name='a', dimensions=dim, scope='shared', dtype=np.uint8)
+
+    assert str(Definition(a)) == "extern  unsigned char a[];"
+
+
+def test_list_inline():
+    expr0 = DummyExpr(Symbol(name='a'), 1)
+    expr1 = DummyExpr(Symbol(name='b'), 2)
+
+    lst = List(body=[expr0, expr1], inline=True)
+    assert str(lst) == """a = 1; b = 2;"""