Fix APIC array reduction replay

[shi-eric]

Description

Closes #1663.

Record wp.utils.array_sum() and wp.utils.array_inner() in APIC graphs so CPU graph replay and CPU/CUDA save/load recompute reductions from current inputs. Previously, the native reduction helpers did not append replayable operation metadata, so captured graphs could omit the reduction or retain capture-time results.

Changes

• Record array reduction metadata and bindings in APIC streams, then reconstruct the operations for CPU replay and CUDA graph loading.
• Support whole-array reductions, explicit counts, composite data types, and positive-stride axis reductions.
• Reject unsupported counts, negative strides, unaligned layouts, and non-empty captures without an explicit output array.
• Bump the APIC writer format to version 15 while retaining reader compatibility with versions 13 through 15.
• Document the supported reduction contract and add CPU/CUDA capture, replay, save/load, metadata, empty-input, and validation coverage.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• CHANGELOG.md is updated for any user-facing changes under the Unreleased section.

Validation summary

Built the release native libraries with CUDA 13.0 on driver 13.0, verifying that the C++ and CUDA additions compile and link.

Ran WARP_CACHE_PATH=/tmp/warp-cache-warp-worktree-4 uv run warp/tests/test_apic.py; all 216 tests passed on CPU and two CUDA devices, including the new array reduction capture, replay, save/load, metadata, and error-handling cases.

Ran pre-commit on all changed files; every applicable hook passed.

Built the HTML documentation successfully with no build errors.

Bug fix

import numpy as np

import warp as wp


with wp.ScopedDevice("cpu"):
    a = wp.array([1.0, 2.0], dtype=wp.float32)
    eager_out = wp.zeros(1, dtype=wp.float32)
    captured_out = wp.zeros(1, dtype=wp.float32)

    wp.utils.array_inner(a, a, out=eager_out)

    with wp.ScopedCapture() as capture:
        wp.utils.array_inner(a, a, out=captured_out)

    captured_out.zero_()
    wp.capture_launch(capture.graph)

    np.testing.assert_allclose(captured_out.numpy(), eager_out.numpy())

Summary by CodeRabbit

• New Features
  • Experimental APIC reduction capture/replay for wp.utils.array_sum() and wp.utils.array_inner() (CPU and CUDA), with correct save/load persistence and parameter updates.
  • APIC graph format update for reduction ops (writer emits a newer version; reader remains backward-compatible).
• Bug Fixes
  • Captured reductions no longer recompute on replay; improved handling for empty/zero-count cases and stricter validation of counts, strides, and required out.
  • CUDA capture behavior clarified so reduction work is deferred appropriately.
• Documentation
  • Updated runtime guidance for APIC CPU/GPU reduction limitations and save/load version compatibility.
• Tests
  • Added new correctness, error-handling, empty-input, and save/load metadata coverage for these reductions.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: b2b5f2f1-82c7-4778-b43a-efa15046be4e

📥 Commits

Reviewing files that changed from the base of the PR and between 4a8d056 and 3e16144.

📒 Files selected for processing (11)

• CHANGELOG.md
• docs/user_guide/runtime.rst
• warp/_src/apic/types.py
• warp/_src/utils.py
• warp/native/apic.cpp
• warp/native/apic.cu
• warp/native/apic_internal.h
• warp/native/apic_types.h
• warp/native/reduce.cpp
• warp/native/reduce.cu
• warp/tests/test_apic.py

🚧 Files skipped from review as they are similar to previous changes (9)

• CHANGELOG.md
• warp/native/reduce.cpp
• warp/_src/apic/types.py
• warp/native/apic_internal.h
• warp/native/apic.cpp
• docs/user_guide/runtime.rst
• warp/native/apic.cu
• warp/_src/utils.py
• warp/tests/test_apic.py

---

📝 Walkthrough

Walkthrough

Changes

Array reductions now participate in CPU graph and CPU/CUDA APIC capture, serialization, validation, and replay. The change adds reduction records, capture constraints, native replay dispatch, save/load coverage, and runtime documentation.

APIC reduction support

Layer / File(s)	Summary
Reduction record contract warp/native/apic_types.h, warp/native/apic_internal.h, warp/native/apic.cpp, warp/_src/apic/types.py	Defines reduction opcodes, serialized records, reduction kinds, metadata validation, byte-span checks, and operation recording.
Python reduction capture rules warp/_src/utils.py	Adds APIC validation for array_sum() and array_inner(), including explicit outputs, count and stride limits, alignment checks, array tracking, and empty-reduction handling.
Native capture and replay warp/native/reduce.cpp, warp/native/reduce.cu, warp/native/apic.cpp, warp/native/apic.cu	Records host and device reductions and dispatches serialized sum or inner operations during CPU and CUDA replay.
Reduction validation and compatibility coverage warp/tests/test_apic.py, docs/user_guide/runtime.rst, CHANGELOG.md	Adds capture, error, empty-input, CUDA, save/load, and metadata tests, with matching runtime documentation and changelog updates.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant ReductionAPI as array_sum/array_inner
  participant APICCapture
  participant APICReplay
  Caller->>ReductionAPI: invoke reduction
  ReductionAPI->>APICCapture: validate and record metadata
  APICCapture->>APICReplay: serialize reduction operation
  APICReplay->>ReductionAPI: dispatch reduction during replay
  APICReplay-->>Caller: write output

Loading

Suggested reviewers: nvtw, adenzler-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 45.16% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and directly describes the main fix: APIC array reduction replay.
Linked Issues check	✅ Passed	The PR records array_inner reductions for APIC replay and save/load, resolving the stale replay behavior described in #1663.
Out of Scope Changes check	✅ Passed	The added docs, tests, format bump, and array_sum support all align with the stated APIC reduction replay fix.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[greptile-apps]

Greptile Summary

This PR fixes a bug where wp.utils.array_sum() and wp.utils.array_inner() called inside an APIC-captured graph were not replayed on graph re-execution; instead, the captured graph either skipped or froze the capture-time result. The fix teaches the APIC byte-stream serialization to record a new APIC_OP_REDUCTION record for each native call, and teaches both the CPU replay loop and the CUDA graph-rebuild loop to reissue the reduction with current array pointers.

• Adds APICReductionRecord (64 bytes, static-asserted) and APIC_OP_REDUCTION = 16 to the type system; bumps the APIC writer format to v15 while remaining backward-compatible with v13–14 readers.
• Injects apic_capture_reduction_host (early-return guard, records + skips execution) and apic_capture_reduction_device (records then falls-through so CUDA kernels reach the capture stream) at the top of all four native reduction entry points.
• Adds Python-side validation in array_sum and array_inner to reject captures with missing out, negative counts, strides that overflow signed int32, or addresses/strides misaligned to the scalar type, and adds 8 new APIC test functions covering capture, errors, empty, save/load, and metadata cases.

Confidence Score: 5/5

This PR is safe to merge. The changes are well-scoped to the APIC reduction serialization path, all validation is layered (Python pre-flight, native metadata check, stream validator, replay pointer resolution), and the 8 new test functions cover capture correctness, save/load round-trips, empty-input edge cases, and all documented error paths.

The implementation follows the same patterns as the existing scan/sort/BSR extended ops. The CPU path correctly skips execution during capture and re-executes on replay; the CUDA path correctly falls through to the CUB dispatch so kernels reach the capture stream as native graph nodes. Overflow and alignment checks are applied at every layer. All 216 tests passed on CPU and two CUDA devices per the author's validation summary.

No files require special attention.

Important Files Changed

Filename	Overview
warp/native/apic_types.h	Adds APICReductionRecord (64 bytes, static_asserted), APICReductionKind enum, APIC_OP_REDUCTION=16, and bumps APIC_FORMAT_VERSION from 14 to 15 while keeping APIC_MIN_SUPPORTED_FORMAT_VERSION at 13.
warp/native/apic_internal.h	Adds apic_reduction_input_bytes() (safe overflow-checked span calculation) and apic_reduction_metadata_valid() (validates kind, dtype, count, strides, alignment, and per-kind field constraints) as inline helpers shared by both the CPU and CUDA paths.
warp/native/reduce.cpp	Injects apic_capture_reduction_host() at the top of all four host entry points; records the APIC record and returns true (caller skips actual compute). On invalid metadata a poisoned record (region_id=-1) is written to ensure validation fails closed.
warp/native/reduce.cu	Injects apic_capture_reduction_device() before the CUB dispatch in all four device entry points; unlike the host path, does not return early so CUDA kernels are still submitted to the capture stream and become native graph nodes.
warp/native/apic.cpp	Adds apic_record_reduction(), apic_add_check() for uint64 overflow detection, APIC_OP_REDUCTION case in the stream validator (region/alignment/span checks), and APIC_OP_REDUCTION case in apic_cpu_replay_stream that resolves region pointers and re-invokes the appropriate host function.
warp/native/apic.cu	Adds APIC_OP_REDUCTION case in apic_replay_ops_into_cuda_capture that resolves device region pointers and calls the device reduction entry points while CUDA stream capture is active, so the CUB kernels become native graph nodes.
warp/_src/utils.py	Adds APIC-capture validation and tracking to array_sum() and array_inner(): validates negative/oversized counts, missing out, negative strides, int32 stride overflow, and scalar alignment, with an early-return fast path for zero-count whole-array captures in array_sum.
warp/_src/apic/types.py	Adds APIC_OP_REDUCTION = 16 to keep the Python type constants in sync with the C++ enum.
warp/tests/test_apic.py	Adds 8 new test functions covering capture+replay correctness, all capture-time error paths, empty-array edge cases, CPU-ops-not-recorded-during-CUDA-capture, and save/load with updated inputs including strided, composite-dtype, and axis-reduction metadata.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant PY as Python (utils.py)
    participant NAT as Native (reduce.cpp/.cu)
    participant APIC as APIC State (apic.cpp/.cu)
    participant RT as Replay/Rebuild

    Note over PY,RT: APIC Capture Phase (CPU or CUDA)
    PY->>PY: _get_apic_capture_for_device()
    PY->>PY: "validate (count, strides, out != None)"
    PY->>APIC: track_array(input / out)
    PY->>NAT: wp_array_sum_float_host/device()

    alt CPU host path
        NAT->>APIC: apic_capture_reduction_host()
        APIC->>APIC: apic_reduction_metadata_valid()
        APIC->>APIC: "apic_resolve_host_ptr() -> region_id + offset"
        APIC->>APIC: apic_record_reduction()
        NAT-->>PY: return (skip actual compute)
    else CUDA device path
        NAT->>APIC: apic_capture_reduction_device()
        APIC->>APIC: "apic_resolve_live_ptr() -> region_id + offset"
        APIC->>APIC: apic_record_reduction()
        NAT->>NAT: "array_sum_device_dispatch() (CUB -> captured by stream)"
    end

    Note over PY,RT: Replay / Graph Rebuild Phase
    RT->>RT: apic_cpu_replay_stream() or apic_replay_ops_into_cuda_capture()
    RT->>APIC: read APICReductionRecord
    RT->>RT: "resolve region_id+offset -> ptr"
    RT->>NAT: wp_array_sum_float_host/device(current ptrs)
    NAT->>NAT: Execute reduction on current data

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant PY as Python (utils.py)
    participant NAT as Native (reduce.cpp/.cu)
    participant APIC as APIC State (apic.cpp/.cu)
    participant RT as Replay/Rebuild

    Note over PY,RT: APIC Capture Phase (CPU or CUDA)
    PY->>PY: _get_apic_capture_for_device()
    PY->>PY: "validate (count, strides, out != None)"
    PY->>APIC: track_array(input / out)
    PY->>NAT: wp_array_sum_float_host/device()

    alt CPU host path
        NAT->>APIC: apic_capture_reduction_host()
        APIC->>APIC: apic_reduction_metadata_valid()
        APIC->>APIC: "apic_resolve_host_ptr() -> region_id + offset"
        APIC->>APIC: apic_record_reduction()
        NAT-->>PY: return (skip actual compute)
    else CUDA device path
        NAT->>APIC: apic_capture_reduction_device()
        APIC->>APIC: "apic_resolve_live_ptr() -> region_id + offset"
        APIC->>APIC: apic_record_reduction()
        NAT->>NAT: "array_sum_device_dispatch() (CUB -> captured by stream)"
    end

    Note over PY,RT: Replay / Graph Rebuild Phase
    RT->>RT: apic_cpu_replay_stream() or apic_replay_ops_into_cuda_capture()
    RT->>APIC: read APICReductionRecord
    RT->>RT: "resolve region_id+offset -> ptr"
    RT->>NAT: wp_array_sum_float_host/device(current ptrs)
    NAT->>NAT: Execute reduction on current data

Loading

_{Reviews (4): Last reviewed commit: "Record APIC array reductions (GH-1663)" | Re-trigger Greptile}