Skip to content

Add a design document for API Capture and CPU Graphs #1639

Description

@shi-eric

Summary

Add a design document under design/ for API Capture (APIC), serialized .wrp graphs, and CPU graph capture and replay.

APIC and CPU graph replay were introduced as experimental features in Warp 1.13 and expanded through 1.14 and 1.15. What began as kernel launch recording now includes a versioned operation stream, memory-region tracking, argument relocation, backward launches, conditional control flow, host helper operations, CPU replay, CUDA graph reconstruction, and standalone C++ loading.

This is enough machinery that the code is no longer a reasonable substitute for a design document.

Why we need this

The user guide does a good job of showing how to use the APIs and listing their current limitations. It does not explain the design choices or the contracts that new code must preserve.

Today, a contributor has to read context.py, warp/_src/apic/, several native APIC files, the memory and launch hooks, and a large test suite to answer questions such as:

  • What does a Graph own after CPU capture, CUDA capture, or capture_load()?
  • Which operations run during capture, and which are deferred until replay?
  • How do array aliases, temporary allocations, gradients, structs, and object handles remain valid when pointers change?
  • What work is required to make another Warp operation recordable?
  • When must a .wrp file be recaptured after a format change?

The original design work exists mostly in external gists linked from #1349 and #1431. Several Phase 2 implementation commits also say that work was tracked in design/graph-capture-newton-coverage.md, but that file is not in the repository. We should bring the relevant reasoning into the repo and update it to match the implementation on main.

Suggested scope

The document should start by defining the terms used throughout this feature: graph capture, APIC recording, the operation byte stream, native CUDA graphs, live graphs, loaded graphs, CPU replay, and CUDA graph reconstruction. A diagram should show how a capture is recorded and later consumed on CPU and CUDA, both with and without capture_save().

It should explain the capture state and execution model. In particular, CPU capture defers operations into the APIC stream, while CUDA capture records work in the native CUDA graph and may also record the APIC representation. The document should cover pause and resume, thread and device scoping, nested-capture restrictions, and error handling when capture ends unsuccessfully.

Memory ownership needs its own treatment. That includes region IDs, base allocations, array views and offsets, previously unknown pointers, capture-time allocations, temporary buffers retained by a graph, named input and output bindings, and the difference between snapshotted and size-only regions. It should also say what the graph retains automatically and what the caller must keep alive.

Kernel arguments have another set of non-obvious rules. The document should describe value blobs, ABI alignment, pointer relocations, gradient pointers and flags, indexed arrays, nested structs, forward and adjoint arguments, module and kernel lookup, and object-handle remapping. Readers should not have to reconstruct these rules from the packed C structs.

The operation model should explain why some host helpers need their own APIC records even though they are not ordinary kernel launches. It should describe the supported CPU and CUDA paths, how unsupported operations fail, and the rule that capture must not silently omit work. Conditional branches and loops should be covered here too, including nested operation streams, replay-time condition evaluation, loop limits, and the differences from CUDA conditional graph nodes.

The serialization section should cover the .wrp file, the companion _modules directory, operation-stream validation, device and architecture restrictions, PTX and CUBIN portability, CPU object loading, and standalone C++ replay. Most importantly, it should state the compatibility policy. Contributors need to know when a format change requires a version bump, whether older files remain readable, and when users must recapture a graph.

Finally, the document should preserve the reasoning behind the current approach. Why is recording done in C++ instead of Python? Why does CPU replay interpret the byte stream directly? Why do CPU and CUDA share a recording format even though their execution paths differ? It should also include a short checklist for adding an operation, covering its recording hook, wire record, validation, CPU replay, CUDA reconstruction where applicable, serialization, tests, documentation, and versioning.

Known gaps belong in the document, but they should be clearly separated from established behavior. That includes unsupported array layouts, stream events, texture copies, additional object types, multi-GPU graphs, and deferred compilation.

Acceptance criteria

  • Add a design document under design/ using the conventions in design/README.md.
  • Mark the feature as implemented but still experimental.
  • Describe the implementation on main and identify older decisions that have been superseded.
  • Include a diagram of the recording and replay paths.
  • Explain graph state, memory ownership, and lifetime rules.
  • Document the process for adding another recorded operation.
  • State the .wrp format and native C API compatibility policy.
  • Link to the relevant implementation files, tests, and tracking issues.
  • File any behavior or documentation inconsistencies found during the audit as separate follow-up issues.

References

Metadata

Metadata

Assignees

Labels

docsImprovements or additions to documentation

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions