Add MAP_CHUNKS2, make parameter name related to input

jkeiser · jkeiser · commit f4fa5b734002 · 2019-08-26T09:46:49.000-07:00
diff --git a/src/arm64/simd_input.h b/src/arm64/simd_input.h
@@ -63,18 +63,32 @@ struct simd_input<Architecture::ARM64> {
     );
   }
 
+  template <typename F>
+  really_inline simd_input<Architecture::ARM64> map(simd_input<Architecture::ARM64> b, F const& map_chunk) {
+    return simd_input<Architecture::ARM64>(
+      map_chunk(this->i0, b.i0),
+      map_chunk(this->i1, b.i1),
+      map_chunk(this->i2, b.i2),
+      map_chunk(this->i3, b.i3)
+    );
+  }
+
   really_inline uint64_t to_bitmask() {
     return neon_movemask_bulk(this->i0, this->i1, this->i2, this->i3);
   }
 
   really_inline uint64_t eq(uint8_t m) {
     const uint8x16_t mask = vmovq_n_u8(m);
-    return this->MAP_BITMASK( vceqq_u8(chunk, mask) );
+    return this->map( [&](auto a) {
+      return vceqq_u8(a, mask);
+    }).to_bitmask();
   }
 
   really_inline uint64_t lteq(uint8_t m) {
     const uint8x16_t mask = vmovq_n_u8(m);
-    return this->MAP_BITMASK( vcleq_u8(chunk, mask) );
+    return this->map( [&](auto a) {
+      return vcleq_u8(a, mask);
+    }).to_bitmask();
   }
 
 }; // struct simd_input
diff --git a/src/arm64/stage1_find_marks.h b/src/arm64/stage1_find_marks.h
@@ -30,7 +30,7 @@ really_inline void find_whitespace_and_structurals(
       (uint8x16_t){8, 0, 18, 4, 0, 1, 0, 1, 0, 0, 0, 3, 2, 1, 0, 0};
   const uint8x16_t low_nib_and_mask = vmovq_n_u8(0xf);
 
-  simd_input<ARCHITECTURE> v = in.map([&](auto chunk) {
+  auto v = in.map([&](auto chunk) {
     uint8x16_t nib_lo = vandq_u8(chunk, low_nib_and_mask);
     uint8x16_t nib_hi = vshrq_n_u8(chunk, 4);
     uint8x16_t shuf_lo = vqtbl1q_u8(low_nibble_mask, nib_lo);
@@ -39,10 +39,10 @@ really_inline void find_whitespace_and_structurals(
   });
 
   const uint8x16_t structural_shufti_mask = vmovq_n_u8(0x7);
-  structurals = v.MAP_BITMASK( vtstq_u8(chunk, structural_shufti_mask) );
+  structurals = MAP_BITMASK( v, vtstq_u8(_v, structural_shufti_mask) );
 
   const uint8x16_t whitespace_shufti_mask = vmovq_n_u8(0x18);
-  whitespace = v.MAP_BITMASK( vtstq_u8(chunk, whitespace_shufti_mask) );
+  whitespace = MAP_BITMASK( v, vtstq_u8(_v, whitespace_shufti_mask) );
 }
 
 #include "generic/stage1_find_marks_flatten.h"
diff --git a/src/haswell/simd_input.h b/src/haswell/simd_input.h
@@ -31,6 +31,14 @@ struct simd_input<Architecture::HASWELL> {
     );
   }
 
+  template <typename F>
+  really_inline simd_input<Architecture::HASWELL> map(simd_input<Architecture::HASWELL> b, F const& map_chunk) {
+    return simd_input<Architecture::HASWELL>(
+      map_chunk(this->lo, b.lo),
+      map_chunk(this->hi, b.hi)
+    );
+  }
+
   really_inline uint64_t to_bitmask() {
     uint64_t r_lo = static_cast<uint32_t>(_mm256_movemask_epi8(this->lo));
     uint64_t r_hi =                       _mm256_movemask_epi8(this->hi);
@@ -39,12 +47,16 @@ struct simd_input<Architecture::HASWELL> {
 
   really_inline uint64_t eq(uint8_t m) {
     const __m256i mask = _mm256_set1_epi8(m);
-    return this->MAP_BITMASK( _mm256_cmpeq_epi8(chunk, mask) );
+    return this->map( [&](auto a) {
+      return _mm256_cmpeq_epi8(a, mask);
+    }).to_bitmask();
   }
 
   really_inline uint64_t lteq(uint8_t m) {
     const __m256i maxval = _mm256_set1_epi8(m);
-    return this->MAP_BITMASK( _mm256_cmpeq_epi8(_mm256_max_epu8(maxval, chunk), maxval) );
+    return this->map( [&](auto a) {
+      return _mm256_cmpeq_epi8(_mm256_max_epu8(maxval, a), maxval);
+    }).to_bitmask();
   }
 
 }; // struct simd_input
diff --git a/src/haswell/stage1_find_marks.h b/src/haswell/stage1_find_marks.h
@@ -62,22 +62,20 @@ really_inline void find_whitespace_and_structurals(simd_input<ARCHITECTURE> in,
     // clang-format off
     const __m256i structural_table =
         _mm256_setr_epi8(44, 125, 0, 0, 0xc0u, 0, 0, 0, 0, 0, 0, 0, 0, 0, 58, 123,
-                          44, 125, 0, 0, 0xc0u, 0, 0, 0, 0, 0, 0, 0, 0, 0, 58, 123);
+                         44, 125, 0, 0, 0xc0u, 0, 0, 0, 0, 0, 0, 0, 0, 0, 58, 123);
     const __m256i white_table = _mm256_setr_epi8(
         32, 100, 100, 100, 17, 100, 113, 2, 100, 9, 10, 112, 100, 13, 100, 100,
         32, 100, 100, 100, 17, 100, 113, 2, 100, 9, 10, 112, 100, 13, 100, 100);
     // clang-format on
     const __m256i struct_offset = _mm256_set1_epi8(0xd4u);
     const __m256i struct_mask = _mm256_set1_epi8(32);
 
-    whitespace = in.MAP_BITMASK( _mm256_cmpeq_epi8(chunk, _mm256_shuffle_epi8(white_table, chunk)) );
-    auto struct_r1 = in.MAP_CHUNKS( _mm256_add_epi8(struct_offset, chunk) );
-    auto struct_r2 = in.MAP_CHUNKS( _mm256_or_si256(chunk, struct_mask) );
-    auto struct_r3 = struct_r1.MAP_CHUNKS( _mm256_shuffle_epi8(structural_table, chunk) );
-    structurals = simd_input<ARCHITECTURE>(
-      _mm256_cmpeq_epi8(struct_r2.lo, struct_r3.lo),
-      _mm256_cmpeq_epi8(struct_r2.hi, struct_r3.hi)
-    ).to_bitmask();
+    whitespace = MAP_BITMASK( in, _mm256_cmpeq_epi8(_in, _mm256_shuffle_epi8(white_table, _in)) );
+
+    auto r1 = MAP_CHUNKS( in, _mm256_add_epi8(struct_offset, _in) );
+    auto r2 = MAP_CHUNKS( in, _mm256_or_si256(_in, struct_mask) );
+    auto r3 = MAP_CHUNKS( r1, _mm256_shuffle_epi8(structural_table, _r1) );
+    structurals = MAP_BITMASK2( r2, r3, _mm256_cmpeq_epi8(_r2, _r3) );
 
   #endif // else SIMDJSON_NAIVE_STRUCTURAL
 }
diff --git a/src/simd_input.h b/src/simd_input.h
@@ -14,6 +14,9 @@ struct simd_input {
     // Map through each simd register in this input, producing another simd_input.
     template <typename F>
     really_inline simd_input<T> map(F const& map_chunk);
+    // Map through each simd register across two inputs, producing a single simd_input.
+    template <typename F>
+    really_inline simd_input<T> map(simd_input<T> b, F const& map_chunk);
     // turn this bytemask (usually the result of a simd comparison operation) into a bitmask.
     uint64_t to_bitmask();
     // a straightforward comparison of a mask against input.
@@ -22,8 +25,10 @@ struct simd_input {
     uint64_t lteq(uint8_t m);
 }; // struct simd_input
 
-#define MAP_CHUNKS(EXPR) map([&](auto chunk) { return EXPR; })
-#define MAP_BITMASK(EXPR) map([&](auto chunk) { return EXPR; }).to_bitmask()
+#define MAP_CHUNKS(A, EXPR) A.map([&](auto _##A) { return (EXPR); })
+#define MAP_BITMASK(A, EXPR) MAP_CHUNKS(A, EXPR).to_bitmask()
+#define MAP_CHUNKS2(A, B, EXPR) A.map((B), [&](auto _##A, auto _##B) { return (EXPR); })
+#define MAP_BITMASK2(A, B, EXPR) MAP_CHUNKS2(A, B, EXPR).to_bitmask()
 
 } // namespace simdjson
 
diff --git a/src/westmere/simd_input.h b/src/westmere/simd_input.h
@@ -40,6 +40,16 @@ struct simd_input<Architecture::WESTMERE> {
     );
   }
 
+  template <typename F>
+  really_inline simd_input<Architecture::WESTMERE> map(simd_input<Architecture::WESTMERE> b, F const& map_chunk) {
+    return simd_input<Architecture::WESTMERE>(
+      map_chunk(this->v0, b.v0),
+      map_chunk(this->v1, b.v1),
+      map_chunk(this->v2, b.v2),
+      map_chunk(this->v3, b.v3)
+    );
+  }
+
   really_inline uint64_t to_bitmask() {
     uint64_t r0 = static_cast<uint32_t>(_mm_movemask_epi8(this->v0));
     uint64_t r1 =                       _mm_movemask_epi8(this->v1);
@@ -50,12 +60,16 @@ struct simd_input<Architecture::WESTMERE> {
 
   really_inline uint64_t eq(uint8_t m) {
     const __m128i mask = _mm_set1_epi8(m);
-    return this->MAP_BITMASK( _mm_cmpeq_epi8(chunk, mask) );
+    return this->map( [&](auto a) {
+      return _mm_cmpeq_epi8(a, mask);
+    }).to_bitmask();
   }
 
   really_inline uint64_t lteq(uint8_t m) {
     const __m128i maxval = _mm_set1_epi8(m);
-    return this->MAP_BITMASK( _mm_cmpeq_epi8(_mm_max_epu8(maxval, chunk), maxval) );
+    return this->map( [&](auto a) {
+      return _mm_cmpeq_epi8(_mm_max_epu8(maxval, a), maxval);
+    }).to_bitmask();
   }
 
 }; // struct simd_input
diff --git a/src/westmere/stage1_find_marks.h b/src/westmere/stage1_find_marks.h
@@ -23,22 +23,17 @@ really_inline void find_whitespace_and_structurals(simd_input<ARCHITECTURE> in,
 
   const __m128i structural_table =
       _mm_setr_epi8(44, 125, 0, 0, 0xc0u, 0, 0, 0, 0, 0, 0, 0, 0, 0, 58, 123);
-  const __m128i white_table = _mm_setr_epi8(32, 100, 100, 100, 17, 100, 113, 2,
-                                              100, 9, 10, 112, 100, 13, 100, 100);
+  const __m128i white_table = _mm_setr_epi8(32, 100, 100, 100,  17, 100, 113,   2,
+                                           100,   9,  10, 112, 100,  13, 100, 100);
   const __m128i struct_offset = _mm_set1_epi8(0xd4u);
   const __m128i struct_mask = _mm_set1_epi8(32);
 
-  whitespace = in.MAP_BITMASK( _mm_cmpeq_epi8(chunk, _mm_shuffle_epi8(white_table, chunk)) );
-
-  auto r1 = in.MAP_CHUNKS( _mm_add_epi8(struct_offset, chunk) );
-  auto r2 = in.MAP_CHUNKS( _mm_or_si128(chunk, struct_mask) );
-  auto r3 = r1.MAP_CHUNKS( _mm_shuffle_epi8(structural_table, chunk) );
-  structurals = simd_input<ARCHITECTURE>(
-    _mm_cmpeq_epi8(r2.v0, r3.v0),
-    _mm_cmpeq_epi8(r2.v1, r3.v1),
-    _mm_cmpeq_epi8(r2.v2, r3.v2),
-    _mm_cmpeq_epi8(r2.v3, r3.v3)
-  ).to_bitmask();
+  whitespace = MAP_BITMASK( in, _mm_cmpeq_epi8(_in, _mm_shuffle_epi8(white_table, _in)) );
+
+  auto r1 = MAP_CHUNKS( in, _mm_add_epi8(struct_offset, _in) );
+  auto r2 = MAP_CHUNKS( in, _mm_or_si128(_in, struct_mask) );
+  auto r3 = MAP_CHUNKS( r1, _mm_shuffle_epi8(structural_table, _r1) );
+  structurals = MAP_BITMASK2( r2, r3, _mm_cmpeq_epi8(_r2, _r3) );
 }
 
 #include "generic/stage1_find_marks_flatten.h"
