Add vectorization in elementwise_util #9432
Draft
Conversation
This file contains 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
swolchok
requested review from
JacobSzwejbka,
lucylq and
manuelcandales
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/9432
Note: Links to docs will display an error until the docs builds have been completed. ❌ 70 New FailuresAs of commit 0beabbb with merge base 1572381 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Mar 20, 2025
Merged
swolchok
added a commit
that referenced
this pull request
Mar 20, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 30d2311bed080c3a5390ab00ca20a1e33563f077 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
facebook-github-bot
added
the
CLA Signed
swolchok
added a commit
that referenced
this pull request
Mar 20, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: d546d5d595929e84814aa38833c8a07bf3cf6ec5 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
This was referenced Mar 25, 2025
swolchok
changed the title
|
This is draft status because
I have confirmed from local measurements with op_mul that this does in fact cause vectorization and it is a significant perf win (matches optimized op_mul for the treat-as-1d case, not too surprisingly). |
|
I cant find the pr where we were discussing using vecotrized class for scalar compute, but if we fix the issue in vec_base.h that you highlighted, that is, https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/cpu/vec/vec_base.h#L157, size is > 1?, would that make a better solution? Then we can just use Vectorized class everywhere. |
We can't use Vectorized everywhere unless we want to commit to keeping the ExecuTorch copy/paste of at::vec::Vectorized in sync and up to date. I would very much like to delete it instead. It is also probably going to end up being worse code to use it everywhere, though I admit that the cost of not having to write Vectorized everywhere may end up being too high. |
Not suggesting to copy/paste it everywhere, but I am not sure what you mean,
Worse from readability perspective? |
https://github.com/pytorch/executorch/tree/main/kernels/optimized/vec is a (partial, out-of-sync) copy/paste of https://github.com/pytorch/pytorch/tree/main/aten/src/ATen/cpu/vec . I would like to delete the ExecuTorch version, not further cement its use.
Yes, and from being able to cleanly control when vectorization is used. It's not necessary to pollute everything with Vectorized; this diff shows how to control Vectorization centrally and I am working on making the rollout to operators clean as well. |
where do you expect we would have to say not want to vectorize something? |
we should not bother generating vectorized loops at least when |
This was referenced Apr 2, 2025
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 8d76653f819dc58a0c93540f3d71a89bfdb7cd26 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
|
This PR is not yet complete (I need to go through and make sure no ops are needlessly held back from vectorization, and I need to overload unary - on at::vec::Vectorized so that op_sigmoid.cpp can vectorize nicely), but it builds and should vectorize a large chunk of elementwise portable ops. Hopefully this makes the direction/vision more clear @kimishpatel (CI failures are expected; the builds are contingent on pytorch/pytorch#150380 which is not yet actually merged into PyTorch core. We'll need a pin bump.) |
swolchok
commented
Apr 2, 2025
Comment on lines
+50
to
52
| [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b, const CTYPE_COMPUTE val_c) { | ||
| return val_c ? val_a : val_b; | ||
| }, |
There was a problem hiding this comment.
(ATen, and therefore our own optimized op_where, doesn't vectorize this)
I am not fully sold but I havent looked at the PR in detail so I will further comment after that. Maybe you are right in that not everything will fit the Vectorized pattern and for those cases your pattern of checking callable for vectorized is probably a better way to do it. |
This was referenced Apr 2, 2025
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 8d76653f819dc58a0c93540f3d71a89bfdb7cd26 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
This was referenced Apr 2, 2025
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 033b63ce3bee8a0136efdab3e03905cafb79b915 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 033b63ce3bee8a0136efdab3e03905cafb79b915 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
kimishpatel
reviewed
Apr 2, 2025
| template <typename T> \ | ||
| auto func_name(at::vec::Vectorized<T> vec) { \ | ||
| if constexpr (!::executorch::runtime::is_floating_point<T>::value) { \ | ||
| return at::vec::convert<float>(vec).func_name(); \ |
There was a problem hiding this comment.
is this a valid thing to do? that is convert say an instance of at::vec::Vectorized<int8_t> to float and apply the func_name? Maybe i am misunderstanding how this works
| */ | ||
| #define ET_INTERNAL_VECTORIZED_FLOAT_UNARY_FUNC(func_name) \ | ||
| namespace executorch { \ | ||
| inline namespace math { \ |
There was a problem hiding this comment.
what does inline to namespace do?
Comment on lines
+62
to
+63
| * corresponding operator is a "float op" in TensorIterator parlance | ||
| * (i.e., uses something like build_borrowing_binary_float_op()), |
There was a problem hiding this comment.
I dont know if anyone reading this code would understand what this means?
There was a problem hiding this comment.
But this provides answer to my earlier question
| @@ -47,7 +47,7 @@ Tensor& where_out( | |||
| CTYPE_COMPUTE, | |||
| op_name, | |||
| utils::SupportedTensorDtypes::SAME_AS_COMMON>( | |||
| [](const auto val_a, const auto val_b, const auto val_c) { | |||
| [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b, const CTYPE_COMPUTE val_c) { | |||
There was a problem hiding this comment.
is this unrelated change?
There was a problem hiding this comment.
no, we can't vectorize this op
| @@ -51,6 +56,34 @@ inline int64_t scalar_to<int64_t>(const Scalar& s) { | |||
| } | |||
|
|
|||
| namespace internal { | |||
| template <typename Ignore, typename T> | |||
| using ignore_first_yield_second = T; | |||
| if constexpr (std::is_invocable_v< | ||
| Op, | ||
| ignore_first_yield_second<Args, Vec>...>) { | ||
| // For bool, we will get a false positive if we rely on only the |
There was a problem hiding this comment.
you mean bool return type?
There was a problem hiding this comment.
bool ctype_compute
Comment on lines
+103
to
+105
| ((inputs.first->scalar_type() == | ||
| CppTypeToScalarType<CTYPE_COMPUTE>::value) && | ||
| ...)); |
There was a problem hiding this comment.
Man, this is a good reminder of all the template meta programming magic
| ...); | ||
| if (!any_is_broadcasted) { | ||
| using Vec = at::vec::Vectorized<CTYPE_COMPUTE>; | ||
| ::executorch::extension::parallel_for( |
There was a problem hiding this comment.
I think doing this blindly for each op is a bit risky in that, no all multithreading is always better. some ops benefit from smaller grain size while others with larger
There was a problem hiding this comment.
IIRC xnnpack blindly parallelizes absolutely everything; we're doing strictly better here
There was a problem hiding this comment.
yeah I am not comparing with xnnpack. In fact one bit part of the reason why we ended up leveraging optimized op lib for some of the llama stuff for exactly that reason. That it blindly parallelized everything and that actually hurt perf
Comment on lines
+125
to
+126
| const auto vectorized_begin = | ||
| begin + (Vec::size() - begin % Vec::size()) % Vec::size(); |
There was a problem hiding this comment.
Feels like something that has chances of bug. Hope we test this enough. I would doubt if our test cases will exercise this code path.
There was a problem hiding this comment.
although I do see you treat scalar left overs of both head and tails separately
| inputs.first->sizes(), out.sizes()) && | ||
| ...); | ||
| if (!any_is_broadcasted) { | ||
| using Vec = at::vec::Vectorized<CTYPE_COMPUTE>; |
There was a problem hiding this comment.
So the one point of contention for me is that why do we need vectorized_math.h which largely is doing trampoline to underlying vectorized methods. Mainly you dont even need to use can_use_vectorized, because on non accelerated platforms Vectorized falls back to scalar impl even if Vec::size() != `. Maybe you said that the generated code would be worse if forced Vectorized, but I am not sure why. Rest makes sense.
However, place where I can potentially see that being useful is for dtype that Vectorized doesnt support but for float I am not sure. So maybe if you can clarify that it would help.
There was a problem hiding this comment.
why do we need vectorized_math.h which largely is doing trampoline to underlying vectorized methods
without it, you can't take the same lambda you already wrote for scalars and reuse it for Vectorized (the change isn't zero because you have to point at executorch::math, but crucially it doesn't require separate code)
|
Ok we synced offline. So the major value prop as I understood: You can write your lambdas without having to explicitly use Vectorized. Why explicitly using Vectorized might not be good? Because maybe Vectorized does not have everything you need to implement your lambda. So as op author you dont have to worry about Vectorized when writing your lambda (although do have to use executorch::math ops). And later if Vectorized added support for the ::math.. ops you get vectorization for free without having to go back and rewrite your lambda with Vectorized. So this is nice. |
Sign up for free
to 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.
This is a first cut at #9241 . In this PR I've vectorized op_mul to make sure that vectorization doesn't break tests; a follow-up PR will make all existing portable ops vectorized-capable. I've left covering ops that use the
unary_ufunc_*utilities in pattern.h for a follow-up push, because pattern.h and elementwise_util need some work before we can migrate pattern.h's utilities to be backed by elementwise_util.