[lgulich]

Summary

Two Warp tracing fixes found while exporting a Warp geometric fabric (nvblox_next GlobalSurfelFabric) through LEAPP. In both cases the export either failed or exported a model that silently gave wrong results after the first call.

1. wp.from_torch of a view of a traced tensor (traced[5:10], surfels[..., 0:3])

• Warp APIC (1.18+) records a Torch-backed array as its whole storage allocation. LEAPP fed the runner only the view, so validation failed with APIC param 'input_0' byte size mismatch: metadata/tensor has 20 bytes, APIC expects 40.
• Fix:
  • A traced view now remembers which traced tensor covered its storage when it was taken, along with both proxies at that moment (datatypes/torch/storage_root.py).
  • A segment input over part of a Torch storage is then replaced by that tensor, so all views of one buffer become one runner input.
• Fails closed:
  • If the view or the base was modified in place after the view was taken, or no traced tensor covers the storage, tracing stops with an explanation. This avoids exporting a stale buffer, since mutations through one alias do not update the other alias's proxy.
  • Older Warp versions record each view as its own region, so they are rejected explicitly.

2. zero_() / fill_() of an untraced Warp buffer inside a node

• A common pattern is resetting a persistent accumulator before kernels atomic_add into it, e.g. metric.zero_() and then a launch. Automatic segments only open at the first Warp call that touches traced data. The reset therefore ran once during tracing and never on replay, and the runner accumulated across calls.
  • LEAPP's single-example validation still reported PASS, so this was silent.
• Fix:
  • Everything between input_tensors and output_tensors belongs to that node. An overwrite of an untraced array there now opens or joins that node's segment.
  • The array joins the node, and its value is bound to the runner output, so later readers depend on the runner through a graph edge. That includes segments split off by Torch work in between.
• Scope:
  • Only fill_ and zero_ are attributed. A kernel launch on an untraced array may read its previous contents, which is undeclared state, so it stays a user error. This keeps test_undeclared_persistent_alias_attaches_to_nothing as is.
  • When several nodes are open, the overwrite cannot be attributed, so LEAPP warns instead of guessing.
  • Open-node tracking is reset by leapp.start(), so a trace aborted by an exception does not leak into the next session.

Tests

New functional tests in tests/functional_tests/test_warp_tracing.py. Inference uses inputs that differ from the trace inputs, and the reset tests run inference three times:

• TestWarpTorchViews.test_offset_view_reads_live_input
• TestWarpTorchViews.test_strided_views_of_one_storage_are_one_segment_input, which also asserts that there is one segment input
• TestWarpTorchViews.test_view_modified_in_place_fails
• TestWarpUntracedOverwrite.test_reset_in_same_segment_runs_every_call, which covers one segment
• TestWarpUntracedOverwrite.test_reset_before_torch_op_feeds_later_segment, which covers two segments joined by a graph edge

All 5 fail without the fix.

Full suite, run in a fresh venv (torch 2.12.1+cu130, onnxruntime-gpu, Warp main / 1.19.0.dev0, RTX 6000 Ada): pytest tests/ gives 874 passed.

With the released warp-lang==1.18.0 wheel on this machine, 37 tests fail on unmodified develop as well, with Failed to load module ... .ptx. The wheel's CUDA 13.4 toolchain emits PTX that this machine's 13.0 driver cannot load. My new inference tests hit the same error there, so they are only verified on Warp main.

Notes

• Rebased on develop after fix: trace Warp writes into plain Torch aliases #23. fix: trace Warp writes into plain Torch aliases #23 makes the Warp array's _tensor the traced tensor itself, so the storage checks unwrap to a plain tensor before querying the layout. Otherwise the query would be recorded as graph operations.
• Related, not in this PR: the ONNX export optimizer (onnxscript add_0/mul_by_1 no-op rules, which use abs_tol=1e-8) removes guards such as sqrt(x + 1e-12), producing inf/NaN in exported models. That needs an upstream onnxscript fix.

🤖 Generated with Claude Code

[frlai]

@lgulich, Issue 1 you defined is fixed as part of #33. it provides larger capabilities of tracing mutations through views of tensors.

issue 2 is a real problem however it might be possible to fix it through annotate.register_buffer()? this is the cannonical way to identify buffers like this. however register buffer needs to be declared after input tensors. Can you try to see if that is good enough? if not I can see what I can do. the solution you proposed might be hazardous especially if someone decides to open multiple model captures at the same time and might introduce overhead.

if issue 2 can't be fully solved with register buffers i can approve this PR and later update it with #33 to solve the first problem to unblock you.

[lgulich]

thanks, seems like register_buffer does solve the problem. So then we can close this MR and merge #33 instead