Add the CUDA implementation to the Rotated Box Iou Calculation by zy1git · Pull Request #9404 · pytorch/vision

zy1git · 2026-02-25T08:03:57Z

Implementation Details (Adapted from Detectron2)
Uses shared memory optimization with 32×16 thread blocks (512 threads per block)
Handles transpose for large box counts (>65535 in Y dimension)
CUDA kernel only supports float32 (consistent with Detectron2)
Shared header allows both CPU and CUDA kernels to use the same math utilities

Testing
All existing tests now run on both CPU and CUDA
Added test_cuda_cpu_consistency to verify CUDA and CPU produce the same results
float64 tests are skipped on CUDA (kernel only supports float32)

pytorch-bot · 2026-02-25T08:04:01Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9404

📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 2 Unrelated Failures

As of commit 66db05b with merge base 1cc5693 ():

NEW FAILURES - The following jobs have failed:

Build Windows Wheels / pytorch/vision / build-wheel-py3_10-cuda12_6 (gh)
Build Windows Wheels / pytorch/vision / upload / upload-wheel-py3_10-cuda12_8 (gh)
Unable to download artifact(s): Artifact not found for name: pytorch_vision__3.10_cu128_x64

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Build Windows Wheels / pytorch/vision / build-wheel-py3_10-cuda12_8 (gh) (trunk failure)
Process completed with exit code 1.
Build Windows Wheels / pytorch/vision / upload / upload-wheel-py3_10-cuda12_6 (gh) (trunk failure)
Unable to download artifact(s): Artifact not found for name: pytorch_vision__3.10_cu126_x64

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

NicolasHug · 2026-02-25T17:50:27Z

        """Core test: IoU computation with different formats, dtypes, and rotations."""
+        # CUDA kernel only supports float32, skip float64 on CUDA
+        if device == "cuda" and dtype == torch.float64:
+            pytest.skip("CUDA kernel only supports float32")


use xfail, here and everywhere else where relevant

Actually, I implemented this:https://github.com/pytorch/vision/pull/9404/changes#r2857140067

Then we don't need this pytest.skip/xfail thing.

NicolasHug · 2026-02-25T17:52:02Z

+            ious, torch.tensor([[expected_iou]], dtype=torch.float32, device=device), atol=1e-4, rtol=1e-4
+        )
+
+    # ==================== CUDA-specific Tests ====================


Thanks for adding the test below, it's a good test. Just remove this comment please

Removed in the new commit.

NicolasHug

Thanks @zy1git , made a quick first pass

NicolasHug · 2026-02-25T18:03:49Z

+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  using scalar_t = float;
+  AT_ASSERTM(


I think we should use TORCH_CHECK instead of AT_ASSERTM for consistency with the rest

Thanks for pointing this out. I have updated it in the new commit.

NicolasHug · 2026-02-25T18:12:02Z

+    // input must be contiguous
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  using scalar_t = float;


Hm, is there any reason double isn't support on CUDA? I'm hoping we can make it work through AT_DISPATCH_FLOATING_TYPES like what was done for the CPU part. Can you try that and see if the tests are complaining?

Great point!

In the new commit, I can have the Input to be float64 for CUDA, but the output is still float32 to match the CPU behavior. And I removed the pytest.skip/xfail parts in the test file.

NicolasHug · 2026-02-25T18:14:01Z

+  __shared__ float block_boxes1[BLOCK_DIM_X * 5];
+  __shared__ float block_boxes2[BLOCK_DIM_Y * 5];


This seems wrong, shouldn't it be T?

Yeah, I agree. The previous implementation from Detectron2 used float. T is better.

Fixed in the new commit.

…ance for xyxyxyxy format on CUDA

zy1git · 2026-02-27T18:09:38Z

+  for (int i = 1; i < num_in - 1; i++) {
+    for (int j = i + 1; j < num_in; j++) {
+      T crossProduct = cross_2d<T>(q[i], q[j]);
+      if ((crossProduct < -1e-4) ||


This 1e-4 is different from the cpp implementation PR: #9379 to address the test failures caused by numerical precision issues.

NicolasHug · 2026-03-02T13:49:59Z

+template <typename T>
+struct Point {
+  T x, y;
+  HOST_DEVICE Point(const T& px = 0, const T& py = 0) : x(px), y(py) {}


These, and others, are HOST_DEVICE_INLINE in the detectron2 implem. Is there a particular reason we're using HOST_DEVICE here instead?

Fixed in the new commit. Good catch! HOST_DEVICE_INLINE is better here because it eliminates function call overhead for these small, frequently-called helper functions.

NicolasHug · 2026-03-02T13:53:00Z

+  // Sort point 1 ~ num_in according to their relative cross-product values
+  // (essentially sorting according to angles)
+  // If the angles are the same, sort according to their distance to origin
+  T dist[24];


detectron2 has a different path for CUDA and for CPU here: https://github.com/facebookresearch/detectron2/blob/fd27788985af0f4ca800bca563acdb700bb890e2/detectron2/layers/csrc/box_iou_rotated/box_iou_rotated_utils.h#L202

Any specific reason we're not doing something similar?

The CUDA and CPU paths are almost identical, except the CPU path recalculates the distance after sorting, which is redundant since we swap both q[] and dist[] together during sorting. I think Detectron2 kept this step because they originally intended to use std::sort (now commented out) for CPU, which wouldn't swap dist[] alongside q[].

In the new commit, I separated the CUDA and CPU paths to match Detectron2 exactly, and reverted the tolerance from 1e-4 to 1e-6. The test failures returned in my local run (not CI yet), confirming that the two-path structure is not the cause.

I think we can still combine the two paths into one, given the analysis above. Would you prefer I do that, or keep them separated to match Detectron2?

github-actions · 2026-03-10T16:20:06Z

Hey @NicolasHug!

You merged this PR, but no labels were added.
The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py

Add the CUDA implementation to the Rotated Box Iou Calculation

7e3d64e

meta-cla Bot added the cla signed label Feb 25, 2026

Merge branch 'main' into rotated-cuda

0fbf751

NicolasHug reviewed Feb 25, 2026

View reviewed changes

Zhitao Yu added 8 commits February 25, 2026 22:44

address comments

fdaebe4

remove unnecessary blank lines

4e95443

address the test failures by adding a clamping and use a higher toler…

1db13ec

…ance for xyxyxyxy format on CUDA

Revert workarounds to expose CUDA numerical precision test failures

ac2597a

delete unnecessary comment

eba567b

fix the CUDA test failures

61b3264

Fix CUDA test failures by relaxing crossProduct tolerance in Graham scan

4c5d62a

remove unnecessary comment

c3e325f

zy1git commented Feb 27, 2026

View reviewed changes

Comment thread torchvision/csrc/ops/box_iou_rotated_utils.h

NicolasHug reviewed Mar 2, 2026

View reviewed changes

Zhitao Yu and others added 3 commits March 3, 2026 01:01

address the comments to match Detectron2

a596846

Temporarily fix rotated box IoU CUDA tests with clamping and xfail

987720f

Merge branch 'main' into rotated-cuda

6716b65

NicolasHug approved these changes Mar 10, 2026

View reviewed changes

Merge branch 'main' into rotated-cuda

66db05b

NicolasHug merged commit cf080e9 into pytorch:main Mar 10, 2026
65 of 73 checks passed

		__shared__ float block_boxes1[BLOCK_DIM_X * 5];
		__shared__ float block_boxes2[BLOCK_DIM_Y * 5];

Uh oh!

Conversation

zy1git commented Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9404

❌ 2 New Failures, 2 Unrelated Failures

Uh oh!

NicolasHug Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NicolasHug left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Feb 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Feb 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Mar 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zy1git commented Feb 25, 2026 •

edited

Loading

pytorch-bot Bot commented Feb 25, 2026 •

edited

Loading

NicolasHug Feb 25, 2026 •

edited

Loading

zy1git Feb 26, 2026 •

edited

Loading

zy1git Feb 26, 2026 •

edited

Loading

zy1git Feb 27, 2026 •

edited

Loading

zy1git Mar 3, 2026 •

edited

Loading