Guard roi_align/ps_roi_align bilinear sampling against non-finite boxes - #9629
Guard roi_align/ps_roi_align bilinear sampling against non-finite boxes#9629fjankovi wants to merge 1 commit into
Conversation
A non-finite box coordinate makes every comparison in the "inverse
elements are out of feature map boundary" check false, so the sample
point is neither rejected by the empty branch nor clamped by the
y <= 0 / x <= 0 clamps below it. The index is then computed from
(int)y with y == NaN, which is undefined and yields INT_MIN on x86,
so pos1..pos4 become large negative offsets and the forward kernels
read far outside the input buffer:
x = torch.zeros(1, 1, 4, 4)
b = torch.tensor([[0., nan, nan, nan, nan]])
torchvision.ops.roi_align(x, b, (1, 1), 1.0, 1, True) # SIGSEGV
NaN is not the only way in, so filtering the boxes up front would not
be enough: [inf, -inf, inf, -inf] makes roi_width NaN via inf - inf,
and [-3.4e38, -3.4e38, 3.4e38, 3.4e38] -- entirely finite -- overflows
roi_end - roi_start to inf and then computes 0 * inf.
Write the guard as !(in range) instead of (out of range) so that a NaN
coordinate takes the empty branch. For finite inputs the two forms are
exactly complementary; comparing the PreCalc output of both over
20,055,768 sample points gives zero bitwise differences.
The backward kernels reach the same conversion but already gate every
write on x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0, which
INT_MIN fails, so no memory is corrupted -- this is an out-of-bounds
read only. It is also not observable: all four interpolation weights
are NaN whenever the coordinate is, so whatever is read is multiplied
by NaN.
CUDA/MPS did not crash, since float-to-int conversion returns 0 rather
than INT_MIN there, but the guard is updated at all ten sites so the
four copies stay identical and CPU and GPU agree on the result.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9629
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @fjankovi! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
1 participant
Problem
A non-finite box coordinate makes every comparison in the "inverse elements are out of feature map boundary" check false, so the sample point is neither rejected by the empty branch nor clamped by the
y <= 0/x <= 0clamps. The index is then computed from(int)ywithy == NaN, which is undefined behaviour and yieldsINT_MINon x86.pos1..pos4become large negative offsets and the forward kernels read far outside the input buffer:UBSan and ASan on the unmodified kernel:
It is not only NaN that gets there, and an input-side
isfinitefilter would not be enough. Measured onx = torch.zeros(1, 1, 4, 4),roi_align(..., (1, 1), 1.0, 1, True):[nan, nan, nan, nan][inf, -inf, inf, -inf]inf - infmakes the width NaN[-3.4e38, -3.4e38, 3.4e38, 3.4e38]roi_end - roi_startoverflows toinf, thenph * bin_size_his0 * infThe backward kernels reach the same undefined conversion but do not corrupt memory: every write is already gated on
x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0, whichINT_MINfails. So this is an out-of-bounds read only, and because all four interpolation weights are NaN whenever the coordinate is, whatever is read is multiplied by NaN and cannot be observed in the output. The impact is a crash, not disclosure.Fix
Write the guard as
!(in range)instead of(out of range)so a NaN coordinate takes the empty branch. For every finite input the two forms are exactly complementary, so this is not a behaviour change outside the non-finite case.The same guard is duplicated at 10 sites in 6 files, and all are updated:
Verification
Built torchvision from source (
v0.23.0, same kernels, against torch 2.8.0) with and without this change and ran each case in its own process:roi_align0.0, forward + backward cleanps_roi_align0.0, forward + backward cleanSix of six crash before, six of six clean after. Note the regression test kills the interpreter on unpatched code rather than reporting an assertion failure, as memory-safety tests do.
Other checks:
roi_align_common.hside by side in one binary and comparing the resultingPreCalcstructs bitwise over 20,055,768 sample points (random geometries, spatial scales1e-3/1/1e6, 12,865,481 of them taking the empty branch) gives 0 mismatches.pytest test/test_ops.py -k "TestRoIAlign or TestPSRoIAlign or TestRoIPool or TestPSRoIPool"on the patched build: 143 passed, 0 failed.float-cast-overflow+ ASan SEGV to a clean0.0result.0for NaN rather thanINT_MIN(measured:(int)NaNfis0on gfx1201,INT_MINon x86-64). The GPU therefore sampled index0and returnednanwhere the CPU faulted. After this change both return0, which is what an out-of-range sample point has always produced. This is the one user-visible output change, and only for input that was undefined behaviour before.test_non_finite_boxesinRoIOpTester, so it runs forroi_align,ps_roi_align,roi_poolandps_roi_pool, forward and backward, on every device: 12 passed on CPU (the pool ops were already safe; they are covered so the whole family is pinned).Not fixed here
int roi_batch_ind = offset_rois[0];in the same kernels is an unchecked conversion too, and is not range-checked against the batch size anywhere (24 sites). The same NaN payload placed in column 0 still reads out of bounds at some shapes, and so does a plain out-of-range integer index — that is #4828, open since 2021. It needs a host-side check rather than a kernel-side predicate, so it is left for a separate change.Disclosure per
AI_POLICY.md: the analysis and this patch were prepared with AI assistance (Claude). Every measurement quoted above was produced by actually building and running the code described, not inferred.