switch back to ispc impls

jfeser · Oct 20, 2022 · 532ddfd · 532ddfd
1 parent 33c4f7e
commit 532ddfd
Show file tree

Hide file tree

Showing 6 changed files with 125 additions and 96 deletions.
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -1,6 +1,6 @@
 {
   inputs = {
-    nixpkgs.url = "github:NixOS/nixpkgs/nixos-22.05";
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
     flake-utils.url = "github:numtide/flake-utils";
   };
   outputs = { self, flake-utils, nixpkgs }@inputs:

diff --git a/lib/bitarray.ispc b/lib/bitarray.ispc
@@ -29,13 +29,54 @@ export void bitarray_or(uniform int32 a[], uniform int32 b[], uniform int32 c[],
 }
 
 export uniform int bitarray_hamming_weight(uniform int32 a[], uniform int len) {
+  /* uniform int part = 0; */
+  /* for (uniform int i = 0; i < len; i++) { */
+  /*   part += popcnt(a[i]); */
+  /* } */
+  /* return part; */
+
+  int part = 0;
+  foreach (i = 0 ... len) {
+    part += popcnt(a[i]);
+  }
+  return reduce_add(part);
+}
+
+export uniform int bitarray_hamming_distance(uniform int32 a[],
+                                             uniform int32 b[],
+                                             uniform int len) {
   uniform int part = 0;
   for (uniform int i = 0; i < len; i++) {
-    part += popcnt(a[i]);
+    part += popcnt(a[i] ^ b[i]);
   }
   return part;
 }
 
+export uniform double bitarray_jaccard_distance(uniform int32 a[],
+                                                uniform int32 b[],
+                                                uniform int len) {
+  if (len == 0) {
+    return 0.;
+  }
+  uniform int inter = 0;
+  uniform int union_ = 0;
+  for (uniform int i = 0; i < len; i++) {
+    union_ += popcnt(a[i] | b[i]);
+    inter += popcnt(a[i] & b[i]);
+  }
+  return union_ == 0 ? 0.0 : 1.0 - ((double)inter / (double)union_);
+}
+
+/* export uniform double bitarray_hash(uniform int32 a[], uniform int len) { */
+/*   uniform int inter = 0; */
+/*   uniform int union_ = 0; */
+/*   for (uniform int i = 0; i < len; i++) { */
+/*     union_ += popcnt(a[i] | b[i]); */
+/*     inter += popcnt(a[i] & b[i]); */
+/*   } */
+/*   return union_ == 0 ? 0.0 : 1.0 - ((double)inter / (double)union_); */
+/* } */
+
 inline bool read_bit(const uniform int32 x[], const int b,
                      const uniform int len) {
   int i = b >> 5;       // div by word_size = 32

diff --git a/lib/bitarray_stubs.c b/lib/bitarray_stubs.c
@@ -111,35 +111,17 @@ CAMLprim value bitarray_hamming_weight_stub_byte(value b) {
 }
 
 CAMLprim intnat bitarray_hamming_distance_stub(value b1, value b2) {
-  const word_t *p1 = (const word_t *)String_val(b1),
-               *p2 = (const word_t *)String_val(b2);
-  int count = 0;
-
-#pragma GCC unroll 8
-  for (int i = 0; i < len(b1); i++) {
-    count += __builtin_popcount(p1[i] ^ p2[i]);
-  }
-  return count;
+  return bitarray_hamming_distance((word_t *)(String_val(b1)),
+                                   (word_t *)(String_val(b2)), len(b1));
 }
 
 CAMLprim value bitarray_hamming_distance_stub_byte(value b1, value b2) {
   return Val_int(bitarray_hamming_distance_stub(b1, b2));
 }
 
 CAMLprim double bitarray_jaccard_stub(value b1, value b2) {
-  const word_t *p1 = (const word_t *)String_val(b1),
-               *p2 = (const word_t *)String_val(b2);
-  int union_ = 0, inter = 0, len = len(b1);
-
-  if (len == 0) {
-    return 0.;
-  }
-  for (int i = 0; i < len; i++) {
-    union_ += __builtin_popcount(p1[i] | p2[i]);
-    inter += __builtin_popcount(p1[i] & p2[i]);
-  }
-
-  return union_ == 0 ? 0.0 : 1.0 - ((double)inter / (double)union_);
+  return bitarray_jaccard_distance((word_t *)(String_val(b1)),
+                                   (word_t *)(String_val(b2)), len(b1));
 }
 
 CAMLprim value bitarray_jaccard_stub_byte(value b1, value b2) {

diff --git a/lib/dune b/lib/dune
@@ -4,22 +4,28 @@
  (foreign_stubs
   (language c)
   (names bitarray_stubs)
-  (flags :standard -O2 -march=native -Wall -g))
+  (flags :standard -O2 -Wall))
  (foreign_archives bitarray)
  (preprocess
   (pps ppx_jane)))
 
 (rule
  (deps bitarray.ispc)
- (targets bitarray.o bitarray.h)
- (action (run ispc -g --pic -o bitarray.o -h bitarray.h bitarray.ispc)))
+ (targets bitarray.o bitarray_avx2.o bitarray_avx512skx.o bitarray_sse4.o
+   bitarray.h)
+ (action
+  (run ispc --target=avx512skx-x8,avx2-i32x8,sse4-i32x8, --pic -o bitarray.o
+    -h bitarray.h bitarray.ispc)))
 
 (rule
  (deps bitarray.o)
  (targets libbitarray.a)
- (action (run ar -rcs libbitarray.a bitarray.o)))
+ (action
+  (run ar -rcs libbitarray.a bitarray.o bitarray_avx2.o bitarray_avx512skx.o
+    bitarray_sse4.o)))
 
 (rule
  (deps bitarray.o)
  (targets dllbitarray.so)
- (action (run gcc -shared -o dllbitarray.so bitarray.o)))
+ (action
+  (run gcc -shared -o dllbitarray.so bitarray.o)))
diff --git a/test/test.ml b/test/test.ml
@@ -264,41 +264,41 @@ let%expect_test "pow" =
   Fmt.pr "%a\n" Bit_matrix.pp (VM.to_matrix (VM.pow input 4));
   [%expect
     {|
-       .█..
-       ..█.
-       ...█
-       ....
-
-       .█..
-       ..█.
-       ...█
-       ....
-
-       ..█.
-       ...█
-       ....
-       ....
-
-       ...█
-       ....
-       ....
-       ....
-
-       ....
-       ....
-       ....
-       .... |}]
+          .█..
+          ..█.
+          ...█
+          ....
+
+          .█..
+          ..█.
+          ...█
+          ....
+
+          ..█.
+          ...█
+          ....
+          ....
+
+          ...█
+          ....
+          ....
+          ....
+
+          ....
+          ....
+          ....
+          .... |}]
 
 let%expect_test "upper-triangle" =
   let module VM = V.Blocked_matrix in
   Fmt.pr "%a" Bit_matrix.pp (VM.to_matrix @@ VM.upper_triangle 5);
   [%expect
     {|
-       █████
-       .████
-       ..███
-       ...██
-       ....█ |}]
+          █████
+          .████
+          ..███
+          ...██
+          ....█ |}]
 
 let%expect_test "to_matrix" =
   let module VM = V.Blocked_matrix in
@@ -313,41 +313,41 @@ let%expect_test "to_matrix" =
   Fmt.pr "%a\n" Bit_matrix.pp VM.(to_matrix O.(id * id));
   [%expect
     {|
-       ........
-       ........
-       ........
-       ........
-       ........
-       ........
-       ........
-       ........
-
-       █.......
-       .█......
-       ..█.....
-       ...█....
-       ....█...
-       .....█..
-       ......█.
-       .......█
-
-       █.......
-       .█......
-       ..█.....
-       ...█....
-       ....█...
-       .....█..
-       ......█.
-       .......█
-
-       █.......
-       .█......
-       ..█.....
-       ...█....
-       ....█...
-       .....█..
-       ......█.
-       .......█ |}]
+          ........
+          ........
+          ........
+          ........
+          ........
+          ........
+          ........
+          ........
+
+          █.......
+          .█......
+          ..█.....
+          ...█....
+          ....█...
+          .....█..
+          ......█.
+          .......█
+
+          █.......
+          .█......
+          ..█.....
+          ...█....
+          ....█...
+          .....█..
+          ......█.
+          .......█
+
+          █.......
+          .█......
+          ..█.....
+          ...█....
+          ....█...
+          .....█..
+          ......█.
+          .......█ |}]
 
 let%expect_test "all" =
   for i = 0 to 100 do