symbolics: specialize sizeof

devito/__init__.py

@@ -65,7 +65,7 @@ def reinit_compiler(val):
 # Setup target platform and compiler
 configuration.add('platform', 'cpu64', list(platform_registry),
                   callback=lambda i: platform_registry[i]())
-configuration.add('compiler', 'custom', list(compiler_registry),
+configuration.add('compiler', 'custom', compiler_registry,
                   callback=lambda i: compiler_registry[i](name=i))
 
 # Setup language for shared-memory parallelism

devito/arch/compiler.py

@@ -988,12 +988,16 @@ class CompilerRegistry(dict):
     """
 
     def __getitem__(self, key):
+        if isinstance(key, Compiler):
+            key = key.name
+
         if key.startswith('gcc-'):
             i = key.split('-')[1]
             return partial(GNUCompiler, suffix=i)
+
         return super().__getitem__(key)
 
-    def __contains__(self, k):
+    def __contains__(self, key):
         if isinstance(k, Compiler):
             k = k.name
         return k in self.keys() or k.startswith('gcc-')

devito/passes/iet/definitions.py

@@ -17,7 +17,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, Keyword, pow_to_mul)
+                              SizeOf, VOID, pow_to_mul)
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap,
                           DeviceRM, Eq, Symbol)
@@ -279,7 +279,7 @@ def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage):
 
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
-        nbytes = SizeOf(Keyword('%s*' % obj._C_typedata))*obj.dim.symbolic_size
+        nbytes = SizeOf(obj._C_typedata, stars='*')*obj.dim.symbolic_size
         alloc0 = self.lang['host-alloc'](memptr, alignment, nbytes)
 
         free0 = self.lang['host-free'](obj._C_symbol)

devito/symbolics/extended_sympy.py

@@ -766,11 +766,12 @@ def __new__(cls, base=''):
 # DefFunction, unlike sympy.Function, generates e.g. `sizeof(float)`, not `sizeof(float_)`
 class SizeOf(DefFunction):
 
-    __rargs__ = ('intype',)
+    __rargs__ = ('intype', 'stars')
 
-    def __new__(cls, intype, **kwargs):
+    def __new__(cls, intype, stars=None, **kwargs):
         newobj = super().__new__(cls, 'sizeof', arguments=[str(intype)], **kwargs)
         newobj.intype = intype
+        newobj.stars = stars or ''
 
         return newobj
 

devito/symbolics/printer.py

@@ -347,6 +347,9 @@ def _print_DefFunction(self, expr):
             template = ''
         return "%s%s(%s)" % (expr.name, template, ','.join(arguments))
 
+    def _print_SizeOf(self, expr):
+        return f'sizeof({self._print(expr.intype)}{self._print(expr.stars)})'
+
     _print_MathFunction = _print_DefFunction
 
     def _print_Fallback(self, expr):
@@ -359,7 +362,6 @@ def _print_Fallback(self, expr):
     _print_IndexSum = _print_Fallback
     _print_ReservedWord = _print_Fallback
     _print_Basic = _print_Fallback
-    _print_SizeOf = _print_DefFunction
 
 
 # Lifted from SymPy so that we go through our own `_print_math_func`