aboutsummaryrefslogtreecommitdiff
path: root/count_pairs_xor_range.c
blob: 7f0b07f806c1979ea7358d53e5293e74d73c4152 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
// 1803. Count Pairs With XOR in a Range
#include "leetcode.h"

/*
 * given a 0-indexed integer array 'numss' and two integers 'low' and 'high',
 * return the numsber of nice pairs. a nice pair is a pair '(i, j)' where '0 <=
 * i < j < numss.len' and 'low <= (numss[i] ^ numss[j]) <= high'.
 */

#pragma GCC optimize("O3")
#pragma GCC target("avx2")
#include <x86intrin.h>

/*
 * the following code is heavily optimised with avx2 intrinsics and has zero
 * branching. the general structure follows an intuitive brute force approach:
 *
 * int pairs = 0;
 * for (int i = 0; i < numsSize - 1; i++)
 *   for (int j = i + 1; j < n; j++)
 *     if ((nums[i] ^ nums[j]) >= low && (nums[i] ^ nums[j]) <= high)
 *       pairs++;
 * return pairs;
 */

int countPairs(int *nums, int numsSize, int low, int high) {
  int pairs = 0;

  /*
   * store input as 16-bit integers, allowing 16 pairs to be processed per
   * iteration. 32 bytes of extra entries are added to avoid kludge for last
   * entries. the inner loop will process the array in chunks of 16 shorts per
   * execution
   */
  unsigned short __attribute__((aligned(32))) d[numsSize + 16];
  for (int i = 0; i < numsSize; i++)
    d[i] = nums[i];
  for (int i = numsSize; i < numsSize + 16; i++)
    d[i] = 0xFFFF;

  /*
   * create arrays for lower and upper bound. the '_m256i' type is a 256-bit
   * register that can hold 16 shorts at a time. two registers are allocated
   * with each 16-bit value set to low and high respectively
   */
  const __m256i low_mask = _mm256_set1_epi16((short)low);
  const __m256i high_mask = _mm256_set1_epi16((short)high);

  for (int i = 0; i < numsSize - 1; i++) {
    // iterate through elements and create a register array 'a[i]' where 16
    // elements are set to 'd[i]'
    __m256i ai = _mm256_set1_epi16(d[i]);

    // critical section
    for (int j = i + 1; j < numsSize; j += 16) {
      // read 16 numbers to array 'a[j]'
      __m256i aj = _mm256_lddqu_si256((__m256i *)(d + j));

      // XOR them with array 'a[i]', essentially creating array of '{d[j] ^
      // d[i], d[j + 1] ^ d[i], ..., d[j + 15] ^ d[i]}'
      __m256i xij = _mm256_xor_si256(ai, aj);

      /*
       * compare XOR results with upper and lower bounds. the 'cmpgt' intrinsic
       * creates an array of compare results, where an element is set to
       * '0xFFFF' when the corresponding element in 'xij' is greater than the
       * corresponding element in 'high_mask'
       */
      __m256i high_ans = _mm256_cmpgt_epi16(xij, high_mask);
      __m256i low_ans = _mm256_cmpgt_epi16(low_mask, xij);

      /*
       * OR the comparison results to mask, count the number of high bits set,
       * and add it to 'pairs'. note, the comparison sets bits when the XOR is
       * outside the '[low, high]' interval
       */
      __m256i ans = _mm256_or_si256(high_ans, low_ans);
      pairs += 32 - __builtin_popcount(_mm256_movemask_epi8(ans));
    }
  }
  return pairs >> 1;
}

int main() {
  int n1[] = {1, 4, 2, 7}, n2[] = {9, 8, 4, 2, 1};
  printf("%d\n", countPairs(n1, ARRAY_SIZE(n1), 2, 6));  // expect: 6
  printf("%d\n", countPairs(n2, ARRAY_SIZE(n2), 5, 14)); // expect: 8
}