src/wc_avx2.c

   1 /* wc_avx - Count the number of newlines with avx2 instructions.
   2    Copyright (C) 2021-2023 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 "system.h"
  20 #include "error.h"
  21 #include "safe-read.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 extern bool
  31 wc_lines_avx2 (char const *file, int fd, uintmax_t *lines_out,
  32                uintmax_t *bytes_out);
  33
  34 extern bool
  35 wc_lines_avx2 (char const *file, int fd, uintmax_t *lines_out,
  36                uintmax_t *bytes_out)
  37 {
  38   __m256i accumulator;
  39   __m256i accumulator2;
  40   __m256i zeroes;
  41   __m256i endlines;
  42   __m256i avx_buf[BUFSIZE / sizeof (__m256i)];
  43   __m256i *datap;
  44   uintmax_t lines = 0;
  45   uintmax_t bytes = 0;
  46   size_t bytes_read = 0;
  47
  48
  49   if (!lines_out || !bytes_out)
  50     return false;
  51
  52   /* Using two parallel accumulators gave a good performance increase.
  53      Adding a third gave no additional benefit, at least on an
  54      Intel Xeon E3-1231v3.  Maybe on a newer CPU with additional vector
  55      execution engines it would be a win. */
  56   accumulator = _mm256_setzero_si256 ();
  57   accumulator2 = _mm256_setzero_si256 ();
  58   zeroes = _mm256_setzero_si256 ();
  59   endlines = _mm256_set1_epi8 ('\n');
  60
  61   while ((bytes_read = safe_read (fd, avx_buf, sizeof (avx_buf))) > 0)
  62     {
  63       __m256i to_match;
  64       __m256i to_match2;
  65       __m256i matches;
  66       __m256i matches2;
  67
  68       if (bytes_read == SAFE_READ_ERROR)
  69         {
  70           error (0, errno, "%s", quotef (file));
  71           return false;
  72         }
  73
  74       bytes += bytes_read;
  75
  76       datap = avx_buf;
  77       char *end = ((char *)avx_buf) + bytes_read;
  78
  79       while (bytes_read >= 64)
  80         {
  81           to_match = _mm256_load_si256 (datap);
  82           to_match2 = _mm256_load_si256 (datap + 1);
  83
  84           matches = _mm256_cmpeq_epi8 (to_match, endlines);
  85           matches2 = _mm256_cmpeq_epi8 (to_match2, endlines);
  86           /* Compare will set each 8 bit integer in the register to 0xFF
  87              on match.  When we subtract it the 8 bit accumulators
  88              will underflow, so this is equal to adding 1. */
  89           accumulator = _mm256_sub_epi8 (accumulator, matches);
  90           accumulator2 = _mm256_sub_epi8 (accumulator2, matches2);
  91
  92           datap += 2;
  93           bytes_read -= 64;
  94         }
  95
  96       /* Horizontally add all 8 bit integers in the register,
  97          and then reset it */
  98       accumulator = _mm256_sad_epu8 (accumulator, zeroes);
  99       lines +=   _mm256_extract_epi16 (accumulator, 0)
 100                + _mm256_extract_epi16 (accumulator, 4)
 101                + _mm256_extract_epi16 (accumulator, 8)
 102                + _mm256_extract_epi16 (accumulator, 12);
 103       accumulator = _mm256_setzero_si256 ();
 104
 105       accumulator2 = _mm256_sad_epu8 (accumulator2, zeroes);
 106       lines +=   _mm256_extract_epi16 (accumulator2, 0)
 107                + _mm256_extract_epi16 (accumulator2, 4)
 108                + _mm256_extract_epi16 (accumulator2, 8)
 109                + _mm256_extract_epi16 (accumulator2, 12);
 110       accumulator2 = _mm256_setzero_si256 ();
 111
 112       /* Finish up any left over bytes */
 113       char *p = (char *)datap;
 114       while (p != end)
 115         lines += *p++ == '\n';
 116     }
 117
 118   *lines_out = lines;
 119   *bytes_out = bytes;
 120
 121   return true;
 122 }