src/wc_avx2.c

   1 /* wc_avx - Count the number of newlines with avx2 instructions.
   2    Copyright (C) 2021-2024 Free Software Foundation, Inc.
   3
   4    This program is free software: you can redistribute it and/or modify
   5    it under the terms of the GNU General Public License as published by
   6    the Free Software Foundation, either version 3 of the License, or
   7    (at your option) any later version.
   8
   9    This program is distributed in the hope that it will be useful,
  10    but WITHOUT ANY WARRANTY; without even the implied warranty of
  11    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12    GNU General Public License for more details.
  13
  14    You should have received a copy of the GNU General Public License
  15    along with this program.  If not, see <https://www.gnu.org/licenses/>.  */
  16
  17 #include <config.h>
  18
  19 #include "wc.h"
  20
  21 #include "system.h"
  22
  23 #include <x86intrin.h>
  24
  25 /* This must be below 16 KB (16384) or else the accumulators can
  26    theoretically overflow, producing wrong result. This is 2*32 bytes below,
  27    so there is no single bytes in the optimal case. */
  28 #define BUFSIZE (16320)
  29
  30 /* Read FD and return a summary.  */
  31 extern struct wc_lines
  32 wc_lines_avx2 (int fd)
  33 {
  34   intmax_t lines = 0;
  35   intmax_t bytes = 0;
  36
  37   __m256i
  38     zeroes = _mm256_setzero_si256 (),
  39     endlines = _mm256_set1_epi8 ('\n');
  40
  41   while (true)
  42     {
  43       /* Using two parallel accumulators gave a good performance increase.
  44          Adding a third gave no additional benefit, at least on an
  45          Intel Xeon E3-1231v3.  Maybe on a newer CPU with additional vector
  46          execution engines it would be a win. */
  47       __m256i
  48         accumulator = _mm256_setzero_si256 (),
  49         accumulator2 = _mm256_setzero_si256 (),
  50         avx_buf[BUFSIZE / sizeof (__m256i)];
  51
  52       ssize_t bytes_read = read (fd, avx_buf, sizeof avx_buf);
  53       if (bytes_read <= 0)
  54         return (struct wc_lines) { bytes_read == 0 ? 0 : errno, lines, bytes };
  55
  56       bytes += bytes_read;
  57       __m256i *datap = avx_buf;
  58
  59       while (bytes_read >= 64)
  60         {
  61           __m256i
  62             to_match = _mm256_load_si256 (datap),
  63             to_match2 = _mm256_load_si256 (datap + 1),
  64             matches = _mm256_cmpeq_epi8 (to_match, endlines),
  65             matches2 = _mm256_cmpeq_epi8 (to_match2, endlines);
  66
  67           /* Compare will set each 8 bit integer in the register to 0xFF
  68              on match.  When we subtract it the 8 bit accumulators
  69              will underflow, so this is equal to adding 1. */
  70           accumulator = _mm256_sub_epi8 (accumulator, matches);
  71           accumulator2 = _mm256_sub_epi8 (accumulator2, matches2);
  72
  73           datap += 2;
  74           bytes_read -= 64;
  75         }
  76
  77       /* Horizontally add all 8 bit integers in the register.  */
  78       accumulator = _mm256_sad_epu8 (accumulator, zeroes);
  79       lines +=   _mm256_extract_epi16 (accumulator, 0)
  80                + _mm256_extract_epi16 (accumulator, 4)
  81                + _mm256_extract_epi16 (accumulator, 8)
  82                + _mm256_extract_epi16 (accumulator, 12);
  83
  84       accumulator2 = _mm256_sad_epu8 (accumulator2, zeroes);
  85       lines +=   _mm256_extract_epi16 (accumulator2, 0)
  86                + _mm256_extract_epi16 (accumulator2, 4)
  87                + _mm256_extract_epi16 (accumulator2, 8)
  88                + _mm256_extract_epi16 (accumulator2, 12);
  89
  90       /* Finish up any left over bytes */
  91       char *end = (char *) datap + bytes_read;
  92       for (char *p = (char *) datap; p < end; p++)
  93         lines += *p == '\n';
  94     }
  95 }