Uh oh!
There was an error while loading. Please reload this page.
[Kernels] Implement portable CPU kernel for aten::bucketize (fixes #20270) - #21869
Open
RanjithRagavan wants to merge 7 commits into
Open
[Kernels] Implement portable CPU kernel for aten::bucketize (fixes #20270)#21869RanjithRagavan wants to merge 7 commits into
RanjithRagavan wants to merge 7 commits into
Conversation
- Implemented portable kernels for bucketize.Tensor_out and bucketize.Scalar_out in kernels/portable/cpu/op_bucketize.cpp. - Supported both 32-bit (int32) and 64-bit (int64) output tensor dtypes via out_int32 argument. - Implemented binary search using std::upper_bound (right=False) and std::lower_bound (right=True) with logarithmic complexity O(N log M). - Added multi-dtype support across floating-point (Float, Double, Half, BFloat16) and integral dtypes, including cross-dtype comparisons and IEEE-754 NaN handling. - Registered operators in functions.yaml, CMakeLists.txt, targets.bzl, and op_registration_util.bzl. - Added comprehensive GTest unit test suite in kernels/test/op_bucketize_test.cpp covering 1D/2D/3D tensors, scalar inputs, empty inputs/boundaries, right flags, and NaN handling.
RanjithRagavan
requested review from
kirklandsign, larryliu0820 and manuelcandales
as code ownersAugust 15, 2026 05:07
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21869
Note: Links to docs will display an error until the docs builds have been completed.
|
RanjithRagavan
commented
Aug 15, 2026
ContributorAuthor
@pytorchbot label "release notes: ops & kernels" |
RanjithRagavan
commented
Aug 15, 2026
ContributorAuthor
@pytorchbot label "module: kernels" |
RanjithRagavan
commented
Aug 19, 2026
ContributorAuthor
@claude review this PR and suggest if any change to address |
ContributorAuthor
@Gasoonjia can you please review and add your feedback |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes#20270
This PR adds portable CPU kernel implementations for
aten::bucketize.Tensor_outandaten::bucketize.Scalar_out.torch.bucketizeis used by model architectures that ExecuTorch targets but currently cannot run on the portable CPU backend:examples/qualcomm/oss_scripts/llama/model/vision_encoder.py(torch.bucketize(frac_h, boundaries, right=True)).extension/llm/modules/turboquant/kv_cache.py(torch.bucketize(rotated, self.boundaries)).Without these kernels, exported models containing
torch.bucketizefail at load time withkernel 'aten::bucketize.Tensor_out' not found(see #20270 for a minimal repro).Implementation details
std::upper_boundforright=False,std::lower_boundforright=True. Complexity is O(N log M), where N is the number of input elements and M the number of boundaries.out_int32:ScalarType::Int(int32) when true,ScalarType::Long(int64) otherwise, validated against theouttensor.Half/BFloat16are supported for both input and boundaries viaET_SWITCH_REALHBF16_TYPES. Mixed-dtype inputs (e.g., integer input with floating-point boundaries) are compared after promotion todouble, matching ATen's type-promotion behavior.boundaries.numel(), empty boundaries produce all zeros, empty inputs are a no-op, and theScalaroverload produces a 0-dim output tensor.Registration
kernels/portable/functions.yaml:bucketize.Tensor_outandbucketize.Scalar_outshim_et/xplat/executorch/kernels/portable/op_registration_util.bzl:op_bucketizetargetkernels/test/CMakeLists.txtandkernels/test/targets.bzlTest plan
New GoogleTest suite
kernels/test/op_bucketize_test.cpp, registered for bothatenandportablemodes so the portable kernel is validated directly against the ATen reference implementation. Coverage:right=True/right=Falsebucket-boundary semantics on 1-D and multi-dimensional inputsint32viaout_int32=True,int64default)rightmodesScalaroverload with float and integer scalars, both output dtypesTest command: