From c68530d8f1574e74c1aa2d3dcbabf0fc6ad13fdf Mon Sep 17 00:00:00 2001
From: Dominique Belhachemi <db@domibel.de>
Date: Thu, 1 Oct 2026 02:05:21 -0400
Subject: [PATCH] Specialise constant-mask swizzle for lsx and lasx

On lsx the only new case is the identity mask, which returns the batch
unchanged. Every other mask takes the runtime path as before.

On lasx there are four. An identity mask again returns the batch
unchanged. The other three exist because xvshuf.b only shuffles within
each 128 bit half of the register, so permuting bytes across the whole
register is the expensive case and worth avoiding. When no element reads
from the other half, a single xvshuf.b does the whole job, with each
index taken modulo the half width so that it addresses its own half.
When every element reads from the low half, that half is first broadcast
to both halves with xvpermi.q 0x00, after which the same single xvshuf.b
applies; a mask reading only the high half is the same with 0x11. Any
other mask takes the runtime path as before.

Those three cases apply only where the element type is one byte. Wider
element types get the identity check and nothing else.
---
 include/xsimd/arch/xsimd_lasx.hpp | 33 ++++++++++++++++++++++++++++++-
 include/xsimd/arch/xsimd_lsx.hpp  |  5 ++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/include/xsimd/arch/xsimd_lasx.hpp b/include/xsimd/arch/xsimd_lasx.hpp
index 55aa982b7..40415d014 100644
--- a/include/xsimd/arch/xsimd_lasx.hpp
+++ b/include/xsimd/arch/xsimd_lasx.hpp
@@ -1106,7 +1106,38 @@ namespace xsimd
         XSIMD_INLINE std::enable_if_t<std::is_arithmetic_v<T>, batch<T, A>>
         swizzle(batch<T, A> const& self, batch_constant<ITy, A, Is...> mask, requires_arch<lasx>) noexcept
         {
-            return swizzle(self, mask.as_batch(), lasx {});
+            if constexpr (detail::is_identity(mask))
+                return self;
+            else if constexpr (sizeof(T) == 1)
+            {
+                constexpr auto lane_mask = mask % std::integral_constant<ITy, (16 / sizeof(T))>();
+                auto const self_int = detail::lasx_to_int(self);
+
+                __m256i src;
+                if constexpr (!detail::is_cross_lane(mask))
+                {
+                    src = self_int;
+                }
+                else if constexpr (detail::is_only_from_lo(mask))
+                {
+                    src = __lasx_xvpermi_q(self_int, self_int, 0x00);
+                }
+                else if constexpr (detail::is_only_from_hi(mask))
+                {
+                    src = __lasx_xvpermi_q(self_int, self_int, 0x11);
+                }
+                else
+                {
+                    return swizzle(self, mask.as_batch(), lasx {});
+                }
+
+                return detail::lasx_from_int<T, A>(
+                    __lasx_xvshuf_b(src, src, detail::lasx_to_int(lane_mask.as_batch())));
+            }
+            else
+            {
+                return swizzle(self, mask.as_batch(), lasx {});
+            }
         }
     }
 }
diff --git a/include/xsimd/arch/xsimd_lsx.hpp b/include/xsimd/arch/xsimd_lsx.hpp
index 598d01e9f..7f1ca30fe 100644
--- a/include/xsimd/arch/xsimd_lsx.hpp
+++ b/include/xsimd/arch/xsimd_lsx.hpp
@@ -999,7 +999,10 @@ namespace xsimd
         XSIMD_INLINE std::enable_if_t<std::is_arithmetic_v<T>, batch<T, A>>
         swizzle(batch<T, A> const& self, batch_constant<ITy, A, Is...> mask, requires_arch<lsx>) noexcept
         {
-            return swizzle(self, mask.as_batch(), lsx {});
+            if constexpr (detail::is_identity(mask))
+                return self;
+            else
+                return swizzle(self, mask.as_batch(), lsx {});
         }
     }
 }