/**
* Demonstrates how can map a 32-bit integer to a range faster than
* a modulo reduction.
* Assumes x64 processor.
* gcc -O3 -o fastrange fastrange.c
*/
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>

#define RDTSC_START(cycles)                                                   \
    do {                                                                      \
        register unsigned cyc_high, cyc_low;                                  \
        __asm volatile(                                                       \
            "cpuid\n\t"                                                       \
            "rdtsc\n\t"                                                       \
            "mov %%edx, %0\n\t"                                               \
            "mov %%eax, %1\n\t"                                               \
            : "=r"(cyc_high), "=r"(cyc_low)::"%rax", "%rbx", "%rcx", "%rdx"); \
        (cycles) = ((uint64_t)cyc_high << 32) | cyc_low;                      \
    } while (0)

#define RDTSC_FINAL(cycles)                                                   \
    do {                                                                      \
        register unsigned cyc_high, cyc_low;                                  \
        __asm volatile(                                                       \
            "rdtscp\n\t"                                                      \
            "mov %%edx, %0\n\t"                                               \
            "mov %%eax, %1\n\t"                                               \
            "cpuid\n\t"                                                       \
            : "=r"(cyc_high), "=r"(cyc_low)::"%rax", "%rbx", "%rcx", "%rdx"); \
        (cycles) = ((uint64_t)cyc_high << 32) | cyc_low;                      \
    } while (0)

/*
 * Prints the best number of operations per cycle where
 * test is the function call, answer is the expected answer generated by
 * test, repeat is the number of times we should repeat and size is the
 * number of operations represented by test.
 */
#define BEST_TIME(test, answer, repeat, size)                         \
    do {                                                              \
        printf("%s: ", #test);                                        \
        fflush(NULL);                                                 \
        uint64_t cycles_start, cycles_final, cycles_diff;             \
        uint64_t min_diff = (uint64_t)-1;                             \
        int wrong_answer = 0;                                         \
        for (int i = 0; i < repeat; i++) {                            \
            __asm volatile("" ::: /* pretend to clobber */ "memory"); \
            RDTSC_START(cycles_start);                                \
            if (test != answer) wrong_answer = 1;                     \
            RDTSC_FINAL(cycles_final);                                \
            cycles_diff = (cycles_final - cycles_start);              \
            if (cycles_diff < min_diff) min_diff = cycles_diff;       \
        }                                                             \
        uint64_t S = (uint64_t)size;                                  \
        float cycle_per_op = (min_diff) / (float)S;                   \
        printf(" %.2f cycles per operation", cycle_per_op);           \
        if (wrong_answer) printf(" [ERROR]");                         \
        printf("\n");                                                 \
        fflush(NULL);                                                 \
    } while (0)


uint32_t modsum(uint32_t * z, uint32_t N, uint32_t * accesses, uint32_t nmbr) {
  uint32_t sum = 0;
  for(uint32_t j = 0; j < nmbr ; ++j ) {
    sum += z[accesses[j] % N];
  }
  return sum;
}

uint32_t fastsum(uint32_t * z, uint32_t N, uint32_t * accesses, uint32_t nmbr) {
  uint32_t sum = 0;
  uint64_t N64 = (uint64_t) N;
  for(uint32_t j = 0; j < nmbr ; ++j ) {
    sum += z[(accesses[j] * N64)>> 32] ;
  }
  return sum;
}

#ifdef __AVX2__
#include <x86intrin.h>

// assume N is divisible by 4
uint32_t vectorsum(uint32_t * z, uint32_t N, uint32_t * accesses, uint32_t nmbr) {
  __m256i Nvec = _mm256_set1_epi32(N);
  __m128i sum = _mm_setzero_si128();
  for(uint32_t j = 0; j < nmbr ; j+=4) {
     __m256i fourints = _mm256_loadu_si256((const __m256i *)(accesses + j));
     __m256i four64bitsproducts =  _mm256_mul_epu32(fourints, Nvec);
     __m256i fourtop32ints = _mm256_srli_epi64(four64bitsproducts,32);
     __m128i four32ints = _mm256_i64gather_epi32 (z,fourtop32ints , 4);
     sum = _mm_add_epi32(sum, four32ints);
  }
  uint32_t buffer[4];
  _mm_storeu_si128((__m128i *)buffer,sum);
  return buffer[0] + buffer[1] + buffer[2] + buffer[3];
}

// credit: Harold Aptroot
uint32_t maskedvectorsum(uint32_t * z, uint32_t N, uint32_t * accesses,
     uint32_t nmbr) {
  __m256i Nvec = _mm256_set1_epi32(N - 1);
  __m256i sum = _mm256_setzero_si256();
  for(uint32_t j = 0; j < nmbr ; j += 8) {
     __m256i indexes = _mm256_loadu_si256((__m256i*)(accesses + j));
     indexes = _mm256_and_si256(indexes, Nvec);
     __m256i fi = _mm256_i32gather_epi32((int*)z, indexes, 4);
     sum = _mm256_add_epi32(sum, fi);
  }
  __m128i sum128 = _mm_add_epi32(_mm256_extracti128_si256(sum, 0), _mm256_extracti128_si256(sum, 1));
  sum128 = _mm_hadd_epi32(sum128, sum128);
  return _mm_extract_epi32(sum128, 0) + _mm_extract_epi32(sum128, 1);
}
#endif

// N is a power of two
uint32_t maskedsum(uint32_t * z, uint32_t N, uint32_t * accesses, uint32_t nmbr) {
  uint32_t sum = 0;
  for(uint32_t j = 0; j < nmbr ; ++j ) {
    sum += z[accesses[j] & (N-1)] ;
  }
  return sum;
}



void demo(uint32_t N) {
  printf("N = %d\n", N);
  uint32_t * z = malloc(N * sizeof(uint32_t));
  for(uint32_t i = 0 ; i < N; ++i) z[i] = rand(); // some rand. number
  uint32_t nmbr = 500;
  uint32_t * accesses = malloc(nmbr * sizeof(uint32_t));
  for(uint32_t i = 0 ; i < nmbr; ++i) accesses[i] = rand(); // some rand. number
  uint32_t expected1 = modsum(z,N,accesses,nmbr);
  uint32_t expected2 = fastsum(z,N,accesses,nmbr);
  BEST_TIME(modsum(z,N,accesses,nmbr), expected1, 1000, nmbr);
  BEST_TIME(fastsum(z,N,accesses,nmbr), expected2, 1000, nmbr);
#ifdef __AVX2__
  uint32_t expected3 = vectorsum(z,N,accesses,nmbr);
  if(N % 4 == 0) BEST_TIME(vectorsum(z,N,accesses,nmbr), expected3, 1000, nmbr);
  uint32_t expected4 = maskedvectorsum(z,N,accesses,nmbr);
  if(N % 8 == 0) BEST_TIME(maskedvectorsum(z,N,accesses,nmbr), expected4, 1000, nmbr);
#endif
  free(z);
  free(accesses);
}

void demopoweroftwo(uint32_t N) {
  printf("N = %d\n", N);
  uint32_t * z = malloc(N * sizeof(uint32_t));
  for(uint32_t i = 0 ; i < N; ++i) z[i] = rand(); // some rand. number
  uint32_t nmbr = 500;
  uint32_t * accesses = malloc(nmbr * sizeof(uint32_t));
  for(uint32_t i = 0 ; i < nmbr; ++i) accesses[i] = rand(); // some rand. number
  uint32_t expected1 = modsum(z,N,accesses,nmbr);
  uint32_t expected2 = fastsum(z,N,accesses,nmbr);
  BEST_TIME(modsum(z,N,accesses,nmbr), expected1, 1000, nmbr);
  BEST_TIME(fastsum(z,N,accesses,nmbr), expected2, 1000, nmbr);
  BEST_TIME(maskedsum(z,N,accesses,nmbr), expected1, 1000, nmbr);
#ifdef __AVX2__
  uint32_t expected3 = vectorsum(z,N,accesses,nmbr);
  BEST_TIME(vectorsum(z,N,accesses,nmbr), expected3, 1000, nmbr);
  uint32_t expected4 = maskedvectorsum(z,N,accesses,nmbr);
  if(N % 8 == 0) BEST_TIME(maskedvectorsum(z,N,accesses,nmbr), expected4, 1000, nmbr);
#endif

  free(z);
  free(accesses);
}

int main() {
  demo(31);
  demo(1500);
  demo(15000);
  demopoweroftwo(32);
  demopoweroftwo(4096);
  demopoweroftwo(65536);
}
