Vulkan: implement any.dim for boolean tensors - #22790
Open
msluszniak wants to merge 1 commit into
Open
Conversation
`aten.any.dim` has no Vulkan implementation, so a boolean reduction is a graph break. Attention masks hit this hard: an RF-DETR segmentation export reduces a [1, 6, 677, 677] boolean tensor 16 times, and each one copies 2.75M elements out to the CPU and back to produce 4k of output. Over a boolean tensor `any` is a max over the uint8 representation, so the general reduce shader already covers it once it can be instantiated for uint8. Two things were in the way: - The accumulator path hardcoded `vec4` at the load sites, which does not compile when `load_texel` returns `uvec4`. Converting at the load instead keeps the accumulator and the INIT/UPDATE/POSTPROCESS macros in float, so they continue to work on both vec4 and scalar arguments, and the write converts back to the tensor's texel type. This is a no-op for the existing float and half variants. - `any.dim` takes a single int dim rather than the int list the other reduce ops take, so it cannot share DEFINE_REDUCE_FN. The partitioner claims only the general (texture) implementation, since the buffer per-row shaders have no boolean variant, and only keepdim=True, which is what the existing support predicate already requires. Verified on a Mali-G76 (Galaxy S10+) with both branches exercised: the result is bit-identical to eager. On RF-DETR, combined with mul.Scalar, this takes the partition count from 39 delegates to 7 and CPU kernel calls from 55 to 7, with every output bit-identical to the 39-delegate build.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22790
Note: Links to docs will display an error until the docs builds have been completed.
|
pytorch-bot
Bot
added
the
module: vulkan
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
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 free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
aten.any.dimhas no Vulkan implementation, so a boolean reduction is a graph break. Attention masks hit this hard: an RF-DETR segmentation export reduces a[1, 6, 677, 677]boolean tensor 16 times, and each one copies 2.75M elements out to the CPU and back to produce 4k of output.Over a boolean tensor
anyis a max over the uint8 representation, so the general reduce shader already covers it once it can be instantiated for uint8. Two things were in the way:vec4at the load sites, which does not compile whenload_texelreturnsuvec4. Converting at the load instead keeps the accumulator and the INIT/UPDATE/POSTPROCESS macros in float, so they keep working on both vec4 and scalar arguments, and the write converts back to the tensor's texel type. No-op for the existing float and half variants.any.dimtakes a single int dim rather than the int list the other reduce ops take, so it cannot shareDEFINE_REDUCE_FN.The partitioner claims only the general (texture) implementation, since the buffer per-row shaders have no boolean variant, and only
keepdim=True, which the existing support predicate already requires.Verified on a Mali-G76 (Galaxy S10+) with both branches exercised (rows planted so the result is neither all-True nor all-False): bit-identical to eager.
Stacked with #22789 (
mul.Scalar), RF-DETR goes from 39 delegates to 7, CPU kernel calls from 55 to 7, and every output stays bit-identical to the 39-delegate build. End to end on the same device, warmed and interleaved over 4 order-reversed rounds: 1078 ms -> 856 ms, against 2474 ms for XNNPACK fp32.cc @SS-JIA @manuelcandales @digantdesai @cbilgin