Skip to content

fix: preserve mutations through tensor views - #33

Open
frlai wants to merge 8 commits into
developfrom
feature/view-mutation-functionalization
Open

frlai wants to merge 8 commits into
developfrom
feature/view-mutation-functionalization

Conversation

@frlai

@frlai frlai commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

Mutating a separately held slice or reshaped view previously updated eager storage but could disappear from the exported graph, leaving base and sibling outputs incorrect. Preserve these relationships for Torch and NumPy, including supported Torch-view/Warp boundaries.

Record ordered updates on stable tensor handles, then use PyTorch export functionalization before pruning output dependencies. Import only user outputs from the export signature so artifacts return updated values without copying mutations into caller inputs. Preserve independent copies across NumPy dtype conversion and Torch/NumPy/Warp conversions, distinguish new Warp allocations from writes, and recompute runner output masks after graph normalization.

Keep FX/TorchScript support for tensor-valued slice bounds when an exporter cannot functionalize them: retain native mutations on per-call copies rather than freezing runtime indices. PT2/ONNX retain their exporter restrictions. Torch-backed partial Warp/APIC views now bind the full traced allocation in both directions; arbitrary Warp-native views remain outside this support.

For Warp versions exposing allocation-wide Torch storage bounds, record weak allocation identity and layout on existing FX nodes, resolve partial views through their graph ancestors, and deduplicate segment input and output bindings. Keep original kernel offsets/strides in APIC and use the existing mutation functionalization to propagate writes. This handles offset, nested, strided and compound views, including sibling mutations and real Warp writes, without per-view storage-root attributes or retained tensor samples. Reject missing covering graph values and dtype reinterpretation. Keep two representative functional tests in the existing shared-array test file: NumPy view mutation with state feedback, and Torch/Warp sliced-allocation replay through ONNX/PT2. The broader view/copy matrix remains in unit tests.

Keep explicit NumPy casts before the same-dtype copy when integrating the ATen adapter: imported graph intermediates can have different dtype metadata from native NumPy samples. Preserve both the requested dtype and independent storage.

Raise the minimum PyTorch version from 2.6.0 to 2.7.0, with matching dependency metadata, lockfile, README, and installation documentation. The broad version matrix exposed two existing boolean-mask ONNX export regressions only on 2.6; all tested 2.7+ releases pass without a compatibility workaround. Add Torch/NumPy coverage for differing copy layouts and ONNX replay with boolean-mask cardinalities 0, 1, 2, and 4.

Current test-consolidation validation:

  • Torch 2.14 / released Warp 1.18: 165 tests and 7 subtests passed across shared-memory, state, Warp tracing/boundaries, and the unchanged view-mutation unit tests.
  • The self-contained NumPy feedback test passes on minimum Torch 2.7.0 CPU.
  • Ruff formatting and diff checks pass. No production code or unit tests changed; the standalone functional runtime file is removed.

Validation before the functional-test consolidation, on rebased develop 1a749d30409d3196f95558063f200f373fa9aab6 (includes the NumPy ATen adapter):

  • Full Torch 2.14 / Warp 1.16 suite: 1024 tests and 1792 subtests passed; 1 expected allocation-contract test skipped.
  • Focused NumPy adapter/view/conversion checks: 237 tests and 1628 subtests passed.
  • Minimum-supported Torch 2.7.0 CPU checks: 71 passed; 6 explicit CUDA cases deselected.
  • Released Warp 1.18 allocation and boundary validation: 86 tests and 7 subtests passed, including the single new ONNX/PT2 allocation regression.
  • Ruff 0.17.1, strict Sphinx HTML/spelling, documentation linkcheck, and offline lockfile validation pass.

Earlier verification before the NumPy-adapter rebase (retained as historical results, not rerun matrices):

  • CPU compatibility matrix: 424 tests and 136 subtests passed per release on 2.7.0, 2.7.1, 2.8.0, 2.9.0, 2.9.1, 2.10.0, 2.11.0, 2.12.0, 2.12.1, 2.13.0, 2.14.0, and 2.14.1. Four CUDA-only cases skipped per CPU environment. This covers view/mutation, indexing, conversions, state, inference mode, export formats, and Warp-boundary discovery.
  • Complete suite on Torch 2.14.0+cu130 / RTX 4090: 951 passed, 1 skipped, 164 subtests passed on Warp 1.16, including examples, visualization, actual CUDA/Warp capture, and PT2/ONNX artifact reload. CUDA was not tested across the entire version matrix.
  • TensorRT 11.4.0.106 / RTX 4090: all 32 valid ONNX files converted to saved/reloaded engines and matched eager plus ONNX Runtime exactly over 152 fresh-input trials. Includes Torch/NumPy view/copy cases, two public annotation exports, and boolean masks with cardinality 0/1/2/4 using dynamic output allocation. No LEAPP-specific TensorRT plugin or ONNX rewriting. One of 33 attempted examples (diagonal mutation) fails before ONNX generation; the same failure occurs in ordinary PyTorch export without LEAPP. TensorRT conversion/build/runtime had no failures for generated files.
  • Allocation binding: the one new regression passes on released Warp 1.18.0 and source-built Warp 1.19.0.dev1 with matching native adapters; it skips on older Warp without allocation bounds. Both stacks pass 26 ONNX/PT2 export/reload cases and 104 independently varied replay trials each, including real Warp slice writes, compound views, temporary/copy roots, and unchanged caller inputs. Weak storage metadata checks pass on all 12 supported tested Torch releases.
  • Fresh root values and mutation operands are varied independently; checks include untouched elements and unchanged caller inputs.
  • The previous Isaac Lab documentation URL failure is resolved by the updated develop branch; current linkcheck passes.

GPU validation used isolated runtime adapters matched to each Warp release and the installed Torch 2.14 stack (C++20 locally for current Torch headers). Runtime build sources are unchanged by this PR. Negative NumPy slicing remains excluded because it is handled by a separate PR.

frlai added 4 commits October 21, 2026 11:50
Signed-off-by: Frank Lai <frlai@nvidia.com>
Signed-off-by: Frank Lai <frlai@nvidia.com>
Signed-off-by: Frank Lai <frlai@nvidia.com>
Signed-off-by: Frank Lai <frlai@nvidia.com>
@frlai
frlai force-pushed the feature/view-mutation-functionalization branch 2 times, most recently from d704152 to f6a0581 Compare October 21, 2026 23:37
Comment thread tests/functional_tests/test_shared_memory_arrays.py
Comment on lines +92 to +94
if old.node is not new.node:
mutation = old.tracer.create_proxy(
"call_method", "copy_", (old, new), {}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The existing full-replacement path in _handle_plain_assignment() assigns dest_proxy = value.proxy, and _record_assignment() then skips the copy. Consequently, buf.copy_(x) gives two independent tensors the same FX proxy. With this new setter, a subsequent buf += 5 records a write into that shared proxy, so returning (x, buf) produces (x + 5, x + 5) instead of (x, x + 5). Previously, updating the destination's proxy only rebound its reference.

Comment thread leapp/leapp_graph/traced_node.py
frlai added 2 commits October 21, 2026 17:10
Signed-off-by: Frank Lai <frlai@nvidia.com>
Signed-off-by: Frank Lai <frlai@nvidia.com>
Comment thread docs/source/tensor_libraries/torch_how_it_works.rst Outdated
Comment thread docs/source/tensor_libraries/warp_how_it_works.rst Outdated
Comment thread docs/source/tensor_libraries/warp_limitations.rst Outdated
Comment thread leapp/leapp_graph/datatypes/warp/patching.py Outdated
Signed-off-by: Frank Lai <frlai@nvidia.com>
* fix: normalize slice bounds after view functionalization

Signed-off-by: Frank Lai <frlai@nvidia.com>

* refactor: separate ATen rewrites from traced node

Signed-off-by: Frank Lai <frlai@nvidia.com>

* refactor: group functionalization with ATen rewrites

Signed-off-by: Frank Lai <frlai@nvidia.com>

* refactor: separate ATen matching from graph edits

Signed-off-by: Frank Lai <frlai@nvidia.com>

---------

Signed-off-by: Frank Lai <frlai@nvidia.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants