Add compile-time coverage for compiled code

buildscripts/azure/azure-linux-macos.yml

@@ -49,6 +49,12 @@ jobs:
         export PATH=$HOME/miniconda3/bin:$PATH
         buildscripts/incremental/test.sh
       displayName: 'Test'
+
+    - task: PublishCodeCoverageResults@2
+      condition: eq(variables['RUN_COVERAGE'], 'yes')
+      inputs:
+        codeCoverageTool: Cobertura
+        summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/cov.xml'
 
     - task: PublishTestResults@2
       condition: succeededOrFailed()

buildscripts/incremental/test.sh

@@ -134,9 +134,13 @@ python -m numba.runtests -l
 # directive in .coveragerc
 echo "INFO: Running shard of discovered tests: ($TEST_START_INDEX:$TEST_COUNT)"
 if [ "$RUN_COVERAGE" == "yes" ]; then
+    echo "INFO: Running with coverage"
     export PYTHONPATH=.
     coverage erase
     $SEGVCATCH coverage run runtests.py -b -j "$TEST_START_INDEX:$TEST_COUNT" --exclude-tags='long_running' -m $TEST_NPROCS --junit -- numba.tests
+    echo "INFO: Post-process coverage"
+    coverage combine
+    coverage xml -o cov.xml
 elif [ "$RUN_TYPEGUARD" == "yes" ]; then
     echo "INFO: Running with typeguard"
     NUMBA_USE_TYPEGUARD=1 NUMBA_ENABLE_CUDASIM=1 PYTHONWARNINGS="ignore:::typeguard" $SEGVCATCH python runtests.py -b -j "$TEST_START_INDEX:$TEST_COUNT" --exclude-tags='long_running' -m $TEST_NPROCS --junit -- numba.tests

docs/source/user/code_coverage.rst

@@ -0,0 +1,24 @@
+===============================
+Code Coverage for Compiled Code
+===============================
+
+Numba, a just-in-time compiler for Python, transforms Python code into machine 
+code for optimized execution. This process, however, poses a challenge for 
+traditional code coverage tools, as they typically operate within the Python 
+interpreter and thus miss the lines of code compiled by Numba. To address this 
+issue, Numba opts for a compile-time notification to coverage tools, rather than 
+during execution, to minimize performance penalties. This approach helps prevent 
+significant coverage gaps in projects utilizing Numba, without incurring 
+substantial performance costs.
+
+No additional effort is required to generate compile-time coverage data. By 
+running a Numba application under the ``coverage`` tool 
+(e.g. ``coverage run ...``), the compiler automatically 
+detects the active coverage session and emits data accordingly. This mechanism 
+ensures that coverage data is generated seamlessly, without the need for manual 
+intervention.
+
+The coverage data is emitted during the lowering phase, which involves the 
+generation of LLVM-IR. This phase inherently excludes lines of code that are 
+statically identified as dead code, ensuring that the coverage data accurately 
+reflects the executable code paths.

docs/source/user/index.rst

@@ -19,6 +19,7 @@ User Manual
    performance-tips.rst
    threading-layer.rst
    cli.rst
+   code_coverage.rst
    troubleshoot.rst
    faq.rst
    examples.rst

docs/upcoming_changes/9508.new_feature.rst

@@ -0,0 +1,7 @@
+Added compile-time code coverage
+--------------------------------
+
+Support for emitting compile-time coverage data is added.
+This feature is automatically activated when running Python under ``coverage``. 
+It collects data during the compiler's lowering phase, showing source lines 
+compiled into LLVM-IR, excluding dead-code eliminated lines. 

numba/core/lowering.py

@@ -16,6 +16,7 @@
 from numba.core.environment import Environment
 from numba.core.analysis import compute_use_defs, must_use_alloca
 from numba.misc.firstlinefinder import get_func_body_first_lineno
+from numba.misc.coverage_support import get_registered_loc_notify
 
 
 _VarArgItem = namedtuple("_VarArgItem", ("vararg", "index"))
@@ -69,6 +70,9 @@ def __init__(self, context, library, fndesc, func_ir, metadata=None):
                                       cgctx=context,
                                       directives_only=directives_only)
 
+        # Loc notify objects
+        self._loc_notify_registry = get_registered_loc_notify()
+
         # Subclass initialization
         self.init()
 
@@ -139,6 +143,8 @@ def post_lower(self):
         Called after all blocks are lowered
         """
         self.debuginfo.finalize()
+        for notify in self._loc_notify_registry:
+            notify.close()
 
     def pre_block(self, block):
         """
@@ -307,6 +313,13 @@ def setup_function(self, fndesc):
     def typeof(self, varname):
         return self.fndesc.typemap[varname]
 
+    def notify_loc(self, loc: ir.Loc) -> None:
+        """Called when a new instruction with the given `loc` is about to be
+        lowered.
+        """
+        for notify_obj in self._loc_notify_registry:
+            notify_obj.notify(loc)
+
     def debug_print(self, msg):
         if config.DEBUG_JIT:
             self.context.debug_print(
@@ -442,6 +455,7 @@ def post_block(self, block):
     def lower_inst(self, inst):
         # Set debug location for all subsequent LL instructions
         self.debuginfo.mark_location(self.builder, self.loc.line)
+        self.notify_loc(self.loc)
         self.debug_print(str(inst))
         if isinstance(inst, ir.Assign):
             ty = self.typeof(inst.target.name)

numba/misc/coverage_support.py

@@ -0,0 +1,94 @@
+"""
+Implement code coverage support.
+
+Currently contains logic to extend ``coverage`` with line covered by the
+compiler.
+"""
+from typing import Optional, Sequence, Callable
+from collections import defaultdict
+from abc import ABC, abstractmethod
+import atexit
+from functools import cache
+
+from numba.core import ir
+
+
+try:
+    import coverage
+except ImportError:
+    coverage_available = False
+else:
+    coverage_available = True
+
+
+def get_active_coverage() -> Optional["coverage.Coverage"]:
+    """Get active coverage instance or return None if not found.
+    """
+    cov = None
+    if coverage_available:
+        cov = coverage.Coverage.current()
+    return cov
+
+
+_the_registry: Callable[[], Optional["NotifyLocBase"]] = []
+
+
+def get_registered_loc_notify() -> Sequence["NotifyLocBase"]:
+    return list(filter(lambda x: x is not None,
+                       (factory() for factory in _the_registry)))
+
+
+@cache
+def _get_coverage_data():
+    """
+    Make a singleton ``CoverageData``.
+    Avoid writing to disk. Other processes can corrupt the file.
+    """
+    covdata = coverage.CoverageData(no_disk=True)
+    cov = get_active_coverage()
+    assert cov is not None, "no active Coverage instance"
+
+    @atexit.register
+    def _finalize():
+        cov.get_data().update(covdata)
+
+    return covdata
+
+
+class NotifyLocBase(ABC):
+    """Interface for notifying visiting of a ``numba.core.ir.Loc``.
+    """
+    @abstractmethod
+    def notify(self, loc: ir.Loc) -> None:
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        pass
+
+
+class NotifyCompilerCoverage(NotifyLocBase):
+    """
+    Use to notify ``coverage`` about compiled lines.
+
+    The compiled lines under "numba_compiled" context in the coverage data.
+    """
+    def __init__(self):
+        self._arcs_data = defaultdict(set)
+
+    def notify(self, loc: ir.Loc):
+        if loc.filename.endswith(".py"):
+            # The compiler doesn't actually know about arc.
+            self._arcs_data[loc.filename].add((loc.line, loc.line))
+
+    def close(self):
+        covdata = _get_coverage_data()
+        with covdata._lock:
+            covdata.set_context("numba_compiled")
+            covdata.add_arcs(self._arcs_data)
+
+
+@_the_registry.append
+def _register_coverage_notifier():
+    if get_active_coverage() is not None:
+        return NotifyCompilerCoverage()

numba/tests/test_caching.py

@@ -385,9 +385,14 @@ def test_looplifted(self):
             self.assertPreciseEqual(f(4), 6)
             self.check_pycache(0)
 
-        self.assertEqual(len(w), 1)
-        self.assertIn('Cannot cache compiled function "looplifted" '
-                      'as it uses lifted code', str(w[0].message))
+        try:
+            self.assertEqual(len(w), 1)
+            self.assertIn('Cannot cache compiled function "looplifted" '
+                          'as it uses lifted code', str(w[0].message))
+        except Exception:
+            print("Dump warnings")
+            for warnobj in w:
+                print(warnobj)
 
     def test_big_array(self):
         # Code references big array globals cannot be cached
@@ -415,11 +420,16 @@ def test_ctypes(self):
                 self.assertPreciseEqual(f(0.0), 0.0)
                 self.check_pycache(0)
 
-            self.assertEqual(len(w), 1)
-            self.assertIn(
-                'Cannot cache compiled function "{}"'.format(f.__name__),
-                str(w[0].message),
-            )
+            try:
+                self.assertEqual(len(w), 1)
+                self.assertIn(
+                    'Cannot cache compiled function "{}"'.format(f.__name__),
+                    str(w[0].message),
+                )
+            except Exception:
+                print("Dump warnings")
+                for warnobj in w:
+                    print(warnobj)
 
     def test_closure(self):
         mod = self.import_module()

numba/tests/test_linalg.py

@@ -60,12 +60,18 @@ def check_contiguity_warning(self, pyfunc):
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter('always', errors.NumbaPerformanceWarning)
             yield
-        self.assertGreaterEqual(len(w), 1)
-        self.assertIs(w[0].category, errors.NumbaPerformanceWarning)
-        self.assertIn("faster on contiguous arrays", str(w[0].message))
-        self.assertEqual(w[0].filename, pyfunc.__code__.co_filename)
-        # This works because our functions are one-liners
-        self.assertEqual(w[0].lineno, pyfunc.__code__.co_firstlineno + 1)
+        try:
+            self.assertGreaterEqual(len(w), 1)
+            self.assertIs(w[0].category, errors.NumbaPerformanceWarning)
+            self.assertIn("faster on contiguous arrays", str(w[0].message))
+            self.assertEqual(w[0].filename, pyfunc.__code__.co_filename)
+            # This works because our functions are one-liners
+            self.assertEqual(w[0].lineno, pyfunc.__code__.co_firstlineno + 1)
+            raise
+        except Exception:
+            print("Dump warnings")
+            for warnobj in w:
+                print(warnobj)
 
     def check_func(self, pyfunc, cfunc, args):
         with self.assertNoNRTLeak():

numba/tests/test_misc_coverage_support.py

@@ -0,0 +1,74 @@
+import unittest
+from unittest.mock import patch
+
+from numba.tests.support import TestCase
+
+from numba import njit
+from numba.core import ir
+from numba.misc.coverage_support import NotifyLocBase, _the_registry
+
+
+class TestMiscCoverageSupport(TestCase):
+    def test_custom_loc_notifier(self):
+        class MyNotify(NotifyLocBase):
+            records = []
+
+            def notify(self, loc):
+                self.records.append(("NOTIFY", loc))
+
+            def close(self):
+                self.records.append(("CLOSE", None))
+
+        # Patch to install registry for testing
+        new_the_registry = _the_registry + [MyNotify]
+        gv = "numba.misc.coverage_support._the_registry"
+        with patch(gv, new_the_registry):
+
+            @njit
+            def foo():
+                return 123
+
+            res = foo()
+
+        self.assertEqual(res, 123)
+
+        # offset by +2 because:
+        # +1 for the decorator
+        # +1 for the `def` line
+        first_offset = 2
+        offset = foo.__code__.co_firstlineno + first_offset
+        loc = ir.Loc(__file__, 1)
+        self.assertIn(("NOTIFY", loc.with_lineno(offset)), MyNotify.records)
+        self.assertIn(("CLOSE", None), MyNotify.records)
+
+        # Test dead branch pruned
+        with patch(gv, new_the_registry):
+            cond = False
+
+            @njit
+            def foo():
+                if cond:
+                    return 321
+                return 123
+
+            res = foo()
+
+        self.assertEqual(res, 123)
+
+        # `if cond` line is compiled
+        offset = foo.__code__.co_firstlineno + first_offset
+        self.assertIn(("NOTIFY", loc.with_lineno(offset)), MyNotify.records)
+
+        # `    return 321` line is not compiled
+        self.assertNotIn(
+            ("NOTIFY", loc.with_lineno(offset + 1)), MyNotify.records
+        )
+
+        # `    return 123` line is compiled
+        self.assertIn(("NOTIFY", loc.with_lineno(offset + 2)), MyNotify.records)
+
+        self.assertIn(("CLOSE", None), MyNotify.records)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba/tests/test_np_functions.py

@@ -6268,7 +6268,7 @@ def foo():
         result, error = run_in_subprocess(code)
         # Assert that the bytestring "OK" was printed to stdout
         self.assertEqual(b"OK", result.strip())
-        self.assertEqual(b"", error.strip())
+        self.assertEqual(b"", error.strip(), msg=f"--ERROR--\n{error}\n")
 
 
 if __name__ == '__main__':