Compiler Explorer

Source code

{.passC: "-mavx2".}
{.passL: "-mavx2".}

import os

type m256i* {.importc: "__m256i", header: "immintrin.h".} = object
type m256* {.importc: "__m256", header: "immintrin.h".} = object

const avx2width* = 32

proc set1_epi8*(b: int8 | uint8): m256i
  {.importc: "_mm256_set1_epi8", header: "immintrin.h".}

proc loadu_byte*(p: ptr byte): m256i
  {.importc: "_mm256_loadu_si256", header: "immintrin.h".}

proc popcnt_u32*(a: int32): int32
  {.importc: "_mm_popcnt_u32", header: "immintrin.h".}

proc cmpeq_epi8*(a, b: m256i): m256i
  {.importc: "_mm256_cmpeq_epi8", header: "immintrin.h".}

proc cmpgt_epi8*(a, b: m256i): m256i
  {.importc: "_mm256_cmpgt_epi8", header: "immintrin.h".}

proc movemask_epi8*(a: m256i): int32
  {.importc: "_mm256_movemask_epi8", header: "immintrin.h".}

proc extract_epi8*(a: m256i, off: int): byte
  {.importc: "_mm256_extract_epi8", header: "immintrin.h".}

proc countGroupByAVX2*(a: openArray[uint8]): array[256, int] =
  let mask0 = set1_epi8(0)
  let mask1 = set1_epi8(1)
  let mask2 = set1_epi8(2)
  let mask3 = set1_epi8(3)
  var i = 0
  while i < a.len-avx2width:
    let ymm = loadu_byte(unsafeAddr a[i])
    result[0] += popcnt_u32 movemask_epi8 cmpeq_epi8(ymm, mask0)
    result[1] += popcnt_u32 movemask_epi8 cmpeq_epi8(ymm, mask1)
    result[2] += popcnt_u32 movemask_epi8 cmpeq_epi8(ymm, mask2)
    result[3] += popcnt_u32 movemask_epi8 cmpeq_epi8(ymm, mask3)
    i += avx2width
  while i < a.len:
    result[a[i]].inc
    inc i

proc main() =
  let p3 = paramStr(1)
  echo countGroupByAVX2([1'u8,2,uint8(p3[0])])

main()