Proposal load/store masked by DiamonDinoia · Pull Request #1162 · xtensor-stack/xsimd

DiamonDinoia · 2025-08-21T15:49:24Z

I made a rough implementation of masked load/store. Before I hash it out. Can I have some early feedback?

Thanks,
Marco

serge-sans-paille · 2025-08-21T19:49:43Z

interesting. Before going into the details, what are the memory effects of a masked load in terms of read memory and value stored? Are the masked elements set to 0 ? to an undefined value ? AVX etc seems to set the value to zero?

I'm not sold on the common implementation which looks quite heavy in scalar operation. I can see that we can't do a plain load followed by an and because it could lead to access to unallocated memory. If the mask were constant, we could optimize statically some common patterns, but with a dynamic mask as you propose...

DiamonDinoia · 2025-08-22T08:58:25Z

Some thoughts:

Undefined for masked values. Since depending on the operations 0 or 1 might be the correct values. In that case they could use the mask itself to initialize the values. Also, because imagine I want to polulate the even elements from one memory location and the odd from another. Masked loads (I think) are faster than a gather. set to 0 following the x86 convention
We could remove the dynamic mask entirely. I added for completeness.
We could do a la vcl and have a load partial, store partial where we just optimize for head and tail. I preferred this solution as I'm assuming xsimd users know the performance implications of the API.
~~For now this is fast only on avx, av2 but for sse even if it heavy on scalar it is slow only when reading bytes or short. We could optimize these cases.~~
I'm not sure about sve/neon
With static masks if the first and the last element are read it is possible to do load+and

In general, I use these operations when I want to vectorize and inner loop that is not a multiple of the simd width. This is a small inner loop nested in a loop executed a lot of times. Depending on the operations, padding sometimes is slower than masking.

DiamonDinoia · 2025-10-08T20:52:07Z

Hi @serge-sans-paille,

This is starting to take shape now. I still need to clean up. But this is a first working implementation.

I extended the mask operation in batch_bool to mimic std::bit so that i could reuse the code for all masked operations. I feel having those living outside the class in detail:: is not the best. They are also useful to the user I think.

DiamonDinoia · 2025-10-10T17:58:55Z

@serge-sans-paille do you have any clue on why mingw32 fails?

I need to do a final pass make the code coherent (I am mixing reinterpret_casts and c-style casts). Affter that I think will ready for review.

serge-sans-paille · 2025-10-12T05:58:33Z

+            }
+            else
+            {
+                // Fallback to runtime path for non prefix/suffix masks


Unless the compiler optimizes those out, the dynamic load is going to perform serveral checks that are going to be always false (none / all).
I'm still not convinced by the relevancy of dynamic mask support. Maybe we could start small with "just" the static mask support and eventually add a dynamic version later if it appears to be needed? But maybe you have specific case in mind ?

It is just that in my mind I can offer a public runtime mask API
then the compile time mask optimizes as much as it can and if there is no alternative it calls the runtime mask.

Same as the swizzle. So I was not planning of optimizing the runtime at all we just get a maskl_load/store where possible or a common since sse does not support masking.

PS:
This can be optimized in the future with using avx on sse or avx512vl but I will open a discussion once this is merged

I am cleaning this to only offer compile time for now and the runtime can be a separate PR

DiamonDinoia · 2025-10-19T04:03:10Z

I don't think the error is due to anything related to this PR. But, I like the looks of it now. Ready to review!

serge-sans-paille · 2025-10-28T09:30:42Z

@DiamonDinoia we plan to do a release in the forthcoming weeks. Do you want this PR to be part of it?

DiamonDinoia · 2025-10-28T14:28:11Z

@DiamonDinoia we plan to do a release in the forthcoming weeks. Do you want this PR to be part of it?

Only if it does not delay the release by much. I would like to have a release by mid-November. This gives me two weeks to update my downstream projects. For an end of the year release.

PS: I am planning on addressing the swizzle issue after this PR.

JohanMabille · 2025-10-30T09:52:00Z

I would like to have a release by mid-November

That is more or less what we planned ;)

serge-sans-paille

A new wave of comments. I have the feeling that we could simplify the implementation a lot by making it more generic.

serge-sans-paille · 2025-10-30T21:04:16Z

+                return load<A>(mem, Mode {});
+            }
+            // confined to lower 128-bit half (4 lanes) → forward to SSE
+            else XSIMD_IF_CONSTEXPR(mask.countl_zero() >= 4)


Thinking of it, we could remove the need for countl_zero by just masking, something like

if((mask.mask() & 0xF0) == 0)

which would be a decent implementation for

if(mask.lower_half() == 0) // or lower_mask ?

which has the nice property of not encoding any "magical value"

The reason why I went for countl_zero is to maintain parity with std::bit. I feel like this is public API so it is worth for users not having to learn a new naming scheme

serge-sans-paille · 2025-10-30T21:16:15Z

+            else XSIMD_IF_CONSTEXPR(mask.countr_zero() >= 2)
+            {
+                constexpr auto mhi = ::xsimd::detail::upper_half<sse2>(mask);
+                const batch<double, sse2> hi(_mm256_extractf128_pd(src, 1));


that's a upper_half, right?

serge-sans-paille · 2025-10-30T21:16:15Z

+            else XSIMD_IF_CONSTEXPR(mask.countl_zero() >= 2)
+            {
+                constexpr auto mlo = ::xsimd::detail::lower_half<sse2>(mask);
+                const batch<double, sse2> lo(_mm256_castpd256_pd128(src));


that's a lower_half, right?

serge-sans-paille · 2025-10-30T21:18:46Z

+            {
+                _mm256_maskstore_pd(mem, mask.as_batch(), src);
+            }
+        }


epiphany : if you express everything in terms of lower / upper half, you're pretty close to a type-generic implementation that would avoid a lot of redunduncy. just need to wrap the actual _mm256_maskstore* and we're good, right?

I guess you are right. It might be possible to express things in terms of size and halves. This might be recursive though and we might need to arch to half arch, similarly to what happens in widen.

I'll also have to think if the size of the element matters. I'll have a look next week

The reason why I kept the implementation separate is as follows:

The general implementation in terms of upper/lower ais possible. This haves the code in avx, avx2 still has to provide a new implementation so that it can take advantage of the new intrinsics so we for from 2 versions in avx down to 1 and 2 in avx2 down to 1. Same for avx512.

But then for merging we can not use more optimized intrinsics than _mm256_castps128_ps256. Is it worth it?

The solution here is to pass the MaskLoad/Store and the merge function as templates to this general implementation. Then we could have 1 general implementation maybe for all the kernels. But, I think it might be too much templating.

In general also wrapping _mm256_maskstore* also makes a clean way to expose this to the used if we wish to do so.

What do you suggest?

serge-sans-paille · 2025-10-30T21:19:54Z

            return _mm512_loadu_pd(mem);
        }

+        // load_masked


why did you split the load_masked in two parts of the header? The generic epiphany for avx seems to apply here too :-)

serge-sans-paille · 2025-10-31T10:21:04Z

+            {
+                constexpr auto mlo = ::xsimd::detail::lower_half<sse4_2>(mask);
+                const auto lo = load_masked(mem, mlo, convert<float> {}, Mode {}, sse4_2 {});
+                return batch<float, A>(detail::merge_sse(lo, batch<float, sse4_2>(0.f)));


could be a call to _mm256_castps128_ps256 instead of merge + batch of zero

serge-sans-paille · 2025-10-31T10:21:41Z

+            {
+                constexpr auto mhi = ::xsimd::detail::upper_half<sse4_2>(mask);
+                const auto hi = load_masked(mem + 4, mhi, convert<float> {}, Mode {}, sse4_2 {});
+                return batch<float, A>(detail::merge_sse(batch<float, sse4_2>(0.f), hi));


I don't know if we could use a similar trick to _mm256_castps128_ps256 here

Related to #1184 and #1162

serge-sans-paille · 2025-11-03T21:23:38Z

@AntoinePrv once you update that branch we can do a release \o/
EDIT: I meant @DiamonDinoia of course :-)

DiamonDinoia · 2025-11-03T21:26:26Z

@AntoinePrv once you update that branch we can do a release \o/

I will work on this asap. I just have something to deliver on another project this week so it might take a bit for me to get back to this.

serge-sans-paille · 2025-11-03T22:30:15Z

Note that once you're done with Intel support, i still have to implement ARM et cie when possible

DiamonDinoia · 2025-11-14T01:21:56Z

I removed most of the duplication at least I stopped when I thought things were getting out of hand. I think in the future we might offer direct access to masked load with dynamic mask. Which is basically exposing the internal helper functions but adapting them to the xsimd convention.

1. Adds new masked API compile time masks (store_masked and load_masked) 2. General use case optimization 3. New tests 4. x86 kernels 5. Adds new APIs to batch_bool_constant for convenience resembling #include<bit> 6. Tests the new APIs

serge-sans-paille · 2025-11-18T08:09:28Z

Merged as 239fbbd
Thanks a lot @DiamonDinoia for the hard work. It's now on me to implement part of it for other architectures.

AntoinePrv · 2025-11-19T10:01:48Z

+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE batch<int32_t, A> load_masked(int32_t const* mem, batch_bool_constant<int32_t, A, Values...> mask, convert<int32_t>, Mode, requires_arch<avx2>) noexcept
+        {
+            return load_masked<A>(mem, mask, convert<int32_t> {}, Mode {}, avx2 {});
+        }


@DiamonDinoia I am having trouble with this function above. Is this an infinite loop?

I might have meant calling uint -> int32 there.

I am currently travelling. I have a look once I am back.

DiamonDinoia force-pushed the masked-memory-ops branch from 1124f01 to 53da643 Compare August 21, 2025 15:51

DiamonDinoia force-pushed the masked-memory-ops branch from 2ae15bf to 510c335 Compare October 8, 2025 20:49

DiamonDinoia force-pushed the masked-memory-ops branch 8 times, most recently from fd0e777 to ecf4291 Compare October 9, 2025 23:27

DiamonDinoia force-pushed the masked-memory-ops branch from ecf4291 to e88369f Compare October 10, 2025 21:35

serge-sans-paille reviewed Oct 12, 2025

View reviewed changes

DiamonDinoia force-pushed the masked-memory-ops branch 11 times, most recently from e535375 to 4c457ae Compare October 19, 2025 02:18

DiamonDinoia marked this pull request as ready for review October 19, 2025 04:02

DiamonDinoia force-pushed the masked-memory-ops branch from 4c457ae to ec2b824 Compare October 20, 2025 03:27

DiamonDinoia mentioned this pull request Oct 30, 2025

Extend support of batch_cast<...> to upcasting to a type twice as big #1184

Merged

serge-sans-paille reviewed Oct 30, 2025

View reviewed changes

serge-sans-paille reviewed Oct 31, 2025

View reviewed changes

serge-sans-paille added a commit that referenced this pull request Oct 31, 2025

Move from split_avx / split_avx512 to lower_half/ upper_half

f4d7a34

Related to #1184 and #1162

serge-sans-paille mentioned this pull request Oct 31, 2025

Move from split_avx / split_avx512 to lower_half/ upper_half #1188

Merged

serge-sans-paille added a commit that referenced this pull request Oct 31, 2025

Move from split_avx / split_avx512 to lower_half/ upper_half

fce8da2

Related to #1184 and #1162

DiamonDinoia force-pushed the masked-memory-ops branch from 3f3acf5 to 8104455 Compare November 5, 2025 20:56

DiamonDinoia force-pushed the masked-memory-ops branch 7 times, most recently from 856d98e to 9185bb2 Compare November 13, 2025 22:46

DiamonDinoia requested a review from serge-sans-paille November 14, 2025 01:22

Feat: Load/Store masked API

fb16c52

1. Adds new masked API compile time masks (store_masked and load_masked) 2. General use case optimization 3. New tests 4. x86 kernels 5. Adds new APIs to batch_bool_constant for convenience resembling #include<bit> 6. Tests the new APIs

DiamonDinoia force-pushed the masked-memory-ops branch from 9185bb2 to fb16c52 Compare November 14, 2025 19:43

DiamonDinoia mentioned this pull request Nov 17, 2025

Add batch_constant comparison and shift #1208

Merged

serge-sans-paille closed this Nov 18, 2025

AntoinePrv reviewed Nov 19, 2025

View reviewed changes

AntoinePrv mentioned this pull request Nov 19, 2025

Proposal for Pixi workflow #1212

Open

4 tasks

Conversation

DiamonDinoia commented Aug 21, 2025

Uh oh!

serge-sans-paille commented Aug 21, 2025

Uh oh!

DiamonDinoia commented Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DiamonDinoia commented Oct 8, 2025

Uh oh!

DiamonDinoia commented Oct 10, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DiamonDinoia commented Oct 19, 2025

Uh oh!

serge-sans-paille commented Oct 28, 2025

Uh oh!

DiamonDinoia commented Oct 28, 2025

Uh oh!

JohanMabille commented Oct 30, 2025

Uh oh!

serge-sans-paille left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DiamonDinoia Nov 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serge-sans-paille Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serge-sans-paille Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serge-sans-paille commented Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DiamonDinoia commented Nov 3, 2025

Uh oh!

serge-sans-paille commented Nov 3, 2025

Uh oh!

DiamonDinoia commented Nov 14, 2025

Uh oh!

serge-sans-paille commented Nov 18, 2025

Uh oh!

AntoinePrv Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

DiamonDinoia commented Aug 22, 2025 •

edited

Loading

DiamonDinoia Nov 5, 2025 •

edited

Loading

serge-sans-paille Oct 31, 2025 •

edited

Loading

serge-sans-paille Oct 31, 2025 •

edited

Loading

serge-sans-paille commented Nov 3, 2025 •

edited

Loading

AntoinePrv Nov 19, 2025 •

edited

Loading