Merge pull request #2949 from devitocodes/faster-python-rebased-final

compiler: Faster compilation time

Author: mloubout

devito/arch/archinfo.py

@@ -497,6 +497,7 @@ def parse_product_arch():
 def get_visible_devices():
     device_vars = (
         'CUDA_VISIBLE_DEVICES',
+        'NVIDIA_VISIBLE_DEVICES',
         'ROCR_VISIBLE_DEVICES',
         'HIP_VISIBLE_DEVICES'
     )

devito/finite_differences/derivative.py

@@ -6,7 +6,9 @@
 
 import sympy
 
-from devito.tools import Pickable, as_mapper, as_tuple, frozendict, is_integer
+from devito.tools import (
+    Pickable, as_mapper, as_tuple, frozendict, is_integer, memoized_func
+)
 from devito.types.dimension import Dimension
 from devito.types.utils import DimensionTuple
 from devito.warnings import warn
@@ -557,6 +559,7 @@ def _evaluate(self, **kwargs):
     def _eval_deriv(self):
         return self._eval_fd(self.expr)
 
+    @memoized_func(scope='build')
     def _eval_fd(self, expr, **kwargs):
         """
         Evaluate the finite-difference approximation of the Derivative.

devito/finite_differences/differentiable.py

@@ -1028,6 +1028,10 @@ def compare(self, other):
     def base(self):
         return self.expr.func(*[a for a in self.expr.args if a is not self.weights])
 
+    @cached_property
+    def pivot(self):
+        return self.base.subs({d: 0 for d in self.dimensions})
+
     @property
     def weights(self):
         return self._weights

devito/finite_differences/finite_difference.py

@@ -170,14 +170,15 @@ def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coeffici
     # `coefficients` method (`taylor` or `symbolic`)
     if weights is None:
         weights = fd_weights_registry[coefficients](expr, deriv_order, indices, x0)
-    if isinstance(weights, Iterable) and len(weights) != len(indices):
+        _, wdim, _ = process_weights(weights, expr, dim)
+    elif isinstance(weights, Iterable) and len(weights) != len(indices):
         warning(f"Number of weights ({len(weights)}) does not match "
                 f"number of indices ({len(indices)}), reverting to Taylor")
         scale = False
+        wdim = None
         weights = fd_weights_registry['taylor'](expr, deriv_order, indices, x0)
 
     # Did fd_weights_registry return a new Function/Expression instead of a values?
-    _, wdim, _ = process_weights(weights, expr, dim)
     if wdim is not None:
         weights = [weights._subs(wdim, i) for i in range(len(indices))]
 

devito/finite_differences/tools.py

@@ -228,10 +228,14 @@ def make_stencil_dimension(expr, _min, _max):
 
 
 @cacheit
-def numeric_weights(function, deriv_order, indices, x0):
+def _numeric_weights(deriv_order, indices, x0):
     return finite_diff_weights(deriv_order, indices, x0)[-1][-1]
 
 
+def numeric_weights(function, deriv_order, indices, x0):
+    return _numeric_weights(deriv_order, indices, x0)
+
+
 fd_weights_registry = {'taylor': numeric_weights, 'standard': numeric_weights,
                        'symbolic': numeric_weights}  # Backward compat for 'symbolic'
 coeff_priority = {'taylor': 1, 'standard': 1}

devito/ir/cgen/printer.py

@@ -18,10 +18,12 @@
 from devito.arch.compiler import AOMPCompiler
 from devito.symbolics.inspection import has_integer_args, sympy_dtype
 from devito.symbolics.queries import q_leaf
-from devito.tools import ctypes_to_cstr, ctypes_vector_mapper, dtype_to_ctype
+from devito.tools import (
+    ctypes_to_cstr, ctypes_vector_mapper, dtype_to_ctype, memoized_func
+)
 from devito.types.basic import AbstractFunction
 
-__all__ = ['BasePrinter', 'ccode']
+__all__ = ['BasePrinter', 'ccode', 'get_printer']
 
 
 class BasePrinter(CodePrinter):
@@ -449,15 +451,20 @@ def _print_Fallback(self, expr):
     sympy.printing.str.StrPrinter._print_Add = BasePrinter._print_Add
 
 
-def ccode(expr, printer=None, **settings):
+@memoized_func
+def get_printer(printer, dtype):
+    return printer(settings={'dtype': dtype})
+
+
+def ccode(expr, printer=None, dtype=None):
     """Generate C++ code from an expression.
 
     Parameters
     ----------
     expr : expr-like
         The expression to be printed.
-    settings : dict
-        Options for code printing.
+    dtype : data-type, optional
+        Data type used by the printer.
 
     Returns
     -------
@@ -468,4 +475,5 @@ def ccode(expr, printer=None, **settings):
     if printer is None:
         from devito.passes.iet.languages.C import CPrinter
         printer = CPrinter
-    return printer(settings=settings).doprint(expr, None)
+    dtype = printer._default_settings['dtype'] if dtype is None else dtype
+    return get_printer(printer, dtype).doprint(expr, None)

devito/ir/clusters/analysis.py

@@ -101,7 +101,7 @@ def _callback(self, clusters, dim, prefix):
         is_parallel_atomic = False
 
         scope = Scope(flatten(c.exprs for c in clusters))
-        for dep in scope.d_all_gen():
+        for dep in scope.d_all_gen(writes=scope.writes_tensor):
             test00 = dep.is_indep(dim) and not dep.is_storage_related(dim)
             test01 = all(dep.is_reduce_atmost(i) for i in prev)
             if test00 and test01:
@@ -112,10 +112,6 @@ def _callback(self, clusters, dim, prefix):
                 is_parallel_indep &= (dep.distance_mapper.get(dim.root) == 0)
                 continue
 
-            if dep.function in scope.initialized:
-                # False alarm, the dependence is over a locally-defined symbol
-                continue
-
             if dep.is_reduction:
                 is_parallel_atomic = True
                 continue