[Misc] Add coverage report to PRs, including kernels

.github/workflows/linux.yml

@@ -6,6 +6,8 @@ on:
       - reopened
       - synchronize
   workflow_dispatch:
+permissions:
+  contents: read
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: true
@@ -21,6 +23,8 @@ jobs:
     runs-on: ${{ matrix.OS }}
     steps:
       - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
       - name: Python check
         uses: actions/setup-python@v4
         with:
@@ -37,3 +41,109 @@ jobs:
       - name: Linux test
         run: |
           bash .github/workflows/scripts_new/linux/4_test.sh
+      - name: Upload wheel for CUDA job
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: linux-wheel
+          path: dist/*.whl
+      - name: Upload CPU coverage data
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-cpu
+          path: |
+            coverage.xml
+            pytest-coverage.txt
+
+  test-cuda:
+    name: Linux CUDA Test
+    needs: build
+    runs-on: gpu-t4-4-core
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Python check
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install CUDA libraries
+        run: |
+          sudo apt-get install -y libcusolver-dev-12-8 libcusolver-12-8 libcusparse-dev-12-8 libcusparse-12-8 libnvjitlink-12-8 libcublas-12-8
+          echo "/usr/local/cuda/targets/x86_64-linux/lib" | sudo tee /etc/ld.so.conf.d/cuda-targets.conf
+          sudo ldconfig
+      - name: Download wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: linux-wheel
+      - name: Install quadrants
+        run: |
+          set -x
+          mkdir -p dist
+          mv *.whl dist/
+          pip install dist/*.whl
+      - name: Install test requirements
+        run: |
+          pip install --group test
+          pip install -r requirements_test_xdist.txt
+      - name: Run CUDA tests with coverage
+        run: |
+          bash .github/workflows/scripts_new/linux/4_test_cuda.sh
+      - name: Upload CUDA coverage data
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-cuda
+          path: coverage.xml
+
+  coverage-comment:
+    if: github.event_name == 'pull_request' && always()
+    needs: [build, test-cuda]
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Python check
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Download CPU coverage
+        uses: actions/download-artifact@v4
+        continue-on-error: true
+        with:
+          name: coverage-cpu
+          path: coverage-cpu
+      - name: Download CUDA coverage
+        uses: actions/download-artifact@v4
+        continue-on-error: true
+        with:
+          name: coverage-cuda
+          path: coverage-cuda
+      - name: Generate coverage report
+        continue-on-error: true
+        run: |
+          COV_XMLS=""
+          if [ -f coverage-cpu/coverage.xml ]; then
+            COV_XMLS="coverage-cpu/coverage.xml"
+          fi
+          if [ -f coverage-cuda/coverage.xml ]; then
+            COV_XMLS="$COV_XMLS coverage-cuda/coverage.xml"
+          fi
+          if [ -z "$COV_XMLS" ]; then
+            echo "No coverage XML files found, skipping report"
+            exit 0
+          fi
+
+          python tests/coverage_report.py --report-only \
+            --compare-branch=origin/${{ github.base_ref }} \
+            --coverage-xml $COV_XMLS \
+            --format markdown > coverage-comment.md
+      - name: Post coverage comment
+        if: always() && hashFiles('coverage-comment.md') != ''
+        run: gh pr comment ${{ github.event.pull_request.number }} --body-file coverage-comment.md
+        env:
+          GH_TOKEN: ${{ github.token }}

.github/workflows/scripts_new/linux/4_test.sh

@@ -7,9 +7,18 @@ pip install -r requirements_test_xdist.txt
 export QD_LIB_DIR="$(python -c 'import quadrants as ti; print(ti.__path__[0])' | tail -n 1)/_lib/runtime"
 ./build/quadrants_cpp_tests  --gtest_filter=-AMDGPU.*
 
+TEST_EXIT=0
+
 # Phase 1: run all tests except torch-dependent ones
-python tests/run_tests.py -v -r 1 -m "not needs_torch"
+python tests/run_tests.py -v -r 3 --coverage -m "not needs_torch" || TEST_EXIT=$?
 
-# Phase 2: install torch, run only torch tests
 pip install torch --index-url https://download.pytorch.org/whl/cpu
-python tests/run_tests.py -v -r 1 -m needs_torch
+QD_KERNEL_COVERAGE=0 python tests/run_tests.py -v -r 3 --coverage --cov-append -m needs_torch || TEST_EXIT=$?
+
+# Phase 3: run tests that are skipped under kernel coverage (offline cache, snode layout, FE-LL observations,
+# etc.) without --coverage so QD_KERNEL_COVERAGE stays 0.
+QD_KERNEL_COVERAGE=0 python tests/run_tests.py -v -r 3 -m "not needs_torch" || TEST_EXIT=$?
+
+python tests/coverage_report.py --collect-only
+
+exit $TEST_EXIT

.github/workflows/scripts_new/linux/4_test_cuda.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+set -ex
+
+TEST_EXIT=0
+
+# Disable kernel-level coverage on CUDA: it changes field memory layout and breaks dlpack tests
+# (ValueError: Expected zero byte_offset).  Python code coverage (--cov) still runs.
+QD_KERNEL_COVERAGE=0 python tests/run_tests.py -v -r 1 --arch cuda --coverage -m "not needs_torch" || TEST_EXIT=$?
+
+pip install torch --index-url https://download.pytorch.org/whl/cu128
+QD_KERNEL_COVERAGE=0 python tests/run_tests.py -v -r 1 --arch cuda --coverage --cov-append -m needs_torch || TEST_EXIT=$?
+
+# Run kernel coverage tests on CUDA with coverage enabled — these are skipped by the phases above
+# (QD_KERNEL_COVERAGE=0) and include GPU-only tests like test_kernel_coverage_simt_e2e.
+QD_KERNEL_COVERAGE=1 python tests/run_tests.py -v -r 1 --arch cuda --coverage --cov-append test_kernel_coverage.py || TEST_EXIT=$?
+
+python tests/coverage_report.py --collect-only
+
+exit $TEST_EXIT

.gitignore

@@ -65,8 +65,13 @@ __pycache__
 /python/test_env
 /CHANGELOG.md
 /.coverage
+/.coverage.*
+_qd_kcov.*
 /coverage.xml
+/coverage-report.html
 /htmlcov
+/diff-cover.*
+/pytest-coverage.txt
 libpython_path.txt
 .vscode
 _build

docs/source/user_guide/index.md

@@ -57,6 +57,14 @@ graph
 perf_dispatch
 ```
 
+```{toctree}
+:caption: Testing
+:maxdepth: 1
+:titlesonly:
+
+kernel_coverage
+```
+
 ```{toctree}
 :caption: Reference
 :maxdepth: 1

docs/source/user_guide/kernel_coverage.md

@@ -0,0 +1,103 @@
+# Kernel code coverage
+
+Standard Python coverage tools only measure host-side code. Quadrants kernel coverage goes further — it tracks which lines actually execute *inside* compiled kernels on the device (CPU or GPU), including which branches of `if`/`else` blocks are taken at runtime.
+
+The coverage data is written in the standard `coverage.py` format, so it works with `coverage report`, `pytest-cov`, `diff-cover`, and IDE coverage viewers out of the box.
+
+## Prerequisites
+
+Kernel coverage requires the `coverage` Python package:
+
+```bash
+pip install coverage
+```
+
+## Enabling kernel coverage
+
+### Automatic with pytest-cov
+
+If you use `pytest-cov`, kernel coverage is enabled automatically — no configuration needed. Quadrants ships a pytest plugin that detects `--cov` and sets `QD_KERNEL_COVERAGE=1` for you. Just run:
+
+```bash
+pytest --cov=my_package --cov-branch tests/
+```
+
+To disable kernel coverage while still collecting Python coverage, opt out explicitly:
+
+```bash
+QD_KERNEL_COVERAGE=0 pytest --cov=my_package --cov-branch tests/
+```
+
+### Manual with any script
+
+For scripts outside pytest, set the `QD_KERNEL_COVERAGE` environment variable:
+
+```bash
+QD_KERNEL_COVERAGE=1 python my_simulation.py
+```
+
+This works with any script that uses quadrants kernels — no changes to your code are needed.
+
+When the process exits, quadrants writes one or more `_qd_kcov.<pid>` files in the working directory containing the collected coverage data.
+
+## Viewing results
+
+### With coverage.py
+
+Combine the kernel coverage files and produce a report using the standard `coverage` tool:
+
+```bash
+# Combine all kernel coverage files into .coverage
+coverage combine _qd_kcov.*
+
+# Terminal summary
+coverage report --show-missing
+
+# HTML report
+coverage html
+```
+
+### With pytest-cov
+
+When using `pytest-cov`, kernel coverage is enabled automatically (see above). The kernel coverage data is merged with Python coverage after the run:
+
+```bash
+coverage combine _qd_kcov.* .coverage
+```
+
+## Key properties
+
+- **Zero overhead when disabled.** The coverage module is never imported unless `QD_KERNEL_COVERAGE=1` is set. There is no cost in normal operation.
+- **Branch coverage.** Probes inside `if`/`else` bodies only fire when that branch is taken, giving true runtime branch coverage — not just kernel-level coverage, or static conditional coverage.
+- **Works with pytest-xdist.** Each worker writes to a separate file; combine them afterward.
+- **Survives `qd.init()` resets.** Coverage data is accumulated across multiple `qd.init()` calls within the same process.
+
+## Advanced usage
+
+### Probe capacity
+
+There is a limit of 100,000 coverage probes per process (one probe per unique source line per kernel/func). If you hit the limit — for example in a very large codebase with many kernels — increase it via the environment variable:
+
+```bash
+QD_COVERAGE_MAX_PROBES=500000 QD_KERNEL_COVERAGE=1 python my_simulation.py
+```
+
+## Coverage and autodiff
+
+The forward pass is covered. The backward pass is not, because instrumenting it would interfere with gradient computation. This is normally fine — the backward pass is auto-generated and replays the same control flow, so forward coverage is sufficient.
+
+One edge case: kernel calls inside a `qd.ad.Tape` with `validation=True` will not be covered.
+
+## Offline cache interaction
+
+Coverage probes change the compiled kernel, so the offline cache will see them as new kernels and recompile. This is expected and does not affect correctness, but the first run with coverage enabled will be slower if you normally rely on cached kernels.
+
+## CI integration
+
+The CI workflow posts a diff coverage report as a PR comment on each push. A **new comment** is created each time (rather than editing the previous one) so that the PR timeline shows a clear chronological sequence of commits and their corresponding coverage results.
+
+## Under the hood
+
+When `QD_KERNEL_COVERAGE=1` is set, quadrants rewrites the Python AST of each `@qd.kernel` and `@qd.func` before compilation. It inserts lightweight probe statements (`field[probe_id] = 1`) at each source line. These probes compile as ordinary field stores and execute on the device alongside your kernel code.
+
+At process exit, the probe data is read back from the device and written to a `.coverage`-compatible file.

pyproject.toml

@@ -99,6 +99,9 @@ test = [
     "pyright",
 ]
 
+[project.entry-points.pytest11]
+quadrants = "quadrants.pytest_plugin"
+
 [project.urls]
 Homepage = "https://github.com/Genesis-Embodied-AI/quadrants"
 
@@ -120,6 +123,14 @@ requires = [
 # things, without doing full c++ build
 build-backend = "setuptools.build_meta"
 
+[tool.coverage.paths]
+source = [
+    "python/quadrants",
+    "*/site-packages/quadrants",
+]
+
+[tool.coverage.report]
+
 [tool.pytest.ini_options]
 filterwarnings = [
     "ignore:Calling non-taichi function",

python/quadrants/lang/_fast_caching/src_hasher.py

@@ -1,4 +1,5 @@
 import json
+import os
 import warnings
 from typing import Any, Iterable, Sequence
 
@@ -49,6 +50,7 @@ def create_cache_key(
             kernel_source_info.filepath,
             str(kernel_source_info.start_lineno),
             "pruned",
+            "kcov" if os.environ.get("QD_KERNEL_COVERAGE") == "1" else "",
         )
     )
     return cache_key

python/quadrants/lang/_func_base.py

@@ -4,6 +4,7 @@
 import ast
 import inspect
 import math
+import os
 import sys
 import textwrap
 import types
@@ -21,6 +22,11 @@
 
 import numpy as np
 
+
+def _kernel_coverage_enabled() -> bool:
+    return os.environ.get("QD_KERNEL_COVERAGE") == "1"
+
+
 from quadrants._lib import core as _qd_core
 from quadrants._lib.core.quadrants_python import KernelLaunchContext
 from quadrants.lang import _kernel_impl_dataclass, impl
@@ -246,9 +252,21 @@ def get_tree_and_ctx(
 
         autodiff_mode = current_kernel.autodiff_mode
 
+        _kcov = None
+        if _kernel_coverage_enabled() and autodiff_mode == _qd_core.AutodiffMode.NONE:
+            from . import (  # pylint: disable=import-outside-toplevel
+                _kernel_coverage as _kcov,
+            )
+
+            tree = _kcov.rewrite_ast(tree, function_source_info.filepath, function_source_info.start_lineno)
+
         quadrants_callable = current_kernel.quadrants_callable
         is_pure = quadrants_callable is not None and quadrants_callable.is_pure
         global_vars = self._get_global_vars(self.func)
+        if _kcov is not None:
+            cov_field = _kcov.get_field()
+            if cov_field is not None:
+                global_vars[_kcov.FIELD_VAR_NAME] = cov_field
 
         template_vars = {}
         if is_kernel or is_real_function: