Merge pull request #2566 from devitocodes/cast-str

compiler: Allow str cast type

Author: mloubout

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/symbolics/extended_sympy.py

@@ -1,6 +1,7 @@
 """
 Extended SymPy hierarchy.
 """
+import re
 
 import numpy as np
 import sympy
@@ -394,13 +395,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,)
 
@@ -429,7 +446,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/tools/dtypes_lowering.py

@@ -250,7 +250,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)

examples/seismic/tutorials/13_LSRTM_acoustic.ipynb

@@ -304,7 +304,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -329,12 +329,13 @@
     "    dm_true =  (solver.model.vp.data**(-2) - model0.vp.data**(-2))\n",
     "    \n",
     "    objective = 0.\n",
+    "    u0 = None\n",
     "    for i in range(nshots):\n",
     "        \n",
     "        #Observed Data using Born's operator\n",
     "        geometry.src_positions[0, :] = source_locations[i, :]\n",
     "\n",
-    "        _, u0, _ = solver.forward(vp=model0.vp, save=True)\n",
+    "        _, u0, _ = solver.forward(vp=model0.vp, save=True, u=u0)\n",
     "        \n",
     "        _, _, _,_ = solver.jacobian(dm_true, vp=model0.vp, rec = d_obs)\n",
     "        \n",

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_symbolics.py

@@ -419,7 +419,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):
@@ -435,6 +435,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