Merge pull request #70 from pypr/nested-profile

Add minimal support for nested profiling.

compyle/api.py

@@ -13,7 +13,7 @@
 )
 from .profile import (
     get_profile_info, named_profile, profile, profile_ctx, print_profile,
-    profile_kernel
+    profile_kernel, ProfileContext, profile2csv
 )
 from .translator import (
     CConverter, CStructHelper, OpenCLConverter, detect_type, ocl_detect_type,

compyle/profile.py

@@ -7,13 +7,17 @@
 from .config import get_config
 
 
-_profile_info = defaultdict(lambda: {'calls': 0, 'time': 0})
+_current_level = 0
+_profile_info = defaultdict(
+    lambda: defaultdict(lambda: {'calls': 0, 'time': 0})
+)
 
 
 def _record_profile(name, time):
-    global _profile_info
-    _profile_info[name]['time'] += time
-    _profile_info[name]['calls'] += 1
+    global _profile_info, _current_level
+    li = _profile_info[_current_level]
+    li[name]['time'] += time
+    li[name]['calls'] += 1
 
 
 @contextmanager
@@ -25,9 +29,14 @@ def profile_ctx(name):
     with profile_ctx('f'):
         f()
     """
+    global _current_level
+    _current_level += 1
     start = time.time()
-    yield start
-    end = time.time()
+    try:
+        yield start
+        end = time.time()
+    finally:
+        _current_level -= 1
     _record_profile(name, end - start)
 
 
@@ -65,34 +74,83 @@ def wrapper(*args, **kwargs):
         return make_wrapper(method)
 
 
+class ProfileContext:
+    """Used for a low-level profiling context.
+
+    This is typically useful in Cython code where decorators are not usable and
+    using a context manager makes the code hard to read.
+
+    Example
+    -------
+
+    p = ProfileContext('some_func')
+    do_something()
+    p.stop()
+
+    """
+    def __init__(self, name):
+        self.name = name
+        global _current_level
+        _current_level += 1
+        self.start = time.time()
+
+    def stop(self):
+        global _current_level
+        _current_level -= 1
+        _record_profile(self.name, time.time() - self.start)
+
+
 def get_profile_info():
     global _profile_info
     return _profile_info
 
 
 def print_profile():
     global _profile_info
-    profile_data = sorted(_profile_info.items(), key=lambda x: x[1]['time'],
-                          reverse=True)
     hr = '-'*70
     print(hr)
     if len(_profile_info) == 0:
         print("No profiling information available")
         print(hr)
         return
     print("Profiling info:")
-    print("{:<40} {:<10} {:<10}".format('Function', 'N calls', 'Time'))
+    print(
+        "{:<6} {:<40} {:<10} {:<10}".format(
+            'Level', 'Function', 'N calls', 'Time')
+    )
     tot_time = 0
-    for kernel, data in profile_data:
-        print("{:<40} {:<10} {:<10}".format(
-            kernel,
-            data['calls'],
-            data['time']))
-        tot_time += data['time']
+    for level in range(0, min(len(_profile_info), 2)):
+        profile_data = sorted(
+            _profile_info[level].items(), key=lambda x: x[1]['time'],
+            reverse=True
+        )
+        for kernel, data in profile_data:
+            print("{:<6} {:<40} {:<10} {:<10.3g}".format(
+                level, kernel, data['calls'], data['time'])
+            )
+            if level == 0:
+                tot_time += data['time']
     print("Total profiled time: %g secs" % tot_time)
     print(hr)
 
 
+def profile2csv(fname):
+    info = get_profile_info()
+    with open(fname, 'w') as f:
+        f.write("{0},{1},{2},{3}\n".format(
+            'level', 'function', 'calls', 'time')
+        )
+        for level in sorted(info.keys()):
+            profile_data = sorted(
+                info[level].items(), key=lambda x: x[1]['time'],
+                reverse=True
+            )
+            for name, data in profile_data:
+                f.write("{0},{1},{2},{3}\n".format(
+                    level, name, data['calls'], data['time']
+                ))
+
+
 def profile_kernel(kernel, name, backend=None):
     """For profiling raw PyCUDA/PyOpenCL kernels or cython functions
     """

compyle/tests/test_profile.py

@@ -5,7 +5,9 @@
 
 from ..config import get_config, use_config
 from ..array import wrap, zeros, ones
-from ..profile import get_profile_info, named_profile, profile, profile_ctx
+from ..profile import (
+    get_profile_info, named_profile, profile, profile_ctx, ProfileContext
+)
 
 
 def axpb():
@@ -38,6 +40,11 @@ def profiled_axpb():
     axpb()
 
 
+@profile
+def nested():
+    profiled_axpb()
+
+
 @named_profile('prefix_sum', backend='opencl')
 def get_prefix_sum_knl():
     from ..opencl import get_queue, get_context
@@ -56,15 +63,15 @@ def test_profile_ctx():
         axpb()
 
     profile_info = get_profile_info()
-    assert profile_info['axpb']['calls'] == 1
+    assert profile_info[0]['axpb']['calls'] == 1
 
 
 def test_profile():
     for i in range(100):
         profiled_axpb()
 
     profile_info = get_profile_info()
-    assert profile_info['profiled_axpb']['calls'] == 100
+    assert profile_info[0]['profiled_axpb']['calls'] == 100
 
 
 def test_profile_method():
@@ -80,13 +87,13 @@ def test_profile_method():
 
     # Then
     profile_info = get_profile_info()
-    assert profile_info['A.f']['calls'] == 5
+    assert profile_info[0]['A.f']['calls'] == 5
 
     # For b.f(), b.name is my_name.
-    assert profile_info['my_name']['calls'] == 5
+    assert profile_info[0]['my_name']['calls'] == 5
 
     # profile was given an explicit name for b.named()
-    assert profile_info['explicit_name']['calls'] == 5
+    assert profile_info[0]['explicit_name']['calls'] == 5
 
 
 def test_named_profile():
@@ -97,4 +104,18 @@ def test_named_profile():
     knl(x.dev)
 
     profile_info = get_profile_info()
-    assert profile_info['prefix_sum']['calls'] == 1
+    assert profile_info[0]['prefix_sum']['calls'] == 1
+
+
+def test_nesting_and_context():
+    # When
+    p = ProfileContext('main')
+    nested()
+    p.stop()
+
+    # Then
+    prof = get_profile_info()
+    assert len(prof) == 3
+    assert prof[0]['main']['calls'] == 1
+    assert prof[1]['nested']['calls'] == 1
+    assert prof[2]['profiled_axpb']['calls'] == 1