Repository navigation
Add a portable kernel for antialiased bicubic resize - #24400
shoumikhin wants to merge 31 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24400
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 31bf43b with merge base 05092c1 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
`to_edge` cannot lower `F.interpolate(x, mode="bicubic", antialias=True)` today. It exports to `aten._upsample_bicubic2d_aa`, which PyTorch does not decompose into core ATen ops and no ExecuTorch kernel implements, so `to_edge` stops with: ``` SpecViolationError: Operator torch._ops.aten._upsample_bicubic2d_aa.default is not in Core ATen opset ``` ## What this change does `to_edge` now rewrites the op into core ATen ops that follow ATen's algorithm: - Each output is a normalized, weighted sum over a window of input pixels, with the Keys cubic filter (a = -0.5). Shrinking widens the window. - `align_corners`, explicit scale factors and dynamic sizes work. As in ATen's CPU kernel, a dim that keeps its size is copied. - A NaN or inf reaches only the outputs whose window has it, as in ATen. - float16 and bfloat16 are computed in float32 and rounded once. Integer inputs are not rewritten and still fail as above, because ATen rounds them in fixed point. The new private module `exir/_decompositions.py` has bilinear and bicubic functions. `to_edge` uses only the bicubic one, since the bilinear op has a portable kernel. Programs without the bicubic op do not change. To keep the op for your own kernel, list it in `EdgeCompileConfig(preserve_ops=...)` for `to_edge`; `_core_aten_ops_exception_list` alone no longer keeps it. A backend that runs the op can return it from its partitioner's `ops_to_not_decompose`. The rewrite holds every output's window at once, so it needs more memory than a kernel: about 200 MiB to shrink one 3-channel float32 1080x1920 image to 224x224. With dynamic output sizes, memory is planned for a window bound that can be far too large, so prefer a fixed output size, or a scale factor without `align_corners` or `recompute_scale_factor=True`. The CUDA backend keeps its own rewrite, which follows ATen's CUDA kernel, and could switch to these functions later. ## Testing - A new test in `exir/program/test/test_program.py` runs lowered bicubic resizes (static, dynamic, channels_last, NaN inputs, a dim that keeps its size) on the portable kernels against eager, and checks both functions against ATen directly. It fails without this change. - On CPU, over many sizes and scales, float64 matches eager ATen to about 1e-14 and float32 to about 5e-6, except an output width of 1 with a changing height, where ATen's CPU kernel disagrees with its own transposed result. bfloat16 and float16 match float32 ATen rounded once, within one unit in the last place except near zero. - Resizes lowered for the portable kernels and XNNPACK, static and dynamic, match eager. - Apart from the new test, the `exir`, XNNPACK and `backends/transforms` test suites give the same results with and without this change.
5b2469c to
f98cc73
Compare
Nothing calls it: the bilinear op keeps its portable kernel. With bicubic as the only user, the helpers no longer take a filter and its support.
A window never holds more input pixels than the input has, so the extra taps only cost memory. Capping them helps extreme shrinks and bounds the memory planned for dynamic sizes.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 26m 32s —— View job Needs discussion: this lets Full reviewPR Review: #24400SummaryAdds a decomposition of Problems found:
Nothing was run for this review because Python was blocked in this CI sandbox. Numerical agreement with ATen was not tested here. Backward Compatibility
Testing
Performance
Code Quality
RecommendationNeeds Discussion The main open question is whether a rewrite with this memory cost belongs in the default table for every
|
to_edge rejected F.interpolate(mode="bicubic", antialias=True) because the op it exports to, aten._upsample_bicubic2d_aa, is not a core ATen op and ExecuTorch had no kernel for it. The earlier commits on this branch decomposed it into gathers, which planned very large buffers for dynamic sizes. This replaces the decomposition with a portable kernel, the way the antialiased bilinear resize already works: the op is an exception in the edge verifier and runs on _upsample_bicubic2d_aa.out. The kernel follows ATen's CPU algorithm. It resizes the width and then the height with the same weights, and copies a dim whose size does not change. Half and bfloat16 are computed in float. Between the two passes it keeps only the rows that one output row reads, so its temporary memory stays small. Graphs without this op are unchanged.
The note said the resize was skipped because the portable library had no antialiased bicubic kernel. It has one now. The exported preprocessing still only pads, because prepare_image() resizes the image before export, so the note says that instead.
compute_scales_value is a port of ATen's helper, but it dropped ATen's check that a scale is positive. A scale of zero, a negative scale or NaN then became the resize ratio. ATen ignores such a scale and uses the input and output sizes instead, so do the same, and drop the bicubic kernel's own copy of the check, which also let NaN through. Of the other upsample kernels, only the antialiased bilinear one can pass such a scale to the helper, and it now matches ATen too. A test covers it. The others check their scales earlier.
A scale far smaller than the output size implies made the kernel's temporary sizes overflow, and it then wrote past its buffer. Very short outputs also asked for more temporary memory than a whole resized plane. - Reject a scale whose inverse is larger than the input size. F.interpolate never produces one, and it would overflow the window bounds. - Reject an empty input height or width, as ATen does. - Cap the window at the input size, and keep a ring of rows between the passes only when it is smaller than all of them. - Keep the start of an empty window inside the input. - Check that each temporary size fits in size_t, and align each buffer for its element type.
The tests now also run a batch of two, resizes of one dim in both layouts, an output width of one, float64, NaN inputs at both ends of a window, and a uint8 input, which the kernel rejects. A new float16 and bfloat16 case averages two rows of opposite sign, so a kernel that rounded its first pass to the half type would fail it, which the earlier cases could not catch. The to_edge test reports the shape mode and the input shape of a failure. For a contiguous input with an output width of one and a changing height, ATen's CPU kernel gives every row the first row's values. The kernel resizes the height there too, and its comment now says so, so that a later comparison with ATen does not take it for a bug.
The bespoke-recipe sections now list six operators, not five.
windows.h defines IN as an empty macro, so a template parameter named IN breaks the build if that header is ever included first. Use CTYPE_IN, as other portable kernels do.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 26m 49s —— View job Request Changes: listing the op in Full reviewPR Review: #24400SummaryAdds a portable API Design
Testing
RecommendationRequest Changes The opt-out has to change. Either support integer inputs and drop the check, or make the opt-out its own documented setting that partitioners can't turn on by accident. The test items are small. Progress
Nothing was built or run for this review. The findings come from reading the code, and the byte counts were worked out by hand from the kernel's allocation sizes. |
Images often reach the resize as uint8. ATen computes this case with weights in fixed point and rounds after each pass, and the kernel now does the same, so its uint8 results match ATen exactly. The edge verifier now lets uint8 through.
Listing the op in the config's exception list or preserve_ops turned the check off for the whole program, and partitioners fill preserve_ops on their own. Now that uint8 works, the check only caught dtypes that ATen's CPU kernel rejects too. Without it such a resize fails in the kernel's dtype check, as the antialiased bilinear resize does.
Running out of temporary memory for the rows between the passes, and height windows past the input while the width also changes, had no test. The half type tolerance comment now says what the tolerance allows.
A scale that F.interpolate never gives can make a window's weights nearly cancel out, so the normalized weights get too large for 16-bit fixed point. The kernel then shifted by a negative amount. It now fails with InvalidArgument, as ATen's generic CPU path fails on the same input.
The test compares uint8 bit for bit and float32 within 1e-5, so random inputs should not change between runs.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 41m 16s —— View job Approve: two minor uint8 problems are left. A negative weight can overflow Full reviewPR Review: #24400SummaryAdds a portable
Code Quality
Performance
RecommendationApprove Neither point blocks the merge. The overflow needs a scale that Progress
Nothing was built or run: Python needs approval in this sandbox, and there is no PyTorch source here. The findings come from reading the code, and the weights and byte counts were worked out by hand. The checkout is shallow, so the diff was taken against |
The fixed-point precision only bounds the largest weight. With a scale that F.interpolate never gives, a window can sum to a negative value, so its most negative weight is larger than any positive one and wrapped when stored in 16 bits. ATen returns a wrong value there (255 where the float result is 100). The kernel now fails with InvalidArgument instead.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 54m 59s —— View job Approve: a few minor uint8 problems remain: the Full reviewPR Review: #24400SummaryAdds a portable Code Quality
Testing
Performance
RecommendationApprove None of these blocks the merge. The overflow and the untested guard only matter for scales that Progress
Nothing was built or run: Python and the compiler need approval in this sandbox. The findings come from reading the code, and the weights, sums and byte counts were worked out by hand. The checkout is shallow, so the diff was taken against |
The uint8 path computed every double weight before converting them, and the temporary memory held them until the kernel returned. It now computes each window twice, once to find the largest weight and once to convert it, so only one window of doubles is kept. Shrinking 1x3x1080x1920 to 224x224 now needs 43,024 bytes of temporary memory instead of 146,496. The float path computes the same weights as before.
ATen adds up a uint8 window in int. With a scale that F.interpolate never gives, one window can hold thousands of weights, and the sum overflows int. For example, resizing a [1, 1, 1, 8000] input that is 255 from column 6579 on to a width of 3 with scales_w 2.28e-4 gave 0 for the last output, where the float result is about 3660 and so 255 after clamping. The sum is now 64-bit, which gives the same result whenever int does not overflow.
The test for weights too large for fixed point used a scale whose weights are also outside the int16 range, so it passed without the precision check. It now uses a scale whose largest weight leaves no fraction bit while every weight still fits in 16 bits. The uint8 test also resizes only the height now, a path no uint8 case ran. The comment on the precision rule said twice the largest weight stays below 2^15; it is the largest weight itself.
HalfTypesMatchFloat32ATen resized kInput to (3, 5) and kOppositeRows to (1, 5) in both float16 and bfloat16. Both cases take the same code path in both types, and only the second one also catches a first pass rounded to a half type, so the first one adds nothing.
uint8 allocates its fixed-point weights on its own path, which no test ran without temporary memory. FailsWithoutTempMemory now checks it too, with a new context so that the earlier failure does not make the check pass on its own.
The kernel capped the number of weights per window at the input size, which ATen does not do. Windows never hold more inputs than there are, so the cap only saved a little memory when shrinking to one to three outputs, and no result depended on it. The weights and results are unchanged.
Every case left now fails for some change to the kernel that no other case catches. Cases that ran the same path as another case are removed: growing in both dims, an output width of 1, NaN inputs, align_corners with a single row, explicit scales on both dims, empty windows in the width alone, the contiguous kept-dim resizes (the channels-last ones remain) and two of the three uint8 sizes. Two cases now also cover code no case reached: - The scale guard test resizes the width too, so a height error that the width's weights overwrote would be caught, and it uses a ratio whose weights fit in memory, so the guard itself is what fails. - The test that compares one resize with two takes a taller input, so that keeping every resized row between the passes would not fit in the fixture's temporary memory.
The program is exported with the first size and run on both. The other three sizes ran the same program path again.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 25m 55s —— View job Approve: none of the issues raised in this round held up when I checked them against the code. Full reviewPR Review: #24400SummaryAdds a portable RecommendationApprove Nothing blocks the merge. Progress
Nothing was built or run, because commands such as The description's temporary-memory figures for 1080x1920 to 224x224 were worked out by hand from the kernel's allocations, and they match: 96,768 B for float32 and 43,024 B for uint8. The checkout is shallow, so the diff was taken against |
to_edgecannot lowerF.interpolate(x, mode="bicubic", antialias=True)today. It exports toaten._upsample_bicubic2d_aa, which is not a core ATen op, soto_edgestops with:ExecuTorch also has no kernel for it.
What this change does
It adds a portable kernel for
_upsample_bicubic2d_aa.outand lists the op as an exception in the edge verifier, as the antialiased bilinear resize does.Like ATen's CPU kernel, it uses the Keys cubic filter (a = -0.5), widened when shrinking, resizes the width, then the height, and copies a dim that keeps its size. float16 and bfloat16, which ATen's CPU kernel lacks, are computed in float32 and rounded once. uint8 follows ATen's fixed-point path: 16-bit integer weights, and each pass rounds and clamps. Other integer types fail in the kernel, as they do in ATen.
Between passes it keeps only the rows one output row reads, so shrinking 1080x1920 to 224x224 needs about 95 KB of temporary memory (43 KB for uint8). It rejects scales
F.interpolatenever produces: too small for the input, or (uint8 only) with weights too large for 16 bits. Where ATen's generic CPU path is wrong (contiguous input, output width 1, changing height; PyTorch issue 200568), it differs from eager on purpose.The shared upsample helper now ignores a non-positive scale, as ATen does; of the other upsample kernels, only antialiased bilinear can pass one. A stale llava comment is updated, and the Arm all-ops test marks the kernel as needing a recipe. Like several portable kernels, it has no experimental CPU backend route yet.
Testing
align_corners, scales, windows past the input, kept dims, both layouts, a batch, all float types, uint8 bit for bit, a dynamic output, running out of temporary memory, and error cases. Each case catches a bug that no other case catches.to_edgetest compares static and dynamic programs (both layouts) with eager in float32 and, bit for bit, in uint8. It fails without the kernel or the verifier exception.