Merge pull request #2762 from devitocodes/dereference-expr

compiler: Enhance Dereference

Author: FabioLuporini

devito/ir/iet/nodes.py

@@ -1057,10 +1057,11 @@ class Dereference(ExprStmt, Node):
 
     is_Dereference = True
 
-    def __init__(self, pointee, pointer, flat=None):
+    def __init__(self, pointee, pointer, flat=None, offset=None):
         self.pointee = pointee
         self.pointer = pointer
         self.flat = flat
+        self.offset = offset
 
     def __repr__(self):
         return "<Dereference(%s,%s)>" % (self.pointee, self.pointer)
@@ -1088,6 +1089,9 @@ def expr_symbols(self):
         else:
             assert False, f"Unexpected pointer type {type(self.pointer)}"
 
+        if self.offset is not None:
+            ret.append(self.offset)
+
         return tuple(filter_ordered(ret))
 
     @property

devito/ir/iet/visitors.py

@@ -506,6 +506,12 @@ def visit_PointerCast(self, o):
 
     def visit_Dereference(self, o):
         a0, a1 = o.functions
+
+        if o.offset:
+            ptr = f'({a1.name} + {o.offset})'
+        else:
+            ptr = a1.name
+
         if a0.is_AbstractFunction:
             cstr = self.ccode(a0.indexed._C_typedata)
 
@@ -517,17 +523,17 @@ def visit_Dereference(self, o):
 
             if o.flat is None:
                 shape = ''.join(f"[{self.ccode(i)}]" for i in a0.symbolic_shape[1:])
-                rvalue = f'({cstr} (*){shape}) {a1.name}{cdim}'
+                rvalue = f'({cstr} (*){shape}) {ptr}{cdim}'
                 lvalue = c.Value(cstr, f'(*{self._restrict_keyword} {a0.name}){shape}')
             else:
-                rvalue = f'({cstr} *) {a1.name}{cdim}'
+                rvalue = f'({cstr} *) {ptr}{cdim}'
                 lvalue = c.Value(cstr, f'*{self._restrict_keyword} {a0.name}')
 
         else:
             if a1.is_Symbol:
-                rvalue = f'*{a1.name}'
+                rvalue = f'*{ptr}'
             else:
-                rvalue = f'{a1.name}->{a0._C_name}'
+                rvalue = f'{ptr}->{a0._C_name}'
             lvalue = self._gen_value(a0, 0)
 
         return c.Initializer(lvalue, rvalue)

tests/test_iet.py

@@ -9,16 +9,16 @@
                     switchconfig)
 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
+    KernelLaunch, Dereference, 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,
                               String, FLOAT)
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
-from devito.types import CustomDimension, Array, LocalObject, Symbol
+from devito.types import CustomDimension, Array, LocalObject, Symbol, Pointer
 
 
 @pytest.fixture
@@ -496,3 +496,16 @@ def test_list_inline():
 
     lst = List(body=[expr0, expr1], inline=True)
     assert str(lst) == """a = 1; b = 2;"""
+
+
+def test_dereference_base_plus_off():
+    ptr = Pointer(name='p', dtype=np.float32)
+    off = Symbol(name='offs', dtype=np.int32)
+
+    dim0 = CustomDimension(name='d0', symbolic_size=2)
+    dim1 = CustomDimension(name='d1', symbolic_size=3)
+    x = Array(name='x', dimensions=(dim0, dim1), dtype=np.float32)
+
+    deref = Dereference(x, ptr, offset=off)
+
+    assert str(deref) == "float (*restrict x)[3] = (float (*)[3]) (p + offs);"