Perf - Shift coverage handling to whole-word bit math.

Author: benjeffery

c/examples/Makefile

@@ -23,7 +23,7 @@ TSKIT_SOURCE=../tskit/*.c ../subprojects/kastore/kastore.c
 targets = api_structure error_handling \
 	haploid_wright_fisher streaming \
 	tree_iteration tree_traversal \
-	take_ownership
+	take_ownership haplotype_benchmark
 
 all: $(targets)
 
@@ -32,4 +32,3 @@ $(targets): %: %.c
 
 clean:
 	rm -f $(targets)
-

c/examples/haplotype_benchmark.c

@@ -0,0 +1,90 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <time.h>
+
+#include <tskit.h>
+#include <tskit/genotypes.h>
+#include <tskit/tables.h>
+
+#define CHECK_TSK(err)                                                                  \
+    do {                                                                                \
+        if ((err) < 0) {                                                                \
+            fprintf(stderr, "Error: line %d: %s\n", __LINE__, tsk_strerror(err));       \
+            exit(EXIT_FAILURE);                                                         \
+        }                                                                               \
+    } while (0)
+
+#define NUM_ITERATIONS 1
+#define MAX_BENCHMARK_NODES 500
+
+int
+main(int argc, char **argv)
+{
+    int ret;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t treeseq;
+    tsk_haplotype_t haplotype_decoder;
+    int8_t *haplotype = NULL;
+    double elapsed_seconds;
+    clock_t start_clock, end_clock;
+    uint64_t checksum = 0;
+
+    const char *filename = "../../simulated_chrom_21_100k.ts";
+    if (argc > 1) {
+        filename = argv[1];
+    }
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CHECK_TSK(ret);
+
+    ret = tsk_table_collection_load(&tables, filename, 0);
+    CHECK_TSK(ret);
+
+    ret = tsk_treeseq_init(&treeseq, &tables, 0);
+    CHECK_TSK(ret);
+
+    tsk_size_t num_nodes = tsk_treeseq_get_num_nodes(&treeseq);
+    tsk_size_t num_sites = tsk_treeseq_get_num_sites(&treeseq);
+    if (num_sites == 0) {
+        fprintf(stderr, "Tree sequence has no sites\n");
+        exit(EXIT_FAILURE);
+    }
+
+    tsk_id_t node_limit
+        = (tsk_id_t) (num_nodes < MAX_BENCHMARK_NODES ? num_nodes : MAX_BENCHMARK_NODES);
+
+    ret = tsk_haplotype_init(&haplotype_decoder, &treeseq, 0, (tsk_id_t) num_sites);
+    CHECK_TSK(ret);
+
+    haplotype = malloc(num_sites * sizeof(*haplotype));
+    if (haplotype == NULL) {
+        fprintf(stderr, "Failed to allocate haplotype buffer\n");
+        exit(EXIT_FAILURE);
+    }
+
+    start_clock = clock();
+    for (int iter = 0; iter < NUM_ITERATIONS; iter++) {
+        for (tsk_id_t node = 0; node < node_limit; node++) {
+            ret = tsk_haplotype_decode(&haplotype_decoder, node, haplotype);
+            CHECK_TSK(ret);
+            for (tsk_id_t site = 0; site < (tsk_id_t) num_sites; site++) {
+                checksum += (uint64_t) haplotype[site];
+            }
+        }
+    }
+    end_clock = clock();
+
+    elapsed_seconds = (double) (end_clock - start_clock) / CLOCKS_PER_SEC;
+
+    printf("Loaded tree sequence from %s\n", filename);
+    printf("Decoded %d iterations over %lld nodes × %lld sites in %.3f seconds\n",
+        NUM_ITERATIONS, (long long) node_limit, (long long) num_sites, elapsed_seconds);
+    printf("Checksummed haplotypes: %llu\n", (unsigned long long) checksum);
+
+    free(haplotype);
+    tsk_haplotype_free(&haplotype_decoder);
+    tsk_treeseq_free(&treeseq);
+    tsk_table_collection_free(&tables);
+    return EXIT_SUCCESS;
+}

c/meson.build

@@ -113,6 +113,9 @@ if not meson.is_subproject()
       executable('tree_traversal',
           sources: ['examples/tree_traversal.c'], 
           link_with: [tskit_lib], dependencies: lib_deps)
+      executable('haplotype_benchmark',
+          sources: ['examples/haplotype_benchmark.c'],
+          link_with: [tskit_lib], dependencies: lib_deps)
       executable('streaming',
           sources: ['examples/streaming.c'], 
           link_with: [tskit_lib], dependencies: lib_deps)

c/tskit/genotypes.c

@@ -123,6 +123,21 @@ tsk_haplotype_bitset_next(
     return start < limit ? start : limit;
 }
 
+static inline uint64_t
+tsk_haplotype_mask_from_offsets(uint32_t start_offset, uint32_t end_offset)
+{
+    if (start_offset >= end_offset) {
+        return 0;
+    }
+    if (start_offset == 0 && end_offset >= 64) {
+        return UINT64_MAX;
+    }
+    uint64_t high_mask
+        = end_offset >= 64 ? UINT64_MAX : ((UINT64_C(1) << end_offset) - 1);
+    uint64_t low_mask = start_offset == 0 ? 0 : ((UINT64_C(1) << start_offset) - 1);
+    return high_mask & ~low_mask;
+}
+
 static bool
 tsk_haplotype_find_next_uncovered(tsk_haplotype_t *self, tsk_size_t start,
     tsk_size_t end, const int32_t *interval_start, const int32_t *interval_end,
@@ -139,6 +154,7 @@ tsk_haplotype_find_next_uncovered(tsk_haplotype_t *self, tsk_size_t start,
     uint64_t start_mask = UINT64_MAX << (start & 63);
     for (; word <= last_word && word < self->num_bit_words; word++) {
         if (self->unresolved_counts[word] == 0) {
+            start_mask = UINT64_MAX;
             continue;
         }
         uint64_t word_bits = self->unresolved_bits[word];
@@ -149,26 +165,55 @@ tsk_haplotype_find_next_uncovered(tsk_haplotype_t *self, tsk_size_t start,
             uint64_t end_mask = UINT64_MAX >> (63 - ((end - 1) & 63));
             word_bits &= end_mask;
         }
+        if (word_bits == 0) {
+            start_mask = UINT64_MAX;
+            continue;
+        }
+        if (interval_count > 0) {
+            int32_t word_left = (int32_t)(word << 6);
+            int32_t word_right = word_left + 64;
+            uint64_t coverage_mask = 0;
+            for (tsk_size_t p = 0; p < interval_count; p++) {
+                int32_t interval_left = interval_start[p];
+                int32_t interval_right = interval_end[p];
+                if (interval_left >= interval_right) {
+                    continue;
+                }
+                if (interval_right <= word_left || interval_left >= word_right) {
+                    continue;
+                }
+                int32_t clipped_left
+                    = interval_left > word_left ? interval_left : word_left;
+                int32_t clipped_right
+                    = interval_right < word_right ? interval_right : word_right;
+                if ((int32_t) start > clipped_left) {
+                    clipped_left = (int32_t) start;
+                }
+                if ((int32_t) end < clipped_right) {
+                    clipped_right = (int32_t) end;
+                }
+                if (clipped_left >= clipped_right) {
+                    continue;
+                }
+                uint32_t start_offset = (uint32_t)(clipped_left - word_left);
+                uint32_t end_offset = (uint32_t)(clipped_right - word_left);
+                coverage_mask
+                    |= tsk_haplotype_mask_from_offsets(start_offset, end_offset);
+                if (coverage_mask == UINT64_MAX) {
+                    break;
+                }
+            }
+            word_bits &= ~coverage_mask;
+        }
         while (word_bits != 0) {
-            uint64_t lowest_bit = word_bits & (~word_bits + 1);
             tsk_size_t bit = tsk_haplotype_ctz64(word_bits);
-            word_bits ^= lowest_bit;
+            word_bits &= word_bits - 1;
             tsk_size_t bit_index = (word << 6) + bit;
             if (bit_index >= end) {
                 break;
             }
-            bool covered = false;
-            for (tsk_size_t p = 0; p < interval_count; p++) {
-                if (interval_start[p] <= (int32_t) bit_index
-                    && (int32_t) bit_index < interval_end[p]) {
-                    covered = true;
-                    break;
-                }
-            }
-            if (!covered) {
-                *out_index = bit_index;
-                return true;
-            }
+            *out_index = bit_index;
+            return true;
         }
         start_mask = UINT64_MAX;
     }