Merge branch 'devitocodes:main' into master

Author: ZoeLeibowitz

devito/operator/profiling.py

@@ -453,7 +453,7 @@ def add_glb_vanilla(self, key, time):
         if not self.input:
             return
 
-        ops = sum(v.ops for v in self.input.values())
+        ops = sum(v.ops for v in self.input.values() if not np.isnan(v.ops))
         traffic = sum(v.traffic for v in self.input.values())
 
         gflops = float(ops)/10**9

devito/passes/clusters/unevaluate.py

@@ -1,8 +1,5 @@
-import sympy
-
 from devito.ir import cluster_pass
-from devito.symbolics import reuse_if_untouched, q_leaf
-from devito.symbolics.unevaluation import Add, Mul, Pow
+from devito.symbolics import unevaluate as _unevaluate
 
 __all__ = ['unevaluate']
 
@@ -12,22 +9,3 @@ def unevaluate(cluster):
     exprs = [_unevaluate(e) for e in cluster.exprs]
 
     return cluster.rebuild(exprs=exprs)
-
-
-mapper = {
-    sympy.Add: Add,
-    sympy.Mul: Mul,
-    sympy.Pow: Pow
-}
-
-
-def _unevaluate(expr):
-    if q_leaf(expr):
-        return expr
-
-    args = [_unevaluate(a) for a in expr.args]
-
-    try:
-        return mapper[expr.func](*args)
-    except KeyError:
-        return reuse_if_untouched(expr, args)

devito/passes/iet/definitions.py

@@ -16,7 +16,7 @@
 from devito.passes.iet.engine import iet_pass
 from devito.passes.iet.langbase import LangBB
 from devito.symbolics import (Byref, DefFunction, FieldFromPointer, IndexedPointer,
-                              SizeOf, VOID, pow_to_mul)
+                              SizeOf, VOID, pow_to_mul, unevaluate)
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap,
                           DeviceRM, Eq, Symbol)
@@ -119,7 +119,7 @@ def _alloc_array_on_global_mem(self, site, obj, storage):
 
         # Create input array
         name = '%s_init' % obj.name
-        initvalue = np.array([pow_to_mul(i) for i in obj.initvalue])
+        initvalue = np.array([unevaluate(pow_to_mul(i)) for i in obj.initvalue])
         src = Array(name=name, dtype=obj.dtype, dimensions=obj.dimensions,
                     space='host', scope='stack', initvalue=initvalue)
 

devito/symbolics/extended_sympy.py

@@ -1,6 +1,7 @@
 """
 Extended SymPy hierarchy.
 """
+import re
 
 import numpy as np
 import sympy
@@ -426,13 +427,29 @@ def __new__(cls, base, dtype=None, stars=None, reinterpret=False, **kwargs):
             # E.g. void
             pass
 
+        dtype, stars = cls._process_dtype(dtype, stars)
+
         obj = super().__new__(cls, base)
         obj._stars = stars or ''
         obj._dtype = dtype
         obj._reinterpret = reinterpret
 
         return obj
 
+    @classmethod
+    def _process_dtype(cls, dtype, stars):
+        if not isinstance(dtype, str) or stars is not None:
+            return dtype, stars
+
+        # String dtype, e.g. "float", "int*", "foo**"
+        match = re.fullmatch(r'(\w+)\s*(\*+)?', dtype)
+        if match:
+            dtype = match.group(1)
+            stars = match.group(2) or ''
+            return dtype, stars
+        else:
+            return dtype, stars
+
     def _hashable_content(self):
         return super()._hashable_content() + (self._stars,)
 
@@ -461,7 +478,10 @@ def _C_ctype(self):
 
     @property
     def _op(self):
-        return f'({ctypes_to_cstr(self._C_ctype)})'
+        cstr = ctypes_to_cstr(self._C_ctype)
+        if self.stars:
+            cstr = f"{cstr}{self.stars}"
+        return f'({cstr})'
 
     def __str__(self):
         return f"{self._op}{self.base}"

devito/symbolics/manipulation.py

@@ -13,7 +13,9 @@
 from devito.symbolics.extended_sympy import DefFunction, rfunc
 from devito.symbolics.queries import q_leaf
 from devito.symbolics.search import retrieve_indexed, retrieve_functions
-from devito.symbolics.unevaluation import Mul as UMul
+from devito.symbolics.unevaluation import (
+    Add as UnevalAdd, Mul as UnevalMul, Pow as UnevalPow
+)
 from devito.tools import as_list, as_tuple, flatten, split, transitive_closure
 from devito.types.basic import Basic, Indexed
 from devito.types.array import ComponentAccess
@@ -22,7 +24,7 @@
 
 __all__ = ['xreplace_indices', 'pow_to_mul', 'indexify', 'subs_op_args',
            'normalize_args', 'uxreplace', 'Uxmapper', 'subs_if_composite',
-           'reuse_if_untouched', 'evalrel', 'flatten_args']
+           'reuse_if_untouched', 'evalrel', 'flatten_args', 'unevaluate']
 
 
 def uxreplace(expr, rule):
@@ -338,7 +340,7 @@ def pow_to_mul(expr):
             # but at least we traverse the base looking for other Pows
             return expr.func(pow_to_mul(base), exp, evaluate=False)
         elif exp > 0:
-            return UMul(*[pow_to_mul(base)]*int(exp), evaluate=False)
+            return UnevalMul(*[pow_to_mul(base)]*int(exp), evaluate=False)
         elif exp < 0:
             # Reciprocal powers become inverse of the negative power
             # for example Pow(expr, -2) becomes Pow(expr * expr, -1)
@@ -502,3 +504,18 @@ def evalrel(func=min, input=None, assumptions=None):
     except TypeError:
         pass
     return rfunc(func, *input)
+
+
+uneval_mapper = {Add: UnevalAdd, Mul: UnevalMul, Pow: UnevalPow}
+
+
+def unevaluate(expr):
+    if q_leaf(expr):
+        return expr
+
+    args = [unevaluate(a) for a in expr.args]
+
+    try:
+        return uneval_mapper[expr.func](*args)
+    except KeyError:
+        return reuse_if_untouched(expr, args)

devito/tools/dtypes_lowering.py

@@ -256,7 +256,10 @@ class c_restrict_void_p(ctypes.c_void_p):
 
 def ctypes_to_cstr(ctype, toarray=None):
     """Translate ctypes types into C strings."""
-    if ctype in ctypes_vector_mapper.values():
+    if isinstance(ctype, str):
+        # Already a C string
+        return ctype
+    elif ctype in ctypes_vector_mapper.values():
         retval = ctype.__name__
     elif isinstance(ctype, CustomDtype):
         retval = str(ctype)

devito/types/array.py

@@ -4,7 +4,7 @@
 import numpy as np
 from sympy import Expr, cacheit
 
-from devito.tools import (Reconstructable, as_tuple, c_restrict_void_p,
+from devito.tools import (Pickable, as_tuple, c_restrict_void_p,
                           dtype_to_ctype, dtypes_vector_mapper, is_integer)
 from devito.types.basic import AbstractFunction, LocalType
 from devito.types.utils import CtypesFactory, DimensionTuple
@@ -533,10 +533,11 @@ def handles(self):
         return self.components
 
 
-class ComponentAccess(Expr, Reconstructable):
+class ComponentAccess(Expr, Pickable):
 
     _component_names = ('x', 'y', 'z', 'w')
 
+    __rargs__ = ('arg',)
     __rkwargs__ = ('index',)
 
     def __new__(cls, arg, index=0, **kwargs):
@@ -558,7 +559,7 @@ def __str__(self):
 
     __repr__ = __str__
 
-    func = Reconstructable._rebuild
+    func = Pickable._rebuild
 
     def _sympystr(self, printer):
         return str(self)
@@ -567,6 +568,10 @@ def _sympystr(self, printer):
     def base(self):
         return self.args[0]
 
+    @property
+    def arg(self):
+        return self.base
+
     @property
     def index(self):
         return self._index

requirements.txt

@@ -6,7 +6,7 @@ py-cpuinfo<10
 cgen>=2020.1
 codepy>=2019.1
 click<9.0
-multidict
+multidict<6.3
 anytree>=2.4.3,<=2.12.1
 cloudpickle
 packaging

tests/test_pickle.py

@@ -1,7 +1,7 @@
 import ctypes
 import pickle as pickle0
-import cloudpickle as pickle1
 
+import cloudpickle as pickle1
 import pytest
 import numpy as np
 from sympy import Symbol
@@ -19,7 +19,7 @@
 from devito.types import (Array, CustomDimension, Symbol as dSymbol, Scalar,
                           PointerArray, Lock, PThreadArray, SharedData, Timer,
                           DeviceID, NPThreads, ThreadID, TempFunction, Indirection,
-                          FIndexed)
+                          FIndexed, ComponentAccess)
 from devito.types.basic import BoundSymbol, AbstractSymbol
 from devito.tools import EnrichedTuple
 from devito.symbolics import (IntDiv, ListInitializer, FieldFromPointer,
@@ -416,6 +416,20 @@ def test_findexed(self, pickle):
         assert new_fi.indices == (x+1, y, z-2)
         assert new_fi.strides_map == fi.strides_map
 
+    def test_component_access(self, pickle):
+        grid = Grid(shape=(3, 3, 3))
+        x, y, z = grid.dimensions
+
+        f = Function(name='f', grid=grid)
+
+        ca = ComponentAccess(f.indexify(), 1)
+
+        pkl_ca = pickle.dumps(ca)
+        new_ca = pickle.loads(pkl_ca)
+
+        assert new_ca.index == 1
+        assert new_ca.function.name == f.name
+
     def test_symbolics(self, pickle):
         a = Symbol('a')
 

tests/test_symbolics.py

@@ -432,7 +432,7 @@ def test_rvalue():
     assert str(Rvalue(ctype, ns, init)) == 'my::namespace::dummytype{}'
 
 
-def test_cast():
+def test_basecast():
     s = Symbol(name='s', dtype=np.float32)
 
     class BarCast(BaseCast):
@@ -448,6 +448,28 @@ class BarCast(BaseCast):
     assert v != v1
 
 
+def test_str_cast():
+    s = Symbol(name='s', dtype=np.float32)
+
+    v = Cast(s, 'foo')
+    assert not v.stars
+    assert v.dtype == 'foo'
+    assert v._op == '(foo)'
+    assert ccode(v) == '(foo)s'
+
+    v = Cast(s, 'foo*')
+    assert v.stars == '*'
+    assert v.dtype == 'foo'
+    assert v._op == '(foo*)'
+    assert ccode(v) == '(foo*)s'
+
+    v = Cast(s, 'foo **')
+    assert v.stars == '**'
+    assert v.dtype == 'foo'
+    assert v._op == '(foo**)'
+    assert ccode(v) == '(foo**)s'
+
+
 def test_findexed():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions
@@ -474,6 +496,24 @@ def test_findexed():
     assert new_fi.strides_map == strides_map
 
 
+def test_component_access():
+    grid = Grid(shape=(3, 3, 3))
+    x, y, z = grid.dimensions
+
+    f = Function(name='f', grid=grid)
+
+    cf0 = ComponentAccess(f.indexify(), 0)
+    cf1 = ComponentAccess(f.indexify(), 1)
+
+    assert ccode(cf0) == 'f[x][y][z].x'
+    assert ccode(cf1) == 'f[x][y][z].y'
+
+    # Reconstruction
+    cf2 = cf1.func(*cf1.args)
+    assert cf2.index == cf1.index
+    assert cf2 == cf1
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions