Adding debug decorator to print Numba compile stats (#925)

* Adding debug decorator to print Numba compile stats

This adds a decorator to print Numba compilation stats for a decorated function
and all nested functions (literally all compiled overloads). Filtering
can be made using decorator arguments.

NOTE: this may show a bit lesser times than compile time calculated using
common approach via first_run_exec_time - second_run_exec_time.

* Adding tests for compile time log format

* Remove check for CPUDispatcher to work with numba DPPL

Author: kozlov-alexey

sdc/decorators.py

@@ -31,6 +31,9 @@
 import numba
 import sdc
 
+from functools import wraps
+from sdc.utilities.utils import print_compile_times
+
 
 def jit(signature_or_function=None, **options):
 
@@ -44,3 +47,29 @@ def jit(signature_or_function=None, **options):
     Use Numba compiler pipeline
     '''
     return numba.jit(signature_or_function, **options)
+
+
+def debug_compile_time(level=1, func_names=None):
+    """ Decorates Numba Dispatcher object to print compile stats after call.
+        Usage:
+            @debug_compile_time()
+            @numba.njit
+            <decorated function>
+        Args:
+            level: if zero prints only short summary
+            func_names: filters output to include only functions which names include listed strings,
+    """
+
+    def get_wrapper(disp):
+
+        @wraps(disp)
+        def wrapper(*args, **kwargs):
+            res = disp(*args, **kwargs)
+            print('*' * 40, 'COMPILE STATS', '*' * 40)
+            print_compile_times(disp, level=level, func_names=func_names)
+            print('*' * 95)
+            return res
+
+        return wrapper
+
+    return get_wrapper

sdc/tests/__init__.py

@@ -48,6 +48,7 @@
 
 from sdc.tests.test_sdc_numpy import *
 from sdc.tests.test_prange_utils import *
+from sdc.tests.test_compile_time import *
 
 # performance tests
 import sdc.tests.tests_perf

sdc/tests/test_compile_time.py

@@ -0,0 +1,114 @@
+# *****************************************************************************
+# Copyright (c) 2019-2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+import numba
+import numpy as np
+import pandas as pd
+import re
+import unittest
+
+from contextlib import redirect_stdout
+from io import StringIO
+from sdc.tests.test_base import TestCase
+from sdc.decorators import debug_compile_time
+
+
+# regexp patterns for lines in @debug_compile_time output log
+line_heading = r'\*+\s+COMPILE STATS\s+\*+\n'
+line_function = r'Function: [^\s]+\n'
+line_args = r'\s+Args:.*\n'
+line_pipeline = r'\s+Pipeline: \w+\n'
+line_passes = r'(\s+\w+\s+[\d.]+\n)+'
+line_time = r'\s+Time: [\d.]+\n'
+line_ending = r'\*+\n'
+
+
+class TestCompileTime(TestCase):
+
+    @staticmethod
+    def _gen_usecase_data():
+        n = 11
+        S1 = pd.Series(np.ones(n))
+        S2 = pd.Series(2 ** np.arange(n))
+        return S1, S2
+
+    def test_log_format_summary(self):
+        """ Verifies shortened log format when only summary info is printed """
+
+        @debug_compile_time(level=0)
+        @self.jit
+        def test_impl(S1, S2):
+            return S1 + S2
+
+        buffer = StringIO()
+        with redirect_stdout(buffer):
+            S1, S2 = self._gen_usecase_data()
+            test_impl(S1, S2)
+
+        entry_format = fr'{line_function}{line_pipeline}{line_time}\n'
+        log_format = fr'^{line_heading}({entry_format})+{line_ending}$'
+        self.assertRegex(buffer.getvalue(), log_format)
+
+    def test_log_format_detailed(self):
+        """ Verifies detailed log format with passes and args information """
+
+        @debug_compile_time()
+        @self.jit
+        def test_impl(S1, S2):
+            return S1 + S2
+
+        buffer = StringIO()
+        with redirect_stdout(buffer):
+            S1, S2 = self._gen_usecase_data()
+            test_impl(S1, S2)
+
+        entry_format = fr'{line_function}{line_args}{line_pipeline}{line_passes}{line_time}\n'
+        log_format = fr'{line_heading}({entry_format})+{line_ending}'
+        self.assertRegex(buffer.getvalue(), log_format)
+
+    def test_func_names_filter(self):
+        """ Verifies filtering log entries via func_names paramter """
+        searched_name = 'add'
+
+        @debug_compile_time(func_names=[searched_name])
+        @self.jit
+        def test_impl(S1, S2):
+            return S1 + S2
+
+        buffer = StringIO()
+        with redirect_stdout(buffer):
+            S1, S2 = self._gen_usecase_data()
+            test_impl(S1, S2)
+
+        line_function = r'Function: ([^\s]+)\n'
+        match_iter = re.finditer(line_function, buffer.getvalue())
+        next(match_iter)    # skip entry for top-level func
+        for m in match_iter:
+            self.assertIn(searched_name, m.group(1))
+
+
+if __name__ == "__main__":
+    unittest.main()

sdc/utilities/utils.py

@@ -51,6 +51,8 @@
 from numba.extending import overload, overload_method, overload_attribute
 from numba.extending import register_jitable, register_model
 from numba.core.datamodel.registry import register_default
+from functools import wraps
+from itertools import filterfalse, chain
 
 
 # int values for types to pass to C code
@@ -686,3 +688,59 @@ def sdc_overload_attribute(typ, name, jit_options={}, parallel=None, strict=True
     return overload_attribute(
         typ, name, jit_options=jit_options, strict=strict, inline=inline, prefer_literal=prefer_literal
     )
+
+
+def print_compile_times(disp, level, func_names=None):
+
+    def print_times(cres, args):
+        print(f'Function: {cres.fndesc.unique_name}')
+        pad = '  ' * 2
+        if level:
+            print(f'{pad * 1}Args: {args}')
+        times = cres.metadata['pipeline_times']
+        for pipeline, pass_times in times.items():
+            print(f'{pad * 1}Pipeline: {pipeline}')
+            if level:
+                for name, timings in pass_times.items():
+                    print(f'{pad * 2}{name:50s}{timings.run:.13f}')
+
+            pipeline_total = sum(t.init + t.run + t.finalize for t in pass_times.values())
+            print(f'{pad * 1}Time: {pipeline_total}\n')
+
+    # print times for compiled function indicated by disp
+    for args, cres in disp.overloads.items():
+        print_times(cres, args)
+
+    def has_no_cache(ovld):
+        return not (getattr(ovld, '_impl_cache', False) and ovld._impl_cache)
+
+    known_funcs = disp.typingctx._functions
+    all_templs = chain.from_iterable(known_funcs.values())
+    compiled_templs = filterfalse(has_no_cache, all_templs)
+
+    # filter only function names that are in the func_names list
+    if func_names:
+        compiled_templs = filterfalse(
+                                lambda x: not any(f in str(x) for f in func_names),
+                                compiled_templs
+                            )
+
+    dispatchers_list = []
+    for template in compiled_templs:
+        tmpl_cached_impls = template._impl_cache.values()
+        dispatchers_list.extend(tmpl_cached_impls)
+
+    for impl_cache in set(dispatchers_list):
+        # impl_cache is usually a tuple of format (dispatcher, args)
+        # if not just skip these entires
+        if not (isinstance(impl_cache, tuple)
+                and len(impl_cache) == 2
+                and isinstance(impl_cache[0], type(disp))):
+            continue
+
+        fndisp, args = impl_cache
+        if not getattr(fndisp, 'overloads', False):
+            continue
+
+        cres, = list(fndisp.overloads.values())
+        print_times(cres, args)