feat: add unit test coverage

.github/workflows/ci.yaml

@@ -178,6 +178,14 @@ jobs:
           pip3 install -r requirements.txt
           pip3 install "amd-aiter @ git+https://github.com/ROCm/aiter.git@${AITER_VERSION}"
           GPU_ARCHS="gfx942;gfx950" PRIMUS_TURBO_FRAMEWORK="PYTORCH" pip3 install --no-build-isolation -e ".[pytorch]" -v
+      - name: Prepare coverage environment
+        run: |
+          # Enable coverage in test subprocesses (spawned by distributed tests)
+          # and clear stale data files left over from retried/previous runs.
+          echo "COVERAGE_PROCESS_START=$(pwd)/pyproject.toml" >> $GITHUB_ENV
+          SITE_PACKAGES=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
+          echo "import coverage; coverage.process_startup()" > "${SITE_PACKAGES}/coverage_subprocess.pth"
+          rm -f .coverage .coverage.*
       - name: Run single-GPU tests (parallel)
         uses: nick-fields/retry@v3
         with:
@@ -187,7 +195,7 @@ jobs:
           command: >
             timeout -k 60s 130m
             bash -lc
-            "exec pytest -v tests/pytorch -n 8 --timeout=600 --timeout-method=thread --max-worker-restart=8 --log-cli-level=INFO"
+            "exec pytest -v tests/pytorch -n 8 --timeout=600 --timeout-method=thread --max-worker-restart=8 --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append"
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
       - name: Run deterministic tests (parallel)
@@ -199,17 +207,31 @@ jobs:
           command: >
             timeout -k 60s 55m
             bash -lc
-            "exec pytest -v tests/pytorch -n 8 --deterministic-only --timeout=600 --timeout-method=thread --max-worker-restart=8 --log-cli-level=INFO"
+            "exec pytest -v tests/pytorch -n 8 --deterministic-only --timeout=600 --timeout-method=thread --max-worker-restart=8 --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append"
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
       - name: Run distributed tests
         uses: nick-fields/retry@v3
         with:
           timeout_minutes: 60
           max_attempts: 2
-          command: pytest -v tests/pytorch --dist-only --timeout=600 --timeout-method=thread --log-cli-level=INFO
+          command: pytest -v tests/pytorch --dist-only --timeout=600 --timeout-method=thread --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+      - name: Combine and stage coverage data
+        if: always()
+        run: |
+          # Merge any leftover parallel/subprocess data files into a single
+          # .coverage, then stage it for the combined coverage-report job.
+          coverage combine || true
+          cp .coverage coverage_data_pytorch || true
+      - name: Upload coverage data
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-pytorch
+          path: coverage_data_pytorch
+          if-no-files-found: warn
 
   unittest-jax-gfx942:
     needs: code-lint
@@ -262,12 +284,20 @@ jobs:
           unset AITER_ASM_DIR
           pip3 install -r requirements.txt
           GPU_ARCHS="gfx942;gfx950" PRIMUS_TURBO_FRAMEWORK="JAX" pip3 install --no-build-isolation -e ".[jax]" -v
+      - name: Prepare coverage environment
+        run: |
+          # Enable coverage in test subprocesses (spawned by distributed tests)
+          # and clear stale data files left over from retried/previous runs.
+          echo "COVERAGE_PROCESS_START=$(pwd)/pyproject.toml" >> $GITHUB_ENV
+          SITE_PACKAGES=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
+          echo "import coverage; coverage.process_startup()" > "${SITE_PACKAGES}/coverage_subprocess.pth"
+          rm -f .coverage .coverage.*
       - name: Run single-GPU tests (parallel)
         uses: nick-fields/retry@v3
         with:
           timeout_minutes: 60
           max_attempts: 2
-          command: pytest -v tests/jax -n 8 --log-cli-level=INFO
+          command: pytest -v tests/jax -n 8 --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
       # Keep in-process multi-GPU tests and distributed/per-process tests in
@@ -280,14 +310,86 @@ jobs:
         with:
           timeout_minutes: 30
           max_attempts: 2
-          command: pytest -v tests/jax --dist-only --ignore=tests/jax/lax/test_mp_dispatch_combine.py --log-cli-level=INFO
+          command: pytest -v tests/jax --dist-only --ignore=tests/jax/lax/test_mp_dispatch_combine.py --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
       - name: Run distributed tests
         uses: nick-fields/retry@v3
         with:
           timeout_minutes: 30
           max_attempts: 2
-          command: pytest -v tests/jax/lax/test_mp_dispatch_combine.py --log-cli-level=INFO
+          command: pytest -v tests/jax/lax/test_mp_dispatch_combine.py --log-cli-level=INFO --cov=primus_turbo --cov-branch --cov-append
         env:
           XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+      - name: Combine and stage coverage data
+        if: always()
+        run: |
+          # Merge any leftover parallel/subprocess data files into a single
+          # .coverage, then stage it for the combined coverage-report job.
+          coverage combine || true
+          cp .coverage coverage_data_jax || true
+      - name: Upload coverage data
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-jax
+          path: coverage_data_jax
+          if-no-files-found: warn
+
+  coverage-report:
+    needs: [unittest-pytorch-gfx942, unittest-jax-gfx942]
+    # Run whenever the GPU jobs were eligible to run, even if some failed, so we
+    # still merge whatever coverage data is available.
+    if: >-
+      always() &&
+      (github.event_name != 'pull_request' ||
+       contains(github.event.pull_request.labels.*.name, 'ci:gpu'))
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+    steps:
+      - run: echo "🎉 Begin Primus-Turbo Coverage Merge."
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install coverage
+        run: pip install "coverage[toml]"
+      - name: Download coverage artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: coverage-*
+          path: coverage-artifacts
+          merge-multiple: true
+      - name: Merge coverage data and generate reports
+        run: |
+          shopt -s nullglob
+          data_files=(coverage-artifacts/coverage_data_*)
+          if [ ${#data_files[@]} -eq 0 ]; then
+            echo "No coverage data found; skipping merge." >> "$GITHUB_STEP_SUMMARY"
+            exit 0
+          fi
+          # Combine the raw per-framework data files into one line-level dataset.
+          coverage combine "${data_files[@]}"
+          coverage report
+          coverage html -d visual-coverage
+          coverage json -o coverage.json
+          python tools/ci/coverage_summary.py coverage.json \
+            >> "$GITHUB_STEP_SUMMARY" 2>/dev/null || true
+      - name: Upload combined coverage report
+        id: upload_combined
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: visual-coverage
+          path: visual-coverage
+          if-no-files-found: warn
+      - name: Link coverage artifact in run summary
+        if: always()
+        run: |
+          {
+            echo ""
+            echo "Full HTML report: download the **visual-coverage** artifact ([direct link](${{ steps.upload_combined.outputs.artifact-url }})), unzip and open \`visual-coverage/index.html\`."
+          } >> "$GITHUB_STEP_SUMMARY"

.gitignore

@@ -15,3 +15,4 @@ agent/workspace/
 logs/
 primus_turbo/_build_info.py
 primus_turbo/_version.py
+.coverage*

pyproject.toml

@@ -38,3 +38,27 @@ ignore = [
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.coverage.run]
+branch = true
+parallel = true
+concurrency = ["multiprocessing", "thread"]
+source = ["primus_turbo"]
+omit = [
+    "primus_turbo/_version.py",
+    "primus_turbo/_build_info.py",
+    "*/tests/*",
+]
+
+[tool.coverage.report]
+show_missing = true
+omit = [
+    "primus_turbo/triton/*",
+    "primus_turbo/flydsl/*",
+]
+
+[tool.coverage.paths]
+source = ["primus_turbo/", "*/primus_turbo/"]

requirements.txt

@@ -7,6 +7,8 @@ expecttest
 pytest
 pytest-xdist
 pytest-timeout
+pytest-cov
+coverage[toml]
 setuptools==69.5.1
 setuptools_scm[toml]>=8
 wheel

tools/ci/coverage_summary.py

@@ -0,0 +1,143 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Render coverage.py JSON as a compact Markdown table for the CI run summary.
+
+Usage:
+    coverage_summary.py REPORT.json [REPORT.json ...]
+"""
+
+import json
+import sys
+from collections import defaultdict
+
+PKG = "primus_turbo/"
+# Kernel layers excluded from the headline number (JIT-compiled; Python line
+# coverage is not meaningful). Kept in sync with [tool.coverage.report] omit.
+OMIT_MODULES = {"triton", "flydsl"}
+# Top-level groups whose sub-packages are shown as indented detail rows; every
+# other group (common, ...) is a single bold row.
+DETAILED_GROUPS = ("pytorch", "jax")
+
+
+def classify(path: str):
+    """Return (group, detail) for a covered file, or None to skip it.
+
+    group is the top-level row key (e.g. pytorch); detail is the sub-row key for
+    DETAILED_GROUPS (e.g. pytorch/ops), else None.
+    """
+    seg = (path[path.find(PKG) :] if PKG in path else path).split("/")
+    if seg[-1] == "__init__.py":
+        return None
+    if len(seg) < 2:
+        return "(top-level)", None
+    if len(seg) == 2:  # primus_turbo/<file>.py
+        return seg[1], None
+    if seg[1] in DETAILED_GROUPS:
+        return seg[1], seg[1] + "/" + seg[2]
+    return seg[1], None
+
+
+def _pct(covered: int, total: int) -> float:
+    return (100.0 * covered / total) if total else 0.0
+
+
+def _merge_files(reports: list) -> dict:
+    """Union of files across reports; per file take the max covered lines.
+
+    Approximate fallback for the multi-report case only. num_statements is
+    stable across reports for the same source file, so the union denominator is
+    well-defined. Taking the max covered lines avoids double counting a module
+    both jobs import, but undercounts when two reports cover disjoint lines of
+    the same file. CI feeds a single already-combined report (line-level union),
+    so this path is not exercised there.
+    """
+    merged = {}
+    for rep in reports:
+        for fpath, info in rep.get("files", {}).items():
+            s = info["summary"]
+            cur = merged.get(fpath)
+            if cur is None:
+                merged[fpath] = [s["covered_lines"], s["num_statements"]]
+            else:
+                cur[0] = max(cur[0], s["covered_lines"])
+                cur[1] = max(cur[1], s["num_statements"])
+    return merged
+
+
+def _aggregate(merged: dict):
+    """Return {group: {detail|group: [covered, statements]}} for kept modules."""
+    agg = defaultdict(lambda: defaultdict(lambda: [0, 0]))
+    for fpath, (cov, stmts) in merged.items():
+        result = classify(fpath)
+        if result is None:
+            continue
+        group, detail = result
+        if group in OMIT_MODULES:
+            continue
+        a = agg[group][detail or group]
+        a[0] += cov
+        a[1] += stmts
+    return agg
+
+
+def render(reports: list) -> str:
+    agg = _aggregate(_merge_files(reports))
+
+    def row(label, cov, stmts, bold=False):
+        vals = [format(cov, ","), format(stmts, ","), "%.1f%%" % _pct(cov, stmts)]
+        w = "**" if bold else ""
+        return "| " + " | ".join("%s%s%s" % (w, x, w) for x in [label] + vals) + " |"
+
+    def group_totals(group):
+        cov = sum(v[0] for v in agg[group].values())
+        stmts = sum(v[1] for v in agg[group].values())
+        return cov, stmts
+
+    groups = [g for g in agg if group_totals(g)[1] > 0]
+    tc = sum(group_totals(g)[0] for g in groups)
+    tn = sum(group_totals(g)[1] for g in groups)
+    excl = ", ".join(sorted(OMIT_MODULES))
+
+    out = ["## Primus-Turbo coverage\n"]
+    out.append(
+        "**Total line coverage: %.1f%%** (%s / %s statements; excludes %s)\n"
+        % (_pct(tc, tn), format(tc, ","), format(tn, ","), excl)
+    )
+    out += ["| Module | Covered | Stmts | Coverage |", "|---|--:|--:|--:|"]
+
+    # Top-level groups, sorted by coverage (desc).
+    for group in sorted(groups, key=lambda g: -_pct(*group_totals(g))):
+        cov, stmts = group_totals(group)
+        out.append(row("`%s`" % group, cov, stmts, bold=True))
+        if group in DETAILED_GROUPS:
+            details = ((k, v) for k, v in agg[group].items() if v[1] > 0)
+            for k, v in sorted(details, key=lambda kv: -_pct(kv[1][0], kv[1][1])):
+                out.append(row("&emsp;`%s`" % k, v[0], v[1]))
+
+    out.append(row("TOTAL", tc, tn, bold=True))
+    return "\n".join(out)
+
+
+def _load(path: str) -> dict:
+    with open(path) as f:
+        return json.load(f)
+
+
+def main() -> int:
+    if len(sys.argv) < 2:
+        print(
+            "usage: coverage_summary.py REPORT.json [REPORT.json ...]",
+            file=sys.stderr,
+        )
+        return 2
+    reports = [_load(p) for p in sys.argv[1:]]
+    print(render(reports))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())