[libc-commits] [libc] [libc] Fix AVX2 compare_and_mask broadcasting byte needle with epi16 (PR #235246)

[via libc-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libc

Author: Aman Maurya (amanmaurya92)

<details>
<summary>Changes</summary>

In `libc/src/string/memory_utils/x86_64/inline_strlen.h`, the AVX2 specialization of `compare_and_mask` used `_mm256_set1_epi16(c)` instead of `_mm256_set1_epi8(c)`.

Because `compare_and_mask` is used by `find_first_character` (and therefore `memchr`), this caused:
1. Matches at odd offsets to be missed (since odd bytes in the vector mask were 0x00 instead of `c`).
2. False positive matches when 0x00 (or 0xFF) occurred at odd offsets while searching for a nonzero character.

This change fixes the broadcast intrinsic to `_mm256_set1_epi8(c)` (matching the SSE2 and AVX-512 specializations) and adds regression tests to `memchr_test`.

Fixes #<!-- -->226787.

---
Full diff: https://github.com/llvm/llvm-project/pull/235246.diff


2 Files Affected:

- (modified) libc/src/string/memory_utils/x86_64/inline_strlen.h (+1-1) 
- (modified) libc/test/src/string/memchr_test.cpp (+24) 


``````````diff
diff --git a/libc/src/string/memory_utils/x86_64/inline_strlen.h b/libc/src/string/memory_utils/x86_64/inline_strlen.h
index 0343e22f3e58e3..3c31d38706816a 100644
--- a/libc/src/string/memory_utils/x86_64/inline_strlen.h
+++ b/libc/src/string/memory_utils/x86_64/inline_strlen.h
@@ -110,7 +110,7 @@ template <>
 LIBC_INLINE uint32_t
 compare_and_mask<__m256i, uint32_t>(const __m256i *block_ptr, char c) {
   __m256i b = _mm256_load_si256(block_ptr);
-  __m256i set = _mm256_set1_epi16(c);
+  __m256i set = _mm256_set1_epi8(c);
   __m256i cmp = _mm256_cmpeq_epi8(b, set);
   return _mm256_movemask_epi8(cmp);
 }
diff --git a/libc/test/src/string/memchr_test.cpp b/libc/test/src/string/memchr_test.cpp
index a92c5fe80be98b..9a3f21e47ba989 100644
--- a/libc/test/src/string/memchr_test.cpp
+++ b/libc/test/src/string/memchr_test.cpp
@@ -131,6 +131,30 @@ TEST(LlvmLibcMemChrTest, SignedCharacterFound) {
   ASSERT_EQ(actual[0], c);
 }
 
+TEST(LlvmLibcMemChrTest, MatchesAtOddAndEvenOffsets) {
+  alignas(32) unsigned char buf[64];
+  for (size_t i = 0; i < sizeof(buf); ++i)
+    buf[i] = 'x';
+
+  for (size_t offset = 0; offset < sizeof(buf); ++offset) {
+    buf[offset] = 'A';
+    const char *actual = call_memchr(buf, 'A', sizeof(buf));
+    ASSERT_EQ(actual, reinterpret_cast<const char *>(buf + offset));
+    buf[offset] = 'x';
+  }
+}
+
+TEST(LlvmLibcMemChrTest, NoFalseMatchWithZeroOrNegativeByteAtOddOffset) {
+  alignas(32) unsigned char buf[64];
+  for (size_t i = 0; i < sizeof(buf); ++i)
+    buf[i] = 0x11;
+
+  // Odd offset has 0x00 and 0xFF; searching for 'A' (0x41) must not false-match.
+  buf[1] = 0x00;
+  buf[3] = 0xFF;
+  ASSERT_STREQ(call_memchr(buf, 'A', sizeof(buf)), nullptr);
+}
+
 #if defined(LIBC_ADD_NULL_CHECKS)
 
 TEST(LlvmLibcMemChrTest, CrashOnNullPtr) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/235246