Port the comment block scanning SIMD to Arm. (#3300)

This adds an Arm Neon code path for comment block scanning. It also
tries to tightened up the exact architecture-specific coding pattern for
SIMD code here. Notably, moving to always having architecture-specific
code inside a macro that can be used to globally disable SIMD, but any
place where *all* architectures will require some custom code, an `else`
branch with an error.

Overall, this makes large-block lexing >5% faster on an ARM server
I have access to, but that's a bit misleading. The Neon performance
there doesn't seem very good. On my M1 laptop the difference is *much*
larger. There, even 4-line comments show a noticable improvement and the
block speed looks well over 20%. Sadly, I don't have the same nice
scripts to generate good statistical data.

Raw benchmark data from an ARM sever for reference:

```
BM_CommentLines/1/0/0                          19.4ms ± 7%  19.7ms ± 6%    ~     (p=0.121 n=20+20)
BM_CommentLines/4/0/0                          24.8ms ± 7%  24.8ms ± 6%    ~     (p=0.904 n=20+20)
BM_CommentLines/128/0/0                         251ms ± 2%   234ms ± 3%  -6.70%  (p=0.000 n=20+20)
BM_CommentLines/1/30/0                         21.6ms ±10%  21.9ms ±10%    ~     (p=0.157 n=20+20)
BM_CommentLines/4/30/0                         28.8ms ± 9%  28.8ms ±11%    ~     (p=0.779 n=20+20)
BM_CommentLines/128/30/0                        248ms ± 2%   233ms ± 2%  -5.72%  (p=0.000 n=20+20)
BM_CommentLines/1/70/0                         23.4ms ±12%  23.7ms ±12%    ~     (p=0.341 n=20+20)
BM_CommentLines/4/70/0                         30.8ms ± 9%  31.1ms ±11%    ~     (p=0.602 n=20+20)
BM_CommentLines/128/70/0                        302ms ± 4%   292ms ± 4%  -3.46%  (p=0.000 n=18+20)
BM_CommentLines/1/0/2                          19.8ms ± 7%  20.0ms ± 6%    ~     (p=0.149 n=20+20)
BM_CommentLines/4/0/2                          25.1ms ± 7%  25.3ms ± 8%    ~     (p=0.659 n=20+20)
BM_CommentLines/128/0/2                         225ms ± 2%   212ms ± 2%  -5.88%  (p=0.000 n=20+20)
BM_CommentLines/1/30/2                         22.0ms ± 9%  22.2ms ±10%    ~     (p=0.289 n=20+20)
BM_CommentLines/4/30/2                         29.0ms ±11%  29.1ms ±12%    ~     (p=0.738 n=20+20)
BM_CommentLines/128/30/2                        261ms ±10%   243ms ± 3%  -6.85%  (p=0.000 n=20+20)
BM_CommentLines/1/70/2                         23.5ms ±11%  23.8ms ±15%    ~     (p=0.429 n=20+20)
BM_CommentLines/4/70/2                         31.3ms ±10%  31.5ms ±11%    ~     (p=0.478 n=20+20)
BM_CommentLines/128/70/2                        306ms ± 4%   292ms ± 4%  -4.52%  (p=0.000 n=18+19)
BM_CommentLines/1/0/8                          20.9ms ± 8%  21.2ms ± 7%    ~     (p=0.127 n=20+20)
BM_CommentLines/4/0/8                          27.3ms ± 9%  27.5ms ±12%    ~     (p=0.678 n=20+20)
BM_CommentLines/128/0/8                         227ms ± 2%   210ms ± 2%  -7.35%  (p=0.000 n=19+20)
BM_CommentLines/1/30/8                         22.6ms ±11%  23.0ms ±10%    ~     (p=0.114 n=20+20)
BM_CommentLines/4/30/8                         29.4ms ±10%  29.4ms ±12%    ~     (p=0.947 n=20+20)
BM_CommentLines/128/30/8                        275ms ± 4%   257ms ± 7%  -6.59%  (p=0.000 n=19+20)
BM_CommentLines/1/70/8                         23.9ms ±13%  24.3ms ±14%    ~     (p=0.265 n=20+20)
BM_CommentLines/4/70/8                         32.3ms ±11%  32.4ms ± 9%    ~     (p=0.478 n=20+20)
BM_CommentLines/128/70/8                        319ms ± 4%   307ms ± 4%  -3.83%  (p=0.000 n=18+19)
```

---------

Co-authored-by: Richard Smith <richard@metafoo.co.uk>
Co-authored-by: Jon Ross-Perkins <jperkins@google.com>

toolchain/lex/tokenized_buffer.cpp

@@ -21,8 +21,14 @@
 #include "toolchain/lex/numeric_literal.h"
 #include "toolchain/lex/string_literal.h"
 
-#if __x86_64__
+#if __ARM_NEON
+#include <arm_neon.h>
+#define CARBON_USE_SIMD 1
+#elif __x86_64__
 #include <x86intrin.h>
+#define CARBON_USE_SIMD 1
+#else
+#define CARBON_USE_SIMD 0
 #endif
 
 namespace Carbon::Lex {
@@ -49,21 +55,29 @@ auto VariantMatch(V&& v, Fs&&... fs) -> decltype(auto) {
   return std::visit(Overload{std::forward<Fs&&>(fs)...}, std::forward<V&&>(v));
 }
 
-#if __x86_64__
-#define CARBON_USE_SIMD 1
-
+#if CARBON_USE_SIMD
+namespace {
+#if __ARM_NEON
+using SIMDMaskT = uint8x16_t;
+#elif __x86_64__
+using SIMDMaskT = __m128i;
+#else
+#error "Unsupported SIMD architecture!"
+#endif
+using SIMDMaskArrayT = std::array<SIMDMaskT, sizeof(SIMDMaskT) + 1>;
+}  // namespace
 // A table of masks to include 0-16 bytes of an SSE register.
-// TODO: Make this constexpr to avoid dynamic initialization.
-static const std::array<__m128i, sizeof(__m128i) + 1> prefix_masks = [] {
-  std::array<__m128i, sizeof(__m128i) + 1> masks = {};
-  for (auto [i, mask] : llvm::enumerate(masks)) {
-    memset(&mask, 0xFF, i);
+static constexpr SIMDMaskArrayT PrefixMasks = []() constexpr {
+  SIMDMaskArrayT masks = {};
+  for (int i = 1; i < static_cast<int>(masks.size()); ++i) {
+    // The SIMD types and constexpr require a C-style cast.
+    // NOLINTNEXTLINE(google-readability-casting)
+    masks[i] = (SIMDMaskT)(std::numeric_limits<unsigned __int128>::max() >>
+                           ((sizeof(SIMDMaskT) - i) * 8));
   }
   return masks;
 }();
-#else
-#define CARBON_USE_SIMD 0
-#endif
+#endif  // CARBON_USE_SIMD
 
 // Scans the provided text and returns the prefix `StringRef` of contiguous
 // identifier characters.
@@ -435,10 +449,28 @@ class [[clang::internal_linkage]] TokenizedBuffer::Lexer {
     if (CARBON_USE_SIMD &&
         position + 16 < static_cast<ssize_t>(source_text.size()) &&
         indent <= MaxIndent) {
-#if __x86_64__
       // Load a mask based on the amount of text we want to compare.
-      auto mask = prefix_masks[prefix_size];
-      // And use the current line's prefix as the exemplar to compare against.
+      auto mask = PrefixMasks[prefix_size];
+#if __ARM_NEON
+      // Load and mask the prefix of the current line.
+      auto prefix = vld1q_u8(reinterpret_cast<const uint8_t*>(
+          source_text.data() + first_line_start));
+      prefix = vandq_u8(mask, prefix);
+      do {
+        // Load and mask the next line to consider's prefix.
+        auto next_prefix = vld1q_u8(
+            reinterpret_cast<const uint8_t*>(source_text.data() + position));
+        next_prefix = vandq_u8(mask, next_prefix);
+        // Compare the two prefixes and if any lanes differ, break.
+        auto compare = vceqq_u8(prefix, next_prefix);
+        if (vminvq_u8(compare) == 0) {
+          break;
+        }
+
+        skip_to_next_line();
+      } while (position + 16 < static_cast<ssize_t>(source_text.size()));
+#elif __x86_64__
+      // Use the current line's prefix as the exemplar to compare against.
       // We don't mask here as we will mask when doing the comparison.
       auto prefix = _mm_loadu_si128(reinterpret_cast<const __m128i*>(
           source_text.data() + first_line_start));
@@ -458,14 +490,14 @@ class [[clang::internal_linkage]] TokenizedBuffer::Lexer {
 
         skip_to_next_line();
       } while (position + 16 < static_cast<ssize_t>(source_text.size()));
+#else
+#error "Unsupported SIMD architecture!"
+#endif
       // TODO: If we finish the loop due to the position approaching the end of
       // the buffer we may fail to skip the last line in a comment block that
       // has an invalid initial sequence and thus emit extra diagnostics. We
       // should really fall through to the generic skipping logic, but the code
       // organization will need to change significantly to allow that.
-#elif CARBON_USE_SIMD
-#error Unknown target for SIMD comment skipping.
-#endif
     } else {
       while (position + prefix_size <
                  static_cast<ssize_t>(source_text.size()) &&