Fix cuda overflow

[lucylq]

Update copy_ method to check

• this->sizes() == other.sizes()
• this->dim() == other.dim()
  Use overflow-safe arithmetic.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19487

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ 1 Pending, 1 Unclassified Failure

As of commit ce6fc1b with merge base 23a91d5 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Check Labels / Check labels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: GraphQL query

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[lucylq]

@claude review

[claude]

Claude finished @lucylq's task in 1m 6s —— View job

---

Code Review: Fix cuda overflow

• Gather context and read the changed file
• Analyze the diff for correctness, security, and style
• Post review feedback

---

Overall: This is a solid security hardening PR. The changes add proper dimension/size validation and replace unchecked arithmetic with overflow-safe operations in the non-contiguous copy_ path. The approach is correct and well-structured.

Strengths

1. Dimension and size checks (lines 443-449): Good defense-in-depth. Checking dim() and sizes() equality before the existing numel() check catches shape mismatches that happen to have the same total element count (e.g., [2,3] vs [3,2]).

2. Overflow-safe offset computation (lines 474-484): Using c10::mul_overflows and c10::add_overflows with early-exit on overflow is the right pattern. The accumulation loop correctly builds both src and dst offsets simultaneously.

3. Byte offset overflow check (lines 488-498): Good catch — the element-to-byte conversion (offset * elem_size) is another multiplication that can overflow, and the non-negativity assertion on src_offset/dst_offset prevents signed-to-unsigned conversion issues.

Minor observations

1. Redundant numel() check: With dim() and sizes() equality now enforced, the numel() check on line 451 is logically redundant (equal sizes implies equal numel). It's fine to keep as a belt-and-suspenders assertion, but worth noting it can never fire if the prior checks pass.

2. compute_numel in size_util.h is still unchecked (slim/util/size_util.h:23-29): The refresh_numel() path calls compute_numel which does n *= s without overflow protection. If an attacker can control sizes passed to set_sizes_and_strides, the numel could overflow before copy_ is ever called. This is out of scope for this PR but worth a follow-up.

3. nbytes() on line 101: numel() * itemsize() is also a potential overflow site for the contiguous fast path (storage_->copy_(this->data_ptr(), other.data_ptr(), other.nbytes(), ...)). Again out of scope, but relevant to the broader hardening effort.

Verdict

LGTM — the changes are correct, focused, and improve the security posture of the non-contiguous copy path. The new checks are properly ordered (cheapest first) and the overflow detection uses battle-tested compiler builtins via c10::mul_overflows/c10::add_overflows.

---

[Copilot]

Pull request overview

This PR hardens SlimTensor::copy_ in the AOTI slim runtime by enforcing stricter shape compatibility and making strided offset/byte computations overflow-safe (particularly relevant for CUDA pointer arithmetic).

Changes:

• Added precondition checks that dim() and sizes() match before copying.
• Reworked strided element-wise copy offset computation to use c10::{mul,add}_overflows and to compute byte offsets explicitly.
• Updated CPU/CUDA memcpy call sites to use the computed byte offsets.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Gasoonjia]

i feel like we only need two overflow calculation, one for src = counter[d] * this->stride(d), the other for dst_offset = dst_offset + src, thoughts?

[lucylq]

yeah - thank you!

[Gasoonjia]

im not a huge fan of adding this new limitation. Two tensor should be copyable as long as their numel are the same. Maybe we can remove this assertion?

[lucylq]

I think we'd also need to update the loop below if we want tensors to be copyable with different dim/sizes..

[Gasoonjia]

thx for adding the overflow check! plz solve the broken cuda cis as well.

[Gasoonjia]

This check also enforces two tensors having same shape -- can we remove it?

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.

[Gasoonjia]

@claude review the PR

[claude]

Claude finished @Gasoonjia's task in 1m 49s —— View job

---

Code Review: Fix cuda overflow

• Gather context and read the changed file
• Analyze the diff for correctness and security
• Review feedback from other reviewers
• Post review feedback

---

Overall: Good security hardening. The overflow-safe arithmetic is correct and uses the right primitives. There are a few open design questions and minor issues worth addressing.

Correctness & Security

1. Overflow checks are sound — c10::mul_overflows/c10::add_overflows wrap compiler builtins (__builtin_mul_overflow etc.) and correctly detect overflow before the result is used. The pattern of computing into temporary terms then accumulating is correct.

2. static_cast<size_t>(src_offset) after >= 0 check (slim_tensor.h:492-498) — The non-negativity check ensures the cast doesn't reinterpret a negative value as a large unsigned number. On 64-bit targets (which all CUDA platforms are), size_t == uint64_t, so no truncation occurs. On a hypothetical 32-bit build, values > UINT32_MAX would truncate before the overflow check. This is extremely unlikely in practice for SlimTensor's use cases, but if you want to be rigorous, add a static_assert(sizeof(size_t) >= sizeof(int64_t)) or a range check.

3. Negative strides (slim_tensor.h:492) — The src_offset >= 0 && dst_offset >= 0 assertion rejects tensors whose strided traversal produces negative offsets (e.g., torch.flip). This is fine for SlimTensor (which likely never produces negative strides), but should be called out explicitly since it's a behavioral constraint that isn't obvious from the function signature.

Design Feedback

4. dim()/sizes() equality restriction (slim_tensor.h:443-449) — Agreeing with @Gasoonjia's feedback: standard PyTorch copy_ only requires equal numel(), allowing e.g. [6].copy_([2,3]). The current strided loop assumes equal dims (uses this->dim() to iterate both src and dst), so enforcing this is correct for the current implementation. If relaxing this is desired later, the loop would need to iterate dimensions independently per tensor. For now, I'd suggest keeping the check but updating the doc comment to document this restriction explicitly.

5. Redundant numel() check (slim_tensor.h:450-451) — With equal sizes() enforced, numel() equality is mathematically guaranteed. Fine to keep as defense-in-depth, but if you'd like to slim down, it can be removed.

6. Stale doc comment (slim_tensor.h:437) — Says "Currently only supports CPU-to-CPU copy (contiguous tensors only)" but the implementation handles CUDA and non-contiguous. Should be updated to reflect reality.

Follow-up Items (out of scope but worth tracking)

7. compute_numel in size_util.h:23-28 uses unchecked n *= s. If sizes are attacker-controlled, numel can overflow before reaching copy_. This is a broader issue than this PR but worth a follow-up.

8. nbytes() at slim_tensor.h:101 (numel() * itemsize()) is also unchecked and used in the contiguous fast path at line 461. Same follow-up category.

9. CI failures — @Gasoonjia mentioned broken CUDA CIs. These should be investigated separately — the overflow logic itself looks correct for both CPU and CUDA paths.

Verdict

The core security fix (overflow-safe offset arithmetic) is correct and well-implemented. The main discussion point is the dim()/sizes() restriction (items 4, addressed in review thread). Updating the stale doc comment and considering the 32-bit cast edge case would be nice cleanups.

---

[Gasoonjia]

Update the doc as well to show that two tensors should have same sizes
And can you do me a favor to remove Currently only supports CPU-to-CPU copy (contiguous tensors only). line? it has out of date loll😛

[lucylq]

@claude review

[claude]

Claude finished @lucylq's task in 1m 47s —— View job

---

Code Review: Fix cuda overflow

• Gather context and read the current state of the changed file
• Analyze the diff against merge base
• Review against prior feedback (Gasoonjia's comments, copilot suggestions)
• Post review feedback

---

Overall: The security hardening looks correct and well-structured. The overflow-safe arithmetic uses the right primitives, and the doc comment has been properly updated. The PR addresses the prior review feedback well (stale comment removed, sizes restriction documented). A few remaining items:

Looks Good

1. Overflow detection pattern (lines 479-484): Using c10::mul_overflows/c10::add_overflows with short-circuit evaluation is correct — if any operation overflows, the check fires before the corrupted value propagates.

2. Byte offset overflow check (lines 490-498): The non-negativity guard before static_cast<size_t> prevents signed-to-unsigned reinterpretation. On all CUDA platforms (64-bit), sizeof(size_t) == sizeof(int64_t), so no truncation.

3. Doc comment (lines 433-439): Updated to match the new preconditions. Stale "CPU-to-CPU" reference removed per @Gasoonjia's request.

Remaining Items

1. data_ptr() has unchecked arithmetic (slim_tensor.h:212):

   return static_cast<char*>(storage_->data()) + storage_offset_ * itemsize();

   This multiplication can overflow for the same reason the strided offsets could. It's called in the contiguous fast path (line 459-460) — a malicious storage_offset_ could produce an OOB pointer before the new overflow checks even run. Worth a follow-up PR.

2. nbytes() unchecked (slim_tensor.h:100-101): numel() * itemsize() feeds directly into the contiguous storage_->copy_() call at line 460. Same category as above.

3. compute_numel unchecked (size_util.h:23-28): Uses n *= s without overflow detection. If sizes are attacker-controlled via set_sizes_and_strides, numel can silently wrap before copy_ is called, potentially causing the numel() check (line 450) to pass incorrectly.

4. CI failures: @Gasoonjia mentioned broken CUDA CIs — these should be investigated separately as the logic here looks correct for both CPU and CUDA paths.

Items 1-3 are pre-existing issues and out of scope for this PR, but worth tracking as follow-ups to complete the hardening effort.

Verdict

LGTM — the core fix is correct and focused. The new checks are ordered cheaply (dim/sizes comparison before the per-element loop) and the overflow detection uses well-tested compiler builtins. The remaining unchecked paths (data_ptr, nbytes, compute_numel) are pre-existing and appropriate for follow-up work.

---