Fix PyTorch autograd hook dropped after no-grad call (include requires_grad in call-data cache signature)

[nv-slang-bot]

Motivation

The torch-tensor call-data cache signature did not include grad-ness. requires_grad was read into the tensor-info struct but never written to the cache key. The autograd hook is gated at dispatch on the torch_autograd flag frozen into the cached CallData at build time (slangpyfunction.cpp → is_torch_autograd()), which detect_torch_tensors sets once, on the first (cache-miss) build.

So a no-grad-first build cached torch_autograd=False; a later same-rank/dtype/shape grad call hit that entry, the TorchAutoGradHook was bypassed, and the output came back with grad_fn=None. That exactly explains the reporter's asymmetry — no-grad→grad broken, grad→grad fine, grad→no-grad fine.

a_nograd = torch.randn((8, 5), device="cuda", requires_grad=False)
module.square(a_nograd)                      # caches torch_autograd=False

a = torch.randn((8, 5), device="cuda", requires_grad=True)
b = module.square(a)                         # same rank/dtype/shape -> cache hit
assert b.grad_fn is not None                 # FAILED before this change

Fixes #1052

Proposed solution

Append a per-tensor grad bit to the signature, composing it onto the existing shape-compatibility run:

[Dn,Sm,G<0|1>,V...]     n = ndim, m = scalar_type, G = requires_grad, V = one shape-compat char per dim

This is a composition, not a replacement. #1082 (d0ee6b1, "Torch tensors supply vector dims") landed the V run on main after this branch was cut, and it rewrote the same emitter. Taking either side wholesale would regress the other: keeping only G drops #1082's vector-dims keying, keeping only V drops the #1052 grad fix. Both are cache-collision bugs, just on different axes — so the format carries both fields.

The bit is per-tensor rather than folded across arguments: a single call-level flag would collide grad (True, False) with (False, True), which is the same class of collision being fixed.

Implemented in both bridges — native tensor_bridge_get_signature (src/slangpy_torch/torch_bridge_impl.cpp) and the Python get_signature fallback (slangpy/torchintegration/bridge_fallback.py) — so the two paths cannot key the cache differently.

Change summary

9 files changed, +197 / −28. Most of that is test updates: the format change invalidates every hardcoded signature expectation, and there are nineteen of them across three files.

Area	Change
src/slangpy_torch/torch_bridge_impl.cpp	Emit ,G<bit> before the shape-compat run
src/slangpy_torch/tensor_bridge_api.h	TENSOR_BRIDGE_API_VERSION 8 → 9; format contract doc updated
slangpy/torchintegration/bridge_fallback.py	Same composed format in the Python fallback
src/slangpy_ext/utils/torch_bridge.cpp	Doc comment for the composed return format
slangpy/tests/slangpy_tests/test_pytorch.py	Regression test for the no-grad→grad sequence
slangpy/tests/utils/test_torch_bridge.py	Signature expectations; grad-discrimination and shape-preservation tests
slangpy/tests/slangpy_tests/test_torchintegration.py	Updated expected signatures
slangpy/tests/slangpy_tests/test_bridge_fallback_gaps.py	Updated expected signatures (outside the conflict set — see below)

Concepts and vocabulary

• Call-data cache key — the signature string is an opaque key. Nothing parses it back into fields; slangpy.cpp appends it verbatim to a SignatureBuffer. Field order is therefore not a parsing constraint. Per review, every fixed-width field precedes the variable-length V run, which is always last — a property worth holding regardless of whether a parser exists today, because it keeps the format cheap to extend.
• Compat gate (src/slangpy_ext/utils/torch_bridge.h) — decides whether the native bridge is used or the Python fallback. It discriminates on api_version and info_struct_size only.
• Contractual vs actual bound — two different buffer-sizing rules, see below.

Why API_VERSION had to move 8 → 9

main and this branch both defined version 8, for different wire formats: main's 8 is [Dn,Sm,V...] (#1082), this branch's was [Dn,Sm,Gk]. Because TensorBridgeInfo's layout is unchanged, info_struct_size cannot discriminate them — so api_version is the sole signal. Leaving both at 8 would let a stale bridge pass the gate and silently emit the wrong format: no error, just a mis-keyed cache. The unlanded side (this one) renumbers to 9.

test_stale_bridge_version_is_rejected pins this: it shims a bridge with api_version = real - 1 and the correct info_struct_size, so only the version differs, and asserts the gate reports incompatible. That is the guard against a future format change forgetting the bump.

The positive direction is covered inside test_is_torch_bridge_available, which asserts that when slangpy_torch is importable it is also the active path — an incompatible bridge falls back silently, so without that assertion the suite would stay green with the native path dead.

Buffer sizing

The native guard is unchanged: required_size = TENSOR_BRIDGE_SIGNATURE_BASE_SIZE + ndim. It is not padded for the extra ,G<bit>, because BASE_SIZE (64) is a deliberate fixed over-allowance and the composed format still fits inside it:

[D + 10 digits + ,S + 10 digits + ,G + 1 + ,V + ] + NUL  =  31 fixed bytes  (+ ndim for the V run)
native allowance: 64 + ndim                              =>  33 bytes spare at any rank

Stated assumption: fast_itoa takes a plain int with no clamp or range check, so budgeting 10 decimal digits each for ndim and scalar_type is a true bound over representable inputs rather than an assumption about typical values. (Bound derivation credit: reviewer analysis on the tracking thread.)

Padding the guard would also have broken test_native_signature_buffer_size_contract, which pins the bound exactly — required_size - 1 must fail and required_size must succeed.

Bounds behaviour differs between the two paths (pre-existing, tracked separately)

Format is in lockstep; bounds are not. The native path enforces a contractual bound (BASE_SIZE + ndim) while the fallback enforces an actual-length bound (sig.size() + 1). This PR changes neither — src/slangpy_ext/utils/torch_bridge.h is not in its diff and the native guard is untouched.

Tracked as #1091 (P2), out of scope here. Documented in this PR only so the asymmetry is not mistaken for something this change introduced — not as a claim that it is harmless.

Verification

Built and run on an NVIDIA L40S (torch 2.13.0+cu126, slangpy_torch.API_VERSION = 9).

• 731 passed, 0 failed across test_torch_bridge.py, test_bridge_fallback_gaps.py, test_torchintegration.py, and the regression (86 pre-existing skips for unavailable device types).
• Regression passes on all four combinations: [native-vulkan], [native-cuda], [fallback-vulkan], [fallback-cuda].
• Stale-bridge rejection demonstrated by execution, not assertion — faking api_version 8 with the correct info_struct_size 72 against compiled 9 yields fallback_reason='incompatible', using_fallback=True.
• Positive control — removing the grad bit from the fallback emitter makes test_signature_distinguishes_requires_grad[fallback] fail with [D2,S6,V44] against the expected [D2,S6,G0,V44]; restoring it passes. The test discriminates rather than merely passing.
• pre-commit run --all-files clean.

Not verified locally: ninja exits non-zero on 89 .pyi stub-generation commands — a pre-existing environment issue (the stub step invokes /usr/bin/python3, which lacks numpy, instead of the venv). Zero C++ compile or link errors; both extensions import. Left untouched as unrelated to this change.

Notes for reviewers

• Signature expectations live in three files, and one is outside the conflict set. test_bridge_fallback_gaps.py never conflicted because this branch had not touched it, so git did not flag it — but it hardcodes signature literals and would have failed post-merge. Five more expectations are @pytest.mark.parametrize rows asserted ~15 lines from the table.
• Rank-0 is the sharp case. torch.empty(()) has an empty V run, so the signature is [D0,S6,G0,V]. ,G is written unconditionally before the extent loop, so an empty run needs no special case.
• The requires_grad bitfield in TensorBridgeInfo predates this work (50c4656, PyTorch optimizations #759) and is unrelated to the cache key — its presence on main is not evidence of a grad fix.
• NativeTorchTensorDiffPair::read_signature is deliberately untouched: that path is unreachable pre-existing due to a Eliminate heap allocations from cached dispatch hot path #872 regression (documented at test_torchintegration.py:1286-1291, which explicitly does not cover it), so there is no second live signature path to keep in lockstep. The asymmetry is intentional, not an oversight.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[CLAassistant]

All committers have signed the CLA.

[nv-slang-bot]

Thanks for the review, @ccummingsNV — appreciated.

Status: approved on the current HEAD (ebb9f68); all builds + pre-commit are green. No outstanding review threads.

This PR is intentionally held as a draft under our automated-contributor policy — the bot does not self-promote to ready-for-review or self-merge. Whenever you're happy for it to land, a maintainer can mark it ready and squash-merge (it Fixes #1052).

One note: license/cla currently shows pending for the bot author — that check may need a maintainer override or is expected for this identity.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[nv-slang-bot]

Reworked and re-pushed. The earlier approval no longer covers this code, so this needs a fresh look rather than a re-confirmation.

Why it changed. #1082 ("Torch tensors supply vector dims") landed on main after this branch was cut and rewrote the same signature emitter. The two changes key the call-data cache on different axes — #1082 on per-dimension shape classifiers, this one on grad-ness — so taking either side wholesale would have regressed the other. The fix is now a composition of both:

[Dn,Sm,V...,G<0|1>]

What to look at:

• TENSOR_BRIDGE_API_VERSION 8 → 9. main and this branch both defined 8 for different wire formats. Since TensorBridgeInfo's layout is unchanged, info_struct_size cannot tell them apart, so api_version is the only discriminator — two live 8s would let a stale bridge pass the compat gate and silently emit the wrong format.
• The buffer guard is deliberately unchanged (BASE_SIZE + ndim, not padded for the extra field). The PR body carries the bound derivation and the assumption it rests on.
• Bounds behaviour differs between the native and fallback paths. That divergence pre-exists on main and is filed separately as Torch bridge: native and fallback get_signature enforce different buffer-size bounds #1091 — out of scope here, but documented in the body so it isn't mistaken for something this change introduced.
• A third test file was updated outside the conflict set. test_bridge_fallback_gaps.py never conflicted because this branch hadn't touched it, so git didn't flag it — but it hardcodes signature literals and would have failed post-merge.

Verified on an L40S: 206 passed / 0 skipped, the regression green on native/fallback × vulkan/cuda, and the grad-bit test confirmed to actually discriminate (removing the bit makes it fail). Full detail in the PR description.

Staying in draft — promotion is a maintainer call.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[ccummingsNV]

couldn't this just be *p++ = tensor.requires_grad ? '1' : '0'

[nv-slang-bot]

Yes — done. fast_itoa was pointless for a value that is always one character; it is now *p++ = tensor.requires_grad() ? '1' : '0';, and the requires_grad local is gone since the ternary reads the tensor directly.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[ccummingsNV]

I'd prefer for the none-variable-length signature parts (inc this new one) to be before the variable length part

[nv-slang-bot]

Done — format is now [Dn,Sm,Gk,V...], so every fixed-width field precedes the variable-length run.

Worth noting your reasoning is better than what I had. I originally proposed exactly this ordering, then talked myself out of it on the grounds that nothing parses the signature (it is an opaque cache key), so I put G last to keep main's [Dn,Sm,V...] prefix byte-identical and minimise the diff. That optimised for the wrong thing: "all fixed-width parts before the variable-length part" is a property worth holding regardless of whether a parser exists today, because it is what keeps the format cheap to extend.

Applied in the native emitter, the Python fallback (bridge_fallback.py), and all 13 signature expectations across three test files.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[ccummingsNV]

see below - let's make constant length parts before the variable length part of the sig

[nv-slang-bot]

Updated to "[Dn,Sm,Gk,V...]" with the invariant stated explicitly — every fixed-width field precedes the variable-length V run, which is always last. The fallback docstring in bridge_fallback.py carries the same wording so the two cannot drift.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[ccummingsNV]

I don't think we need more of these bridge tests - they're testing the same stuff we already test for in different ways

[nv-slang-bot]

Removed test_native_bridge_version_matches — you are right that it largely duplicates test_is_torch_bridge_available.

One assertion in it was not covered elsewhere, so I folded that single line into the existing test rather than keep a whole new one: is_torch_bridge_using_fallback() is False when slangpy_torch is importable. Without it a version-incompatible bridge falls back silently and the suite stays green with the native path dead — which is the failure mode that made the 8→9 bump necessary in the first place. Three lines inside a test that already exists, instead of a separate one.

test_stale_bridge_version_is_rejected I have kept: it is the only test that exercises the incompatible path deliberately, by shimming a bridge with a stale api_version and the correct info_struct_size so only the version differs. Happy to drop it too if you would rather, but it is the guard against a future format change forgetting the version bump.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}