Merge pull request #2076 from devitocodes/hotfix-minmax

compiler: Fix and improve MIN/MAX codegen

Author: FabioLuporini

devito/finite_differences/elementary.py

@@ -93,17 +93,13 @@ def root(x):
 class Min(sympy.Min, Evaluable):
 
     def _evaluate(self, **kwargs):
-        args = self._evaluate_args(**kwargs)
-        assert len(args) == 2
-        return self.func(args[0], args[1], evaluate=False)
+        return self.func(*self._evaluate_args(**kwargs), evaluate=False)
 
 
 class Max(sympy.Max, Evaluable):
 
     def _evaluate(self, **kwargs):
-        args = self._evaluate_args(**kwargs)
-        assert len(args) == 2
-        return self.func(args[0], args[1], evaluate=False)
+        return self.func(*self._evaluate_args(**kwargs), evaluate=False)
 
 
 def Id(x):

devito/ir/iet/visitors.py

@@ -11,6 +11,7 @@
 
 import cgen as c
 from sympy import IndexedBase
+from sympy.core.function import Application
 
 from devito.exceptions import VisitorException
 from devito.ir.iet.nodes import (Node, Iteration, Expression, ExpressionBundle,
@@ -24,9 +25,9 @@
                           IndexedData, DeviceMap)
 
 
-__all__ = ['FindNodes', 'FindSections', 'FindSymbols', 'MapExprStmts', 'MapNodes',
-           'IsPerfectIteration', 'printAST', 'CGen', 'CInterface', 'Transformer',
-           'Uxreplace']
+__all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols',
+           'MapExprStmts', 'MapNodes', 'IsPerfectIteration', 'printAST', 'CGen',
+           'CInterface', 'Transformer', 'Uxreplace']
 
 
 class Visitor(GenericVisitor):
@@ -953,6 +954,55 @@ def visit_Node(self, o, ret=None):
         return ret
 
 
+class FindApplications(Visitor):
+
+    """
+    Find all SymPy applied functions (aka, `Application`s). The user may refine
+    the search by supplying a different target class.
+    """
+
+    def __init__(self, cls=Application):
+        super().__init__()
+        self.match = lambda i: isinstance(i, cls) and not isinstance(i, Basic)
+
+    @classmethod
+    def default_retval(cls):
+        return set()
+
+    def visit_object(self, o, **kwargs):
+        return self.default_retval()
+
+    def visit_tuple(self, o, ret=None):
+        ret = ret or self.default_retval()
+        for i in o:
+            ret.update(self._visit(i, ret=ret))
+        return ret
+
+    def visit_Node(self, o, ret=None):
+        ret = ret or self.default_retval()
+        for i in o.children:
+            ret.update(self._visit(i, ret=ret))
+        return ret
+
+    def visit_Expression(self, o, **kwargs):
+        return o.expr.find(self.match)
+
+    def visit_Iteration(self, o, **kwargs):
+        ret = self._visit(o.children) or self.default_retval()
+        ret.update(o.symbolic_min.find(self.match))
+        ret.update(o.symbolic_max.find(self.match))
+        return ret
+
+    def visit_Call(self, o, **kwargs):
+        ret = self.default_retval()
+        for i in o.arguments:
+            try:
+                ret.update(i.find(self.match))
+            except (AttributeError, TypeError):
+                continue
+        return ret
+
+
 class IsPerfectIteration(Visitor):
 
     """

devito/mpi/routines.py

@@ -58,13 +58,6 @@ def msgs(self):
     def regions(self):
         return [i for i in self._regions.values() if i is not None]
 
-    @property
-    def headers(self):
-        """
-        No headers needed by default
-        """
-        return {}
-
     def make(self, hs):
         """
         Construct Callables and Calls implementing distributed-memory halo
@@ -517,14 +510,6 @@ class OverlapHaloExchangeBuilder(DiagHaloExchangeBuilder):
         remainder()
     """
 
-    @property
-    def headers(self):
-        """
-        Overlap Mode uses MIN/MAX that need to be defined
-        """
-        return {'headers': [('MIN(a,b)', ('(((a) < (b)) ? (a) : (b))')),
-                            ('MAX(a,b)', ('(((a) > (b)) ? (a) : (b))'))]}
-
     def _make_msg(self, f, hse, key):
         # Only retain the halos required by the Diag scheme
         halos = sorted(i for i in hse.halos if isinstance(i.dim, tuple))

devito/operator/operator.py

@@ -19,7 +19,8 @@
 from devito.operator.registry import operator_selector
 from devito.mpi import MPI
 from devito.parameters import configuration
-from devito.passes import Graph, lower_index_derivatives, generate_implicit, instrument
+from devito.passes import (Graph, lower_index_derivatives, generate_implicit,
+                           generate_macros, instrument)
 from devito.symbolics import estimate_cost
 from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_tuple, flatten,
                           filter_sorted, frozendict, is_integer, split, timed_pass,
@@ -450,6 +451,9 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):
         # specialization further Sections may be introduced
         instrument(graph, profiler=profiler, sregistry=sregistry)
 
+        # Extract the necessary macros from the symbolic objects
+        generate_macros(graph)
+
         return graph.root, graph
 
     # Read-only properties exposed to the outside world

devito/passes/iet/misc.py

@@ -1,12 +1,17 @@
+from functools import singledispatch
+
 import cgen
+import sympy
 
-from devito.ir import (Any, Forward, List, Prodder, FindNodes, Transformer,
-                       filter_iterations, retrieve_iteration_tree)
+from devito.finite_differences import Max, Min
+from devito.ir import (Any, Forward, List, Prodder, FindApplications, FindNodes,
+                       Transformer, filter_iterations, retrieve_iteration_tree)
 from devito.passes.iet.engine import iet_pass
-from devito.symbolics import evalrel
+from devito.symbolics import evalrel, has_integer_args
 from devito.tools import split
 
-__all__ = ['avoid_denormals', 'hoist_prodders', 'relax_incr_dimensions']
+__all__ = ['avoid_denormals', 'hoist_prodders', 'relax_incr_dimensions',
+           'generate_macros']
 
 
 @iet_pass
@@ -124,9 +129,35 @@ def relax_incr_dimensions(iet, options=None, **kwargs):
     if mapper:
         iet = Transformer(mapper, nested=True).visit(iet)
 
-        headers = [('MIN(a,b)', ('(((a) < (b)) ? (a) : (b))')),
-                   ('MAX(a,b)', ('(((a) > (b)) ? (a) : (b))'))]
-    else:
-        headers = []
+    return iet, {}
+
+
+@iet_pass
+def generate_macros(iet):
+    applications = FindApplications().visit(iet)
+    headers = set().union(*[_generate_macros(i) for i in applications])
 
     return iet, {'headers': headers}
+
+
+@singledispatch
+def _generate_macros(expr):
+    return set()
+
+
+@_generate_macros.register(Min)
+@_generate_macros.register(sympy.Min)
+def _(expr):
+    if has_integer_args(*expr.args) and len(expr.args) == 2:
+        return {('MIN(a,b)', ('(((a) < (b)) ? (a) : (b))'))}
+    else:
+        return set()
+
+
+@_generate_macros.register(Max)
+@_generate_macros.register(sympy.Max)
+def _(expr):
+    if has_integer_args(*expr.args) and len(expr.args) == 2:
+        return {('MAX(a,b)', ('(((a) > (b)) ? (a) : (b))'))}
+    else:
+        return set()

devito/passes/iet/mpi.py

@@ -292,7 +292,6 @@ def make_mpi(iet, mpimode=None, **kwargs):
 
     efuncs = sync_heb.efuncs + user_heb.efuncs
     iet = Transformer(mapper, nested=True).visit(iet)
-    headers = user_heb.headers
 
     # Must drop the PARALLEL tag from the Iterations within which halo
     # exchanges are performed
@@ -307,9 +306,8 @@ def make_mpi(iet, mpimode=None, **kwargs):
                                for n in tree[:tree.index(i)+1]})
                 break
     iet = Transformer(mapper, nested=True).visit(iet)
-    headers.update({'includes': ['mpi.h'], 'efuncs': efuncs})
 
-    return iet, headers
+    return iet, {'includes': ['mpi.h'], 'efuncs': efuncs}
 
 
 def mpiize(graph, **kwargs):

devito/symbolics/inspection.py

@@ -11,7 +11,7 @@
 from devito.symbolics.queries import q_routine
 from devito.tools import as_tuple, prod
 
-__all__ = ['compare_ops', 'estimate_cost']
+__all__ = ['compare_ops', 'estimate_cost', 'has_integer_args']
 
 
 def compare_ops(e1, e2):
@@ -235,3 +235,28 @@ def _(expr, estimate):
     flops *= prod(i._size for i in expr.dimensions)
 
     return flops, False
+
+
+def has_integer_args(*args):
+    """
+    True if all `args` are of integer type, False otherwise.
+    """
+    if len(args) == 0:
+        return False
+
+    if len(args) == 1:
+        try:
+            return np.issubdtype(args[0].dtype, np.integer)
+        except AttributeError:
+            return args[0].is_integer
+
+    res = True
+    for a in args:
+        try:
+            if len(a.args) > 0:
+                res = res and has_integer_args(*a.args)
+            else:
+                res = res and has_integer_args(a)
+        except AttributeError:
+            res = res and has_integer_args(a)
+    return res

devito/symbolics/printer.py

@@ -12,6 +12,7 @@
 from sympy.printing.c import C99CodePrinter
 
 from devito.arch.compiler import AOMPCompiler
+from devito.symbolics.inspection import has_integer_args
 
 __all__ = ['ccode']
 
@@ -102,14 +103,16 @@ def _print_Mod(self, expr):
         return '%'.join(args)
 
     def _print_Min(self, expr):
-        """Print Min using devito defined header Min"""
-        func = 'MIN' if has_integer_args(*expr.args) else 'fmin'
-        return "%s(%s)" % (func, self._print(expr.args)[1:-1])
+        if has_integer_args(*expr.args) and len(expr.args) == 2:
+            return "MIN(%s)" % self._print(expr.args)[1:-1]
+        else:
+            return super()._print_Min(expr)
 
     def _print_Max(self, expr):
-        """Print Max using devito defined header Max"""
-        func = 'MAX' if has_integer_args(*expr.args) else 'fmax'
-        return "%s(%s)" % (func, self._print(expr.args)[1:-1])
+        if has_integer_args(*expr.args) and len(expr.args) == 2:
+            return "MAX(%s)" % self._print(expr.args)[1:-1]
+        else:
+            return super()._print_Max(expr)
 
     def _print_Abs(self, expr):
         """Print an absolute value. Use `abs` if can infer it is an Integer"""
@@ -258,30 +261,3 @@ def ccode(expr, **settings):
 # to always use the correct one from our printer
 if Version(sympy.__version__) >= Version("1.11"):
     setattr(sympy.printing.str.StrPrinter, '_print_Add', CodePrinter._print_Add)
-
-
-# Check arguements type
-def has_integer_args(*args):
-    """
-    Check if expression is Integer.
-    Used to choose the function printed in the c-code
-    """
-    if len(args) == 0:
-        return False
-
-    if len(args) == 1:
-        try:
-            return np.issubdtype(args[0].dtype, np.integer)
-        except AttributeError:
-            return args[0].is_integer
-
-    res = True
-    for a in args:
-        try:
-            if len(a.args) > 0:
-                res = res and has_integer_args(*a.args)
-            else:
-                res = res and has_integer_args(a)
-        except AttributeError:
-            res = res and has_integer_args(a)
-    return res

examples/performance/00_overview.ipynb

@@ -1184,10 +1184,9 @@
      "output_type": "stream",
      "text": [
       "#define _POSIX_C_SOURCE 200809L\n",
-      "#define MIN(a,b) (((a) < (b)) ? (a) : (b))\n",
-      "#define MAX(a,b) (((a) > (b)) ? (a) : (b))\n",
       "#define START_TIMER(S) struct timeval start_ ## S , end_ ## S ; gettimeofday(&start_ ## S , NULL);\n",
       "#define STOP_TIMER(S,T) gettimeofday(&end_ ## S, NULL); T->S += (double)(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;\n",
+      "#define MIN(a,b) (((a) < (b)) ? (a) : (b))\n",
       "\n",
       "#include \"stdlib.h\"\n",
       "#include \"math.h\"\n",
@@ -1622,10 +1621,9 @@
      "output_type": "stream",
      "text": [
       "#define _POSIX_C_SOURCE 200809L\n",
-      "#define MIN(a,b) (((a) < (b)) ? (a) : (b))\n",
-      "#define MAX(a,b) (((a) > (b)) ? (a) : (b))\n",
       "#define START_TIMER(S) struct timeval start_ ## S , end_ ## S ; gettimeofday(&start_ ## S , NULL);\n",
       "#define STOP_TIMER(S,T) gettimeofday(&end_ ## S, NULL); T->S += (double)(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;\n",
+      "#define MIN(a,b) (((a) < (b)) ? (a) : (b))\n",
       "\n",
       "#include \"stdlib.h\"\n",
       "#include \"math.h\"\n",

tests/test_linearize.py

@@ -169,9 +169,9 @@ def test_codegen_quality0():
     assert all('const long' in str(i) for i in exprs[:-2])
 
     # Only four access macros necessary, namely `uL0`, `bufL0`, `bufL1`
-    # MIN/MAX for the efunc args
+    # for the efunc args
     # (the other three obviously are _POSIX_C_SOURCE, START_TIMER, STOP_TIMER)
-    assert len(op._headers) == 8
+    assert len(op._headers) == 6
 
 
 def test_codegen_quality1():
@@ -193,7 +193,7 @@ def test_codegen_quality1():
 
     # Only two access macros necessary, namely `uL0` and `r1L0` (the other five
     # obviously are _POSIX_C_SOURCE, MIN, MAX, START_TIMER, STOP_TIMER)
-    assert len(op._headers) == 7
+    assert len(op._headers) == 6
 
 
 def test_pow():