Optimize two-locus site operations

This PR is a combination of three separate modifications. They are
described below and in #3290. Fixes (#3290).

* Two-locus malloc optimizations

This revision moves all malloc operations out of the hot loop in
two-locus statistics, instead providing pre-allocated regions of memory
that the two-locus framework will use to perform work. Instead of simply
passing each pre-allocated array into each function call, we introduce a
simple structure called `two_locus_work_t`, which stores the statistical
results, and provides temporary arrays for storing the normalisation
constants. Setup and teardown methods for this work structure are
provided. Python and C tests are passing and valgrind reports no errors.

* Refactor bit array api, rename to bitset.

As discussed in #2834, this patch renames tsk_bit_array_t to
tsk_bitset_t. Philosophically, we treat these as sets and not arrays,
performing intersections, unions, and membership tests. Therefore, it
makes sense to alter the API to use set theoretic vocabulary, describing
the intent more precisely. Fundamentally, the bitset structure is a list
of N independent bitsets. Each operation on two sets must select the row
on which to operate. The tsk_bitset_t originally tracked `len` only,
which was N, the number of sets. For convenience, we also track the
`row_len`, which is the number of unsigned integers per row. If we
multiply `row_len` by `TSK_BITSET_BITS`, we get the number of bits that
each set (or row) in the list of bitsets will hold.

We had also discussed each set theoretic operation accepting a row index
instead of a pointer to a row within the bitset object. Now, each
operation accepts a row index for each bitset structure passed into the
function. This simplifies the consumption of this API considerably,
removing the need of storing and tracking many intermediate temporary
array pointers. We also see some performance improvements from this
cleanup. For DRY purposes, I've created a private macro,
`BITSET_DATA_ROW`, which abstracts away the pointer arithmetic for
selecting a row out of the list of sets. Because of these changes,
`tsk_bit_array_get_row` is no longer needed and has been removed from
the API.

This change does not change the size of the "chunk", which is the
unsigned integer storing bits. It remains a 32 bit unsigned integer,
which is most performant for bit counting (popcount). I've streamlined
the macros used to determine which integer in the row will be used to
store a particular bit. Everything now revolves around the
TSK_BITSET_BITS macro, which is simply 32 and bitshift operations have
been converted to unsigned integer division.

Testing has been refactored to reflect these changes, removing tests
that operate on a specific rows. Tests in c and python are passing and
valgrind shows no errors.

Fixes (#2834).

* Precompute A/B Counts and Biallelic Summary Func

Precompute A/B counts for each sample set. We were previously computing
them redundantly each for each site pair in our results matrix. The
precomputation happens in a function called `get_mutation_sample_sets`,
which takes our list of sets (`tsk_bitset_t`) for each mutation and
intersects the samples with a particular mutation with the sample sets
passed in by the user. The result is an expanded list of sets with one
set per mutation per sample set. During this operation, we compute the
number of samples containing the given allele for each mutation,
avoiding the need to perform redundant count operations on the data.

In addition to precomputation, we add a non-normalized version of
`compute_general_two_site_stat_result` for situations where we're
computing stats from biallelic loci. We dispatch the computation of the
result based on the number of alleles in the two loci we're comparing.
If the number of alleles in both loci is 2, then we simply perform an LD
computation on the derived alleles for the two loci. As a result, we
remove the need to compute a matrix of LD values, then take a weighted
sum. This is much more efficient and means that we only run the full
multiallelic LD routine on sites that are multiallelic.

Author: lkirk

c/tests/test_core.c

@@ -531,111 +531,107 @@ test_bit_arrays(void)
 {
     // NB: This test is only valid for the 32 bit implementation of bit arrays. If we
     //     were to change the chunk size of a bit array, we'd need to update these tests
-    tsk_bit_array_t arr;
+    tsk_bitset_t arr;
     tsk_id_t items_truth[64] = { 0 }, items[64] = { 0 };
     tsk_size_t n_items = 0, n_items_truth = 0;
 
     // test item retrieval
-    tsk_bit_array_init(&arr, 90, 1);
-    tsk_bit_array_get_items(&arr, items, &n_items);
+    tsk_bitset_init(&arr, 90, 1);
+    CU_ASSERT_EQUAL_FATAL(arr.len, 1);
+    CU_ASSERT_EQUAL_FATAL(arr.row_len, 3);
+    tsk_bitset_get_items(&arr, 0, items, &n_items);
     assert_arrays_equal(n_items_truth, items, items_truth);
 
-    for (tsk_bit_array_value_t i = 0; i < 20; i++) {
-        tsk_bit_array_add_bit(&arr, i);
+    for (tsk_bitset_val_t i = 0; i < 20; i++) {
+        tsk_bitset_set_bit(&arr, 0, i);
         items_truth[n_items_truth] = (tsk_id_t) i;
         n_items_truth++;
     }
-    tsk_bit_array_add_bit(&arr, 63);
-    tsk_bit_array_add_bit(&arr, 65);
+    tsk_bitset_set_bit(&arr, 0, 63);
+    tsk_bitset_set_bit(&arr, 0, 65);
 
     // these assertions are only valid for 32-bit values
     CU_ASSERT_EQUAL_FATAL(arr.data[0], 1048575);
     CU_ASSERT_EQUAL_FATAL(arr.data[1], 2147483648);
     CU_ASSERT_EQUAL_FATAL(arr.data[2], 2);
 
     // verify our assumptions about bit array counting
-    CU_ASSERT_EQUAL_FATAL(tsk_bit_array_count(&arr), 22);
+    CU_ASSERT_EQUAL_FATAL(tsk_bitset_count(&arr, 0), 22);
 
-    tsk_bit_array_get_items(&arr, items, &n_items);
+    tsk_bitset_get_items(&arr, 0, items, &n_items);
     assert_arrays_equal(n_items_truth, items, items_truth);
 
     tsk_memset(items, 0, 64);
     tsk_memset(items_truth, 0, 64);
     n_items = n_items_truth = 0;
-    tsk_bit_array_free(&arr);
+    tsk_bitset_free(&arr);
 
-    // create a length-2 array with 64 bit capacity
-    tsk_bit_array_init(&arr, 64, 2);
-    tsk_bit_array_t arr_row1, arr_row2;
-
-    // select the first and second row
-    tsk_bit_array_get_row(&arr, 0, &arr_row1);
-    tsk_bit_array_get_row(&arr, 1, &arr_row2);
+    // create a length-2 array with 64 bit capacity (two chunks per row)
+    tsk_bitset_init(&arr, 64, 2);
+    CU_ASSERT_EQUAL_FATAL(arr.len, 2);
+    CU_ASSERT_EQUAL_FATAL(arr.row_len, 2);
 
     // fill the first 50 bits of the first row
-    for (tsk_bit_array_value_t i = 0; i < 50; i++) {
-        tsk_bit_array_add_bit(&arr_row1, i);
+    for (tsk_bitset_val_t i = 0; i < 50; i++) {
+        tsk_bitset_set_bit(&arr, 0, i);
         items_truth[n_items_truth] = (tsk_id_t) i;
         n_items_truth++;
     }
 
-    tsk_bit_array_get_items(&arr_row1, items, &n_items);
+    tsk_bitset_get_items(&arr, 0, items, &n_items);
     assert_arrays_equal(n_items_truth, items, items_truth);
 
     tsk_memset(items, 0, 64);
     tsk_memset(items_truth, 0, 64);
     n_items = n_items_truth = 0;
 
     // fill bits 20-40 of the second row
-    for (tsk_bit_array_value_t i = 20; i < 40; i++) {
-        tsk_bit_array_add_bit(&arr_row2, i);
+    for (tsk_bitset_val_t i = 20; i < 40; i++) {
+        tsk_bitset_set_bit(&arr, 1, i);
         items_truth[n_items_truth] = (tsk_id_t) i;
         n_items_truth++;
     }
 
-    tsk_bit_array_get_items(&arr_row2, items, &n_items);
+    tsk_bitset_get_items(&arr, 1, items, &n_items);
     assert_arrays_equal(n_items_truth, items, items_truth);
 
     tsk_memset(items, 0, 64);
     tsk_memset(items_truth, 0, 64);
     n_items = n_items_truth = 0;
 
     // verify our assumptions about row selection
-    CU_ASSERT_EQUAL_FATAL(arr.data[0], 4294967295);
-    CU_ASSERT_EQUAL_FATAL(arr.data[1], 262143);
-    CU_ASSERT_EQUAL_FATAL(arr_row1.data[0], 4294967295);
-    CU_ASSERT_EQUAL_FATAL(arr_row1.data[1], 262143);
-
-    CU_ASSERT_EQUAL_FATAL(arr.data[2], 4293918720);
-    CU_ASSERT_EQUAL_FATAL(arr.data[3], 255);
-    CU_ASSERT_EQUAL_FATAL(arr_row2.data[0], 4293918720);
-    CU_ASSERT_EQUAL_FATAL(arr_row2.data[1], 255);
+    CU_ASSERT_EQUAL_FATAL(arr.data[0], 4294967295); // row1 elem1
+    CU_ASSERT_EQUAL_FATAL(arr.data[1], 262143);     // row1 elem2
+    CU_ASSERT_EQUAL_FATAL(arr.data[2], 4293918720); // row2 elem1
+    CU_ASSERT_EQUAL_FATAL(arr.data[3], 255);        // row2 elem2
 
     // subtract the second from the first row, store in first
-    tsk_bit_array_subtract(&arr_row1, &arr_row2);
+    tsk_bitset_subtract(&arr, 0, &arr, 1);
 
     // verify our assumptions about subtraction
-    CU_ASSERT_EQUAL_FATAL(arr_row1.data[0], 1048575);
-    CU_ASSERT_EQUAL_FATAL(arr_row1.data[1], 261888);
+    CU_ASSERT_EQUAL_FATAL(arr.data[0], 1048575);
+    CU_ASSERT_EQUAL_FATAL(arr.data[1], 261888);
 
-    tsk_bit_array_t int_result;
-    tsk_bit_array_init(&int_result, 64, 1);
+    tsk_bitset_t int_result;
+    tsk_bitset_init(&int_result, 64, 1);
+    CU_ASSERT_EQUAL_FATAL(int_result.len, 1);
+    CU_ASSERT_EQUAL_FATAL(int_result.row_len, 2);
 
     // their intersection should be zero
-    tsk_bit_array_intersect(&arr_row1, &arr_row2, &int_result);
+    tsk_bitset_intersect(&arr, 0, &arr, 1, &int_result);
     CU_ASSERT_EQUAL_FATAL(int_result.data[0], 0);
     CU_ASSERT_EQUAL_FATAL(int_result.data[1], 0);
 
     // now, add them back together, storing back in a
-    tsk_bit_array_add(&arr_row1, &arr_row2);
+    tsk_bitset_union(&arr, 0, &arr, 1);
 
     // now, their intersection should be the subtracted chunk (20-40)
-    tsk_bit_array_intersect(&arr_row1, &arr_row2, &int_result);
+    tsk_bitset_intersect(&arr, 0, &arr, 1, &int_result);
     CU_ASSERT_EQUAL_FATAL(int_result.data[0], 4293918720);
     CU_ASSERT_EQUAL_FATAL(int_result.data[1], 255);
 
-    tsk_bit_array_free(&int_result);
-    tsk_bit_array_free(&arr);
+    tsk_bitset_free(&int_result);
+    tsk_bitset_free(&arr);
 }
 
 static void

c/tskit/core.c

@@ -1260,16 +1260,16 @@ tsk_avl_tree_int_ordered_nodes(const tsk_avl_tree_int_t *self, tsk_avl_node_int_
 }
 
 // Bit Array implementation. Allows us to store unsigned integers in a compact manner.
-// Currently implemented as an array of 32-bit unsigned integers for ease of counting.
+// Currently implemented as an array of 32-bit unsigned integers.
 
 int
-tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length)
+tsk_bitset_init(tsk_bitset_t *self, tsk_size_t num_bits, tsk_size_t length)
 {
     int ret = 0;
 
-    self->size = (num_bits >> TSK_BIT_ARRAY_CHUNK)
-                 + (num_bits % TSK_BIT_ARRAY_NUM_BITS ? 1 : 0);
-    self->data = tsk_calloc(self->size * length, sizeof(*self->data));
+    self->row_len = (num_bits / TSK_BITSET_BITS) + (num_bits % TSK_BITSET_BITS ? 1 : 0);
+    self->len = length;
+    self->data = tsk_calloc(self->row_len * length, sizeof(*self->data));
     if (self->data == NULL) {
         ret = tsk_trace_error(TSK_ERR_NO_MEMORY);
         goto out;
@@ -1278,96 +1278,111 @@ tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length
     return ret;
 }
 
-void
-tsk_bit_array_get_row(const tsk_bit_array_t *self, tsk_size_t row, tsk_bit_array_t *out)
-{
-    out->size = self->size;
-    out->data = self->data + (row * self->size);
-}
+#define BITSET_DATA_ROW(bs, row) ((bs)->data + (row) * (bs)->row_len)
 
 void
-tsk_bit_array_intersect(
-    const tsk_bit_array_t *self, const tsk_bit_array_t *other, tsk_bit_array_t *out)
+tsk_bitset_intersect(const tsk_bitset_t *self, tsk_size_t self_row,
+    const tsk_bitset_t *other, tsk_size_t other_row, tsk_bitset_t *out)
 {
-    for (tsk_size_t i = 0; i < self->size; i++) {
-        out->data[i] = self->data[i] & other->data[i];
+    const tsk_bitset_val_t *restrict self_d = BITSET_DATA_ROW(self, self_row);
+    const tsk_bitset_val_t *restrict other_d = BITSET_DATA_ROW(other, other_row);
+    tsk_bitset_val_t *restrict out_d = out->data;
+    for (tsk_size_t i = 0; i < self->row_len; i++) {
+        out_d[i] = self_d[i] & other_d[i];
     }
 }
 
 void
-tsk_bit_array_subtract(tsk_bit_array_t *self, const tsk_bit_array_t *other)
+tsk_bitset_subtract(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other,
+    tsk_size_t other_row)
 {
-    for (tsk_size_t i = 0; i < self->size; i++) {
-        self->data[i] &= ~(other->data[i]);
+    tsk_bitset_val_t *restrict self_d = BITSET_DATA_ROW(self, self_row);
+    const tsk_bitset_val_t *restrict other_d = BITSET_DATA_ROW(other, other_row);
+    for (tsk_size_t i = 0; i < self->row_len; i++) {
+        self_d[i] &= ~(other_d[i]);
     }
 }
 
 void
-tsk_bit_array_add(tsk_bit_array_t *self, const tsk_bit_array_t *other)
+tsk_bitset_union(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other,
+    tsk_size_t other_row)
 {
-    for (tsk_size_t i = 0; i < self->size; i++) {
-        self->data[i] |= other->data[i];
+    tsk_bitset_val_t *restrict self_d = BITSET_DATA_ROW(self, self_row);
+    const tsk_bitset_val_t *restrict other_d = BITSET_DATA_ROW(other, other_row);
+    for (tsk_size_t i = 0; i < self->row_len; i++) {
+        self_d[i] |= other_d[i];
     }
 }
 
 void
-tsk_bit_array_add_bit(tsk_bit_array_t *self, const tsk_bit_array_value_t bit)
+tsk_bitset_set_bit(tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit)
 {
-    tsk_bit_array_value_t i = bit >> TSK_BIT_ARRAY_CHUNK;
-    self->data[i] |= (tsk_bit_array_value_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i));
+    tsk_bitset_val_t i = (bit / TSK_BITSET_BITS);
+    *(BITSET_DATA_ROW(self, row) + i) |= (tsk_bitset_val_t) 1
+                                         << (bit - (TSK_BITSET_BITS * i));
 }
 
 bool
-tsk_bit_array_contains(const tsk_bit_array_t *self, const tsk_bit_array_value_t bit)
+tsk_bitset_contains(const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit)
 {
-    tsk_bit_array_value_t i = bit >> TSK_BIT_ARRAY_CHUNK;
-    return self->data[i]
-           & ((tsk_bit_array_value_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i)));
+    tsk_bitset_val_t i = (bit / TSK_BITSET_BITS);
+    return *(BITSET_DATA_ROW(self, row) + i)
+           & ((tsk_bitset_val_t) 1 << (bit - (TSK_BITSET_BITS * i)));
 }
 
-tsk_size_t
-tsk_bit_array_count(const tsk_bit_array_t *self)
+static inline uint32_t
+popcount(tsk_bitset_val_t v)
 {
-    // Utilizes 12 operations per bit array. NB this only works on 32 bit integers.
+    // Utilizes 12 operations per chunk. NB this only works on 32 bit integers.
     // Taken from:
     //   https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel
     // There's a nice breakdown of this algorithm here:
     //   https://stackoverflow.com/a/109025
-    // Could probably do better with explicit SIMD (instead of SWAR), but not as
-    // portable: https://arxiv.org/pdf/1611.07612.pdf
     //
-    // There is one solution to explore further, which uses __builtin_popcountll.
-    // This option is relatively simple, but requires a 64 bit bit array and also
-    // involves some compiler flag plumbing (-mpopcnt)
+    // The gcc/clang compiler flag will -mpopcnt will convert this code to a
+    // popcnt instruction (most if not all modern CPUs will support this). The
+    // popcnt instruction will yield some speed improvements, which depend on
+    // the tree sequence.
+    //
+    // NB: 32bit counting is typically faster than 64bit counting for this task.
+    //     (at least on x86-64)
 
-    tsk_bit_array_value_t tmp;
-    tsk_size_t i, count = 0;
+    v = v - ((v >> 1) & 0x55555555);
+    v = (v & 0x33333333) + ((v >> 2) & 0x33333333);
+    return (((v + (v >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24;
+}
+
+tsk_size_t
+tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row)
+{
+    tsk_size_t i = 0;
+    tsk_size_t count = 0;
+    const tsk_bitset_val_t *restrict self_d = BITSET_DATA_ROW(self, row);
 
-    for (i = 0; i < self->size; i++) {
-        tmp = self->data[i] - ((self->data[i] >> 1) & 0x55555555);
-        tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333);
-        count += (((tmp + (tmp >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24;
+    for (i = 0; i < self->row_len; i++) {
+        count += popcount(self_d[i]);
     }
     return count;
 }
 
 void
-tsk_bit_array_get_items(
-    const tsk_bit_array_t *self, tsk_id_t *items, tsk_size_t *n_items)
+tsk_bitset_get_items(
+    const tsk_bitset_t *self, tsk_size_t row, tsk_id_t *items, tsk_size_t *n_items)
 {
     // Get the items stored in the row of a bitset.
-    // Uses a de Bruijn sequence lookup table to determine the lowest bit set. See the
-    // wikipedia article for more info: https://w.wiki/BYiF
+    // Uses a de Bruijn sequence lookup table to determine the lowest bit set.
+    // See the wikipedia article for more info: https://w.wiki/BYiF
 
     tsk_size_t i, n, off;
-    tsk_bit_array_value_t v, lsb; // least significant bit
+    tsk_bitset_val_t v, lsb; // least significant bit
     static const tsk_id_t lookup[32] = { 0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25,
         17, 4, 8, 31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9 };
+    const tsk_bitset_val_t *restrict self_d = BITSET_DATA_ROW(self, row);
 
     n = 0;
-    for (i = 0; i < self->size; i++) {
-        v = self->data[i];
-        off = i * ((tsk_size_t) TSK_BIT_ARRAY_NUM_BITS);
+    for (i = 0; i < self->row_len; i++) {
+        v = self_d[i];
+        off = i * TSK_BITSET_BITS;
         if (v == 0) {
             continue;
         }
@@ -1381,7 +1396,7 @@ tsk_bit_array_get_items(
 }
 
 void
-tsk_bit_array_free(tsk_bit_array_t *self)
+tsk_bitset_free(tsk_bitset_t *self)
 {
     tsk_safe_free(self->data);
 }

c/tskit/core.h

@@ -1104,29 +1104,31 @@ FILE *tsk_get_debug_stream(void);
 
 /* Bit Array functionality */
 
-typedef uint32_t tsk_bit_array_value_t;
+// define a 32-bit chunk size for our bitsets.
+// this means we'll be able to hold 32 distinct items in each 32 bit uint
+#define TSK_BITSET_BITS ((tsk_size_t) 32)
+typedef uint32_t tsk_bitset_val_t;
+
 typedef struct {
-    tsk_size_t size;             // Number of chunks per row
-    tsk_bit_array_value_t *data; // Array data
-} tsk_bit_array_t;
-
-#define TSK_BIT_ARRAY_CHUNK 5U
-#define TSK_BIT_ARRAY_NUM_BITS (1U << TSK_BIT_ARRAY_CHUNK)
-
-int tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length);
-void tsk_bit_array_free(tsk_bit_array_t *self);
-void tsk_bit_array_get_row(
-    const tsk_bit_array_t *self, tsk_size_t row, tsk_bit_array_t *out);
-void tsk_bit_array_intersect(
-    const tsk_bit_array_t *self, const tsk_bit_array_t *other, tsk_bit_array_t *out);
-void tsk_bit_array_subtract(tsk_bit_array_t *self, const tsk_bit_array_t *other);
-void tsk_bit_array_add(tsk_bit_array_t *self, const tsk_bit_array_t *other);
-void tsk_bit_array_add_bit(tsk_bit_array_t *self, const tsk_bit_array_value_t bit);
-bool tsk_bit_array_contains(
-    const tsk_bit_array_t *self, const tsk_bit_array_value_t bit);
-tsk_size_t tsk_bit_array_count(const tsk_bit_array_t *self);
-void tsk_bit_array_get_items(
-    const tsk_bit_array_t *self, tsk_id_t *items, tsk_size_t *n_items);
+    tsk_size_t row_len; // Number of size TSK_BITSET_BITS chunks per row
+    tsk_size_t len;     // Number of rows
+    tsk_bitset_val_t *data;
+} tsk_bitset_t;
+
+int tsk_bitset_init(tsk_bitset_t *self, tsk_size_t num_bits, tsk_size_t length);
+void tsk_bitset_free(tsk_bitset_t *self);
+void tsk_bitset_intersect(const tsk_bitset_t *self, tsk_size_t self_row,
+    const tsk_bitset_t *other, tsk_size_t other_row, tsk_bitset_t *out);
+void tsk_bitset_subtract(tsk_bitset_t *self, tsk_size_t self_row,
+    const tsk_bitset_t *other, tsk_size_t other_row);
+void tsk_bitset_union(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other,
+    tsk_size_t other_row);
+void tsk_bitset_set_bit(tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit);
+bool tsk_bitset_contains(
+    const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit);
+tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row);
+void tsk_bitset_get_items(
+    const tsk_bitset_t *self, tsk_size_t row, tsk_id_t *items, tsk_size_t *n_items);
 
 #ifdef __cplusplus
 }