fix(linalg): make GPU scalar in-place ops correct on CUDA (#988)

[yingjerkao]

Fixes #988. Follow-up to #906 / #980 (based on fix/scalar-inplace-ops-preserve-storage until it merges, then this should retarget master).

Validated on a CUDA device (RTX 4070 Ti, CUDA 13) which the #980 review lacked.

What #988 asked for

1. Regression: non-contiguous GPU LHS threw for *= / /= (scalar), while iAdd/iSub handled it.
2. Efficiency: the shape-{1} scalar wrapper did a per-call H2D copy.

Changes (src/linalg/i{Add,Sub,Mul,Div}.cpp, src/Tensor.cpp)

1. Non-contiguous scalar regression. iMul/iDiv now route a broadcast scalar on a non-contiguous GPU LHS through the kernel — the rconst path scales every physical element in place, so the layout mappers are irrelevant and the result is correct. Because cuMul/cuDiv (unlike cuAdd/cuSub) do not consume the layout mappers, a genuine non-contiguous tensor*=tensor / tensor/=tensor now throws loudly instead of silently pairing mismatched elements.

2. No per-call H2D copy. _scalar_as_rank1_tensor keeps the wrapper CPU-resident. For a length-1 RHS the GPU kernels read the scalar with a host-side dereference and pass it in by value, so no copy is needed; the device check permits a host-resident length-1 scalar against a GPU LHS, and cudaSetDevice targets Lt.device(). (This also removes a latent UB: fix(Tensor): scalar in-place arithmetic mutates storage in place #980 moved the scalar to the GPU, making that _Rin[0] host dereference read a device pointer.)

3. Safety — real op= complex corrupted memory. The GPU path reinterpreted the real output buffer as complex (out-of-bounds write). Added a device-independent guard so it throws like the CPU path.

4. Safety — narrow LHS corrupted memory. A LHS narrower than the promoted dtype (e.g. float_gpu *= 2.0) had the promoted-width result written into the narrower buffer (OOB; float += 1.0 produced 0). Now computed in the promoted dtype and truncated back, matching the CPU element-wise semantics.

Tests (tests/gpu/Tensor_test.cpp, Gpu… — no underscores per #857)

Non-contiguous *=//= regression, host-scalar contiguous values, storage sharing, dtype preservation, real op= complex throws, and the non-contiguous tensor⊗tensor guard.

Verification (RTX 4070 Ti, CUDA 13)

• 7/7 new GPU tests pass.
• 8/8 existing CPU scalar in-place tests still pass (guard doesn't change CPU behavior).
• 224 GPU arithmetic tests pass, 0 failures — including every existing iSub/iMul/iDiv all-types in-place case.

Known limitation (follow-up)

General non-contiguous tensor⊗tensor mul/div on GPU still throws; the real fix is adding non-contiguous kernels to cuMul/cuDiv like cuAdd/cuSub have. Left for a separate PR.

🤖 Generated with Claude Code

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[pcchen]

Reviewed the diff and built the branch + ran it on a CUDA device (Tesla V100, CUDA 12.9) to check the GPU-only paths — note CI has no CUDA runner, so none of this code (incl. the 7 new GPU tests) is exercised by CI; correctness rests on local hardware runs.

Overall: solid, well-motivated fixes. Two issues below, both reproduced on the V100.

What's good

• cudaSetDevice(Rt.device()) → Lt.device()) is required now that the scalar stays host-resident (Rt.device()==cpu==-1 would be an invalid device). ✅
• Dropping the per-call H2D scalar copy (and the latent device-pointer-deref UB it removes). ✅
• Device-independent real op= complex guard — throws early on both paths instead of corrupting the real buffer. ✅
• The non-contiguous scalar regression fix (route size-1 RHS through the rconst kernel, throw loudly on genuine non-contiguous tensor⊗tensor) is the right call, and the comments are genuinely excellent.

---

⚠️ 1. Narrow-LHS scalar path silently breaks storage sharing on GPU (CPU/GPU divergence)

The narrow-LHS guard does Lt = promoted.astype(Lt.dtype()), which rebinds Lt to fresh storage, so an aliased handle stops tracking the update. The CPU path applies lhs[idx] op= rhs[idx] straight into the existing buffer and preserves sharing. Measured:

CPU  shared_after=1  b[0]=2   # aliased handle sees a *= 2.0
GPU  shared_after=0  b[0]=1   # aliased handle does NOT

(a = ones({4}, Float, dev); b = from_storage(a.storage()); a *= 2.0;)

This diverges from the "preserve storage sharing" invariant of the parent chain (#906/#980). The current tests miss it: GpuScalarInplaceKeepsStorageSharing uses a Double LHS (non-narrow) and GpuScalarInplacePreserveDtype uses a Float LHS but never checks aliasing. Suggest copying the truncated result back into Lt's existing storage in place rather than rebinding, plus a narrow+aliased test.

⚠️ 2. The narrow-LHS memory-corruption fix only covers host scalars, not a GPU-resident RHS

The guard is gated on rhs_is_host_scalar, but the real danger condition is Lt.dtype() > Rt.dtype() on GPU. The shared contiguous branch still does tmpo = Lt.clone() (narrow) while the kernel writes promoted width, so a same-shape mixed-dtype op corrupts — float_gpu *= double_gpu (RHS all 3.0):

no throw; a = 0 2.125 0 2.125    # expected 3 3 3 3  → out-of-bounds garbage

This is pre-existing (the clone() predates this PR — not a regression) and technically outside #988's scalar scope, but it's the exact corruption class the PR's point 4 describes, just via a non-host RHS. Fixing at the source — tmpo = Lt.astype(type_promote(Lt, Rt)) in the shared branch — would close all paths at once and let the special-case guard be dropped. At minimum worth tracking, since a *= b with mixed dtype on GPU silently corrupts.

---

Minor

• iAdd(gpu_tensor, cpu_shape{1}_tensor) now silently succeeds (treated as a broadcast scalar) instead of erroring on device mismatch — intended for the wrapper, but a slight public-API broadening worth a doc note.
• Heavy verbatim duplication of the ~14-line guard block across all four files (and the non-contig scalar kernel block between iMul/iDiv); a shared helper would reduce divergence risk.
• Base is the unmerged fix/scalar-inplace-ops-preserve-storage — needs retargeting to master before merge (already noted in the description).

_{🤖 Reviewed with Claude Code; GPU findings reproduced on a Tesla V100 / CUDA 12.9 build of this branch.}

Posted by Claude Code on behalf of @pcchen

[pcchen]

Code Review

LGTM on substance — and this PR materially matters for #980: it fixes a real UB that #980 introduces on CUDA. #980's {1}-wrapper does .to(device), but the GPU kernels read the length-1 scalar with a host-side _Rin[0] dereference — dereferencing a device pointer on the host. There is no GPU CI, so checks can't see it; this PR was validated on real hardware (RTX 4070 Ti, 224 GPU tests green). Consequence: #980 and #992 should land together (or back-to-back) — #980 alone leaves CUDA scalar in-place broken.

Verified in the diff (all four i-ops)

• Host-resident wrapper (drops .to(device)) + rhs_is_host_scalar device-check exemption + cudaSetDevice(Lt.device()) (the old code targeted Rt.device() — wrong once the RHS is host-resident). ✅
• real ⊕= complex device-independent guard — the CPU visit already threw, but the GPU path bypassed it and reinterpreted a real output buffer as complex (memory corruption). Now rejected uniformly, matching fix(Tensor): scalar in-place arithmetic mutates storage in place #980's CPU semantics. ✅
• Narrow-LHS OOB fix: the GPU in-place kernels write the promoted-width result into the LHS buffer → out-of-bounds write when the LHS is narrower. The detour computes at the promoted dtype then truncates back — matching the CPU/fix(Tensor): scalar in-place arithmetic mutates storage in place #980 value semantics — and the recursion terminates (the promoted LHS is no longer narrower). ✅
• Non-contiguous: a scalar RHS on a non-contiguous GPU LHS now routes through the rconst path (layout mappers are irrelevant for scale-every-element — the GPU: scalar in-place ops throw on non-contiguous tensors for *= and /=, and do a per-call H2D copy #988 regression fix); genuine non-contiguous tensor⊗tensor mul/div still throws, now with an accurate message (cuMul/cuDiv ignore the mappers; routing would silently mispair elements). The known-limitation note (add real non-contig kernels later) is honest. ✅

Notes (non-blocking)

1. 🟡 Residual, pre-known (Scalar in-place Tensor arithmetic detaches storage-sharing views #906-tracked): the narrow-LHS GPU path does Lt = promoted.astype(...) — a rebind, so storage-sharing views detach in that specific case (CPU preserves sharing). This extends the already-documented "GPU can rebind on promotion" limitation rather than introducing a new class — but deserves a one-liner in the PR body.
2. 🟠 Process: the branch still contains the pre-rebase fix(Tensor): scalar in-place arithmetic mutates storage in place #980 commits (the base branch was force-pushed during its conflict resolution), so this needs a cherry-pick of its own commit onto the current fix/scalar-inplace-ops-preserve-storage head — happy to do that.
3. GPU tests can't run in CI (Build CI CD for gpu. #538) — local-hardware evidence only; same standing caveat as the MKL paths (Run GitHub action with all allowed configurations and cache the build result #583).

Posted by Claude Code on behalf of @pcchen

[codecov]

Codecov Report

❌ Patch coverage is 11.62791% with 76 lines in your changes missing coverage. Please review.
✅ Project coverage is 32.95%. Comparing base (dfe688b) to head (3fda11f).
⚠️ Report is 15 commits behind head on master.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/Tensor.cpp	4.00%	32 Missing and 16 partials ⚠️
src/linalg/iAdd.cpp	22.22%	4 Missing and 3 partials ⚠️
src/linalg/iDiv.cpp	22.22%	4 Missing and 3 partials ⚠️
src/linalg/iMul.cpp	22.22%	4 Missing and 3 partials ⚠️
src/linalg/iSub.cpp	22.22%	4 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #992      +/-   ##
==========================================
+ Coverage   31.84%   32.95%   +1.10%     
==========================================
  Files         230      230              
  Lines       33117    33109       -8     
  Branches    13846    13813      -33     
==========================================
+ Hits        10547    10910     +363     
+ Misses      15549    14832     -717     
- Partials     7021     7367     +346     

Flag	Coverage Δ
cpp	32.59% <11.62%> (+1.11%)	⬆️
python	61.84% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Components	Coverage Δ
C++ backend	33.70% <11.62%> (+1.13%)	⬆️
Python bindings	25.15% <ø> (+1.11%)	⬆️
Python package	61.84% <ø> (ø)

Files with missing lines	Coverage Δ
src/linalg/iAdd.cpp	37.50% <22.22%> (-12.50%)	⬇️
src/linalg/iDiv.cpp	33.33% <22.22%> (-10.42%)	⬇️
src/linalg/iMul.cpp	33.33% <22.22%> (-10.42%)	⬇️
src/linalg/iSub.cpp	37.50% <22.22%> (-12.50%)	⬇️
src/Tensor.cpp	22.59% <4.00%> (+0.60%)	⬆️

... and 38 files with indirect coverage changes

---

Continue to review full report in Codecov by Harness.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update dfe688b...3fda11f. Read the comment docs.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[pcchen]

Approving. Verified across all four i-ops: the host-resident {1} wrapper removes the CUDA host-deref UB that #980 alone would have shipped (plus the per-call H2D copy), the rhs_is_host_scalar device-check exemption and cudaSetDevice(Lt.device()) fix are correct, the device-independent real ⊕= complex guard closes a GPU memory-corruption path, the narrow-LHS promote-compute-truncate-back detour eliminates the OOB write while matching CPU value semantics, and the non-contiguous scalar regression is fixed via the mapper-irrelevant rconst path (with genuine non-contiguous tensor⊗tensor mul/div now failing loudly instead of silently mispairing — honest follow-up noted for real kernels). Hardware-validated on RTX 4070 Ti (224 GPU tests); CPU-side CI fully green on master base including BuildAndTest on all platforms — the only red is advisory codecov/patch, structural for GPU-only lines with no GPU CI (#538/#583). The residual narrow-LHS rebind (views detach in that GPU case) is the pre-known #906-class limitation.

Posted by Claude Code on behalf of @pcchen