perf(index): parallelize the SA-to-BWT pipeline with OpenMP

api-test/Makefile

@@ -2,6 +2,23 @@ CC=			gcc
 CFLAGS=		-std=gnu99 -Wall -O3
 EXE=		mbmap-one mbmap-batch
 
+# libminibwa.a may contain OpenMP code (parallel BWT construction), so the
+# example programs must link the OpenMP runtime. Mirror the parent Makefile's
+# detection so the setting matches the library that was built.
+UNAME_S=	$(shell uname -s)
+ifeq ($(UNAME_S),Darwin)
+	LIBOMP_PREFIX?=	$(shell test -d /opt/homebrew/opt/libomp && echo /opt/homebrew/opt/libomp || echo /usr/local/opt/libomp)
+	omp_cflags=	-Xpreprocessor -fopenmp -I$(LIBOMP_PREFIX)/include
+	omp_libs=	-L$(LIBOMP_PREFIX)/lib -Wl,-rpath,$(LIBOMP_PREFIX)/lib -lomp
+else
+	omp_cflags=	-fopenmp
+	omp_libs=	-fopenmp
+endif
+omp?=		$(shell printf '\043include <omp.h>\nint main(){return omp_get_max_threads();}' | $(CC) -x c $(omp_cflags) - $(omp_libs) -o /dev/null 2>/dev/null && echo "1" || echo "0")
+ifeq ($(omp),1)
+	LIBS+=$(omp_libs)
+endif
+
 .PHONY:all clean
 
 all:$(EXE)
@@ -10,10 +27,10 @@ all:$(EXE)
 	(cd ..; make libminibwa.a)
 
 mbmap-one:ex-one.c ../libminibwa.a
-	$(CC) $(CFLAGS) -o $@ -I.. $^ -lz -lm
+	$(CC) $(CFLAGS) -o $@ -I.. $^ -lz -lm $(LIBS)
 
 mbmap-batch:ex-batch.c ../libminibwa.a
-	$(CC) $(CFLAGS) -o $@ -I.. $^ -lz -lm
+	$(CC) $(CFLAGS) -o $@ -I.. $^ -lz -lm $(LIBS)
 
 clean:
 	rm -f $(EXE)

bwt.c

@@ -3,6 +3,9 @@
 #include <string.h>
 #include <assert.h>
 #include <stdint.h>
+#ifdef LIBSAIS_OPENMP
+#include <omp.h>
+#endif
 #include "kommon.h"
 #include "kalloc.h"
 #include "bwt.h"
@@ -106,6 +109,119 @@ mb_bwt_t *mb_bwt_init_from_raw(int is_byte, const void *raw_, uint64_t len, uint
 	return bwt;
 }
 
+// Build the rank dict from a byte-packed inverted SA (one BWT char per byte,
+// lower 2 bits used). The single k where a[k]==0 is the implicit $
+// ('primary') and is skipped here. Distinct from mb_bwt_init_from_raw above
+// (which packs 4 chars/byte for the GPL bwtgen path) to keep its hot loop
+// branch-free.
+//
+// Parallel via a 3-phase scan: each thread packs its blocks' BWT bits and a
+// local count starting from zero; serial prefix-sum yields per-thread start
+// offsets; each thread then adds its start offset to the lower 56 bits of
+// every block-count word it owns. The byte-64 delta in the upper 8 bits is
+// invariant under that shift, so phase 3 leaves it alone.
+mb_bwt_t *mb_bwt_init_from_inverted_sa(const uint8_t *a, int64_t len, int64_t primary, int n_thread)
+{
+	mb_bwt_t *bwt;
+	int64_t total_blocks = (len + 127) / 128;
+	uint64_t *delta_c, *start_c;
+	int nt, n_used;
+
+	bwt = mb_bwt_init();
+	bwt->primary = primary;
+	bwt->seq_len = len;
+	bwt->data_len = mb_bwt_data_len(len);
+	bwt->data = kom_malloc(uint64_t, bwt->data_len);
+
+#ifdef LIBSAIS_OPENMP
+	nt = n_thread;
+#else
+	nt = 1;
+#endif
+	n_used = nt; // real thread count; corrected to omp_get_num_threads() inside the region
+	delta_c = kom_malloc(uint64_t, (size_t)nt * 4);
+	start_c = kom_calloc(uint64_t, (size_t)(nt + 1) * 4); // start_c[0..3] must be zero for the prefix-scan base
+
+#ifdef LIBSAIS_OPENMP
+	#pragma omp parallel num_threads(nt)
+#endif
+	{
+#ifdef LIBSAIS_OPENMP
+		int tid = omp_get_thread_num();
+		int actual_nt = omp_get_num_threads();
+#else
+		int tid = 0, actual_nt = 1;
+#endif
+		int64_t blocks_per_t = (total_blocks + actual_nt - 1) / actual_nt;
+		int64_t b_first = (int64_t)tid * blocks_per_t;
+		int64_t b_last  = b_first + blocks_per_t;
+		if (b_first > total_blocks) b_first = total_blocks;
+		if (b_last  > total_blocks) b_last  = total_blocks;
+		int64_t out_start = b_first * 128;
+		int64_t out_end   = b_last  * 128;
+		if (out_end > len) out_end = len;
+		uint64_t local_c[4] = {0,0,0,0};
+		uint64_t local_x[4] = {0,0,0,0};
+		uint64_t *block_count_ptr = 0;
+		int64_t k_data = b_first * 8;
+
+		for (int64_t out = out_start; out < out_end; ++out) {
+			int64_t in = out + (out >= primary);
+			uint8_t b = a[in] & 3;
+			int pos = (int)(out & 0x7f);
+			if (pos == 0) {
+				if (out > out_start) memcpy(&bwt->data[k_data], local_x, 32), k_data += 4;
+				block_count_ptr = &bwt->data[k_data];
+				memcpy(&bwt->data[k_data], local_c, 32), k_data += 4;
+				memset(local_x, 0, 32);
+			} else if (pos == 64) {
+				for (int j = 0; j < 4; ++j)
+					block_count_ptr[j] |= (local_c[j] - block_count_ptr[j]) << BWT_CNT_SHIFT;
+			}
+			++local_c[b];
+			local_x[pos >> 5] |= (uint64_t)b << ((pos & 0x1f) << 1);
+		}
+
+		if (b_last > b_first) memcpy(&bwt->data[k_data], local_x, 32), k_data += 4;
+		assert(k_data == b_last * 8);
+		for (int j = 0; j < 4; ++j) delta_c[(size_t)tid * 4 + j] = local_c[j];
+
+#ifdef LIBSAIS_OPENMP
+		#pragma omp barrier
+		#pragma omp single
+#endif
+		{ // OpenMP may give fewer threads than requested; record the real count
+			n_used = actual_nt;
+			for (int t = 0; t < actual_nt; ++t)
+				for (int j = 0; j < 4; ++j)
+					start_c[(size_t)(t+1) * 4 + j] = start_c[(size_t)t * 4 + j] + delta_c[(size_t)t * 4 + j];
+		}
+
+		uint64_t add[4] = {
+			start_c[(size_t)tid * 4 + 0], start_c[(size_t)tid * 4 + 1],
+			start_c[(size_t)tid * 4 + 2], start_c[(size_t)tid * 4 + 3]
+		};
+		for (int64_t b = b_first; b < b_last; ++b) {
+			uint64_t *p = &bwt->data[b * 8];
+			for (int j = 0; j < 4; ++j) {
+				uint64_t low = p[j] & BWT_CNT_MASK, high = p[j] & ~BWT_CNT_MASK;
+				p[j] = ((low + add[j]) & BWT_CNT_MASK) | high;
+			}
+		}
+	}
+
+	memcpy(&bwt->data[total_blocks * 8], &start_c[(size_t)n_used * 4], 32);
+	bwt->L2[0] = 0;
+	for (int j = 0; j < 4; ++j)
+		bwt->L2[j+1] = bwt->L2[j] + start_c[(size_t)n_used * 4 + j];
+	assert(bwt->L2[4] == (uint64_t)len);
+	assert((uint64_t)(total_blocks * 8 + 4) == bwt->data_len);
+
+	free(delta_c);
+	free(start_c);
+	return bwt;
+}
+
 /********
  * Rank *
  ********/

bwt.h

@@ -53,6 +53,7 @@ mb_bwt_t *mb_bwt_load_raw(const char *fn); // from raw bwt_gen.c output
 int mb_bwt_save(const char *fn, const mb_bwt_t *bwt);
 mb_bwt_t *mb_bwt_load(const char *fn);
 mb_bwt_t *mb_bwt_init_from_raw(int is_byte, const void *raw_, uint64_t len, uint64_t primary);
+mb_bwt_t *mb_bwt_init_from_inverted_sa(const uint8_t *a, int64_t len, int64_t primary, int n_thread);
 
 void mb_bwt_cache(mb_bwt_t *bwt, int32_t len);
 uint64_t mb_bwt_rank11(const mb_bwt_t *bwt, uint64_t k, uint8_t c);

index.c

@@ -1,6 +1,7 @@
 #include <stdlib.h>
 #include <assert.h>
 #include <stdio.h>
+#include <stdint.h>
 #include "libsais.h"
 #include "libsais64.h"
 #include "kommon.h"
@@ -18,63 +19,53 @@ static ko_longopt_t long_opts[] = { // common long options shared across all ind
 static inline uint8_t l2b_c2t(uint8_t b) { return b == 1? 3 : b; } // C(1) -> T(3)
 static inline uint8_t l2b_g2a(uint8_t b) { return b == 2? 0 : b; } // G(2) -> A(0)
 
-// invert the suffix array a[] (32- or 64-bit) to the BWT in seq[], sample the SSA and drop the primary ($); return the primary
-static int64_t sa_to_bwt(void *a, int use_int32, uint8_t *seq, int64_t len, int sa_bit, uint64_t *ssa)
+// Base at position `in` of the (possibly meth-converted) concatenated reference.
+// Non-meth: forward copy then its reverse complement. Meth: c2t_f, g2a_f, g2a_r,
+// c2t_r. Shared by the parallel seq fill and the parallel BWT inversion so the
+// two stay byte-for-byte consistent with the previous serial fill.
+static inline uint8_t l2b_seqbase(const l2b_t *l2b, int is_meth, int both_strand, int64_t n_fwd, int64_t in)
 {
-	int32_t *a32 = a;
-	int64_t *a64 = a, i, primary = -1;
-	uint64_t mask = (1ULL<<sa_bit) - 1;
-	for (i = 0; i <= len; ++i) {
-		int64_t v = use_int32? a32[i] : a64[i];
-		if ((i & mask) == 0) ssa[i>>sa_bit] = v;
-		if (v == 0) primary = i;
-		else if (use_int32) a32[i] = seq[v - 1];
-		else a64[i] = seq[v - 1];
-	}
-	ssa[0] = (uint64_t)-1;
-	for (i = 0; i < primary; ++i) seq[i] = use_int32? a32[i] : a64[i];
-	for (; i < len; ++i) seq[i] = use_int32? a32[i+1] : a64[i+1];
-	return primary;
+	if (!is_meth)
+		return (both_strand && in >= n_fwd)? 3 - l2b_get0(l2b, 2*n_fwd - 1 - in) : l2b_get0(l2b, in);
+	if (in < n_fwd)        return l2b_c2t(l2b_get0(l2b, in));                    // c2t forward
+	else if (in < 2*n_fwd) return l2b_g2a(l2b_get0(l2b, in - n_fwd));            // g2a forward
+	else if (in < 3*n_fwd) return 3 - l2b_g2a(l2b_get0(l2b, 3*n_fwd - 1 - in));  // g2a reverse
+	else                   return 3 - l2b_c2t(l2b_get0(l2b, 4*n_fwd - 1 - in));  // c2t reverse
 }
 
+// Build the FM-index from the libsais suffix array, parallelised across threads
+// when compiled with -DLIBSAIS_OPENMP (no-op otherwise). Pipeline: parallel seq
+// fill -> libsais -> free seq -> fused SSA sampling + BWT inversion (reads bases
+// via l2b so seq[] is already gone) -> compact the wide SA to one byte per BWT
+// char -> mb_bwt_init_from_inverted_sa. A 32-bit suffix array is used when the
+// concatenated length fits in int32_t, halving SA memory and bandwidth. The
+// output is byte-identical to the serial mb_bwt_init_from_raw path for any
+// thread count.
 static mb_bwt_t *mb_bwt_libsais(const l2b_t *l2b, int sa_bit, int both_strand, int is_meth, int n_thread)
 {
 	const int fs = 10000;
-	uint8_t *seq;
-	int64_t i, j, primary, len;
+	uint8_t *seq, *bwt_byte;
+	int64_t primary, len, n_fwd;
 	mb_bwt_t *bwt;
-	uint64_t *ssa, n_ssa;
+	uint64_t *ssa, n_ssa, mask;
 	void *a;
 	int use_int32;
 
-	len = l2b->tot_len * (is_meth? 2 : 1) * (both_strand? 2 : 1);
+	n_fwd = l2b->tot_len;
+	len = n_fwd * (is_meth? 2 : 1) * (both_strand? 2 : 1);
 	// use a 32-bit suffix array (half the memory) when the concatenated length fits in int32_t
 	use_int32 = (len + fs + 1 <= INT32_MAX);
+
 	seq = kom_malloc(uint8_t, len);
-	if (use_int32) a = kom_malloc(int32_t, len + fs + 1);
-	else a = kom_malloc(int64_t, len + fs + 1);
-	if (is_meth) {
-		// c2t forward
-		for (i = 0, j = 0; i < l2b->tot_len; ++i, ++j)
-			seq[j] = l2b_c2t(l2b_get0(l2b, i));
-		// g2a forward
-		for (i = 0; i < l2b->tot_len; ++i, ++j)
-			seq[j] = l2b_g2a(l2b_get0(l2b, i));
-		if (both_strand) {
-			// g2a reverse (reverse complement of g2a converted)
-			for (i = l2b->tot_len - 1; i >= 0; --i, ++j)
-				seq[j] = 3 - l2b_g2a(l2b_get0(l2b, i));
-			// c2t reverse (reverse complement of c2t converted)
-			for (i = l2b->tot_len - 1; i >= 0; --i, ++j)
-				seq[j] = 3 - l2b_c2t(l2b_get0(l2b, i));
-		}
-	} else {
-		for (i = 0, j = 0; i < l2b->tot_len; ++i, ++j)
-			seq[j] = l2b_get0(l2b, i);
-		if (both_strand)
-			for (i = l2b->tot_len - 1; i >= 0; --i, ++j)
-				seq[j] = 3 - l2b_get0(l2b, i);
-	}
+	a = use_int32? (void*)kom_malloc(int32_t, len + fs + 1)
+	             : (void*)kom_malloc(int64_t, len + fs + 1);
+
+#ifdef LIBSAIS_OPENMP
+	#pragma omp parallel for num_threads(n_thread) schedule(static)
+#endif
+	for (int64_t k = 0; k < len; ++k)
+		seq[k] = l2b_seqbase(l2b, is_meth, both_strand, n_fwd, k);
+
 	if (use_int32) {
 		int32_t *a32 = a;
 #ifdef LIBSAIS_OPENMP
@@ -92,15 +83,75 @@ static mb_bwt_t *mb_bwt_libsais(const l2b_t *l2b, int sa_bit, int both_strand, i
 #endif
 		a64[0] = len;
 	}
+	free(seq); // freed early; the fused inversion below reads bases via l2b (4x denser)
+
+	n_ssa = (len + (1ULL<<sa_bit)) >> sa_bit;
+	ssa = kom_malloc(uint64_t, n_ssa);
+	mask = (1ULL << sa_bit) - 1;
+	primary = -1;
 
-	n_ssa = (len + (1<<sa_bit)) >> sa_bit;
-	ssa = kom_calloc(uint64_t, n_ssa);
-	primary = sa_to_bwt(a, use_int32, seq, len, sa_bit, ssa);
+	// Fused SSA sampling + BWT inversion. Reads the BWT character via l2b so
+	// seq[] could be freed above. Exactly one iteration sees a[k]==0 (libsais's
+	// $-marker); the atomic makes that single store well-defined under OpenMP.
+	if (use_int32) {
+		int32_t *a32 = a;
+#ifdef LIBSAIS_OPENMP
+		#pragma omp parallel for num_threads(n_thread) schedule(static)
+#endif
+		for (int64_t k = 0; k <= len; ++k) {
+			int32_t v = a32[k];
+			if (((uint64_t)k & mask) == 0) ssa[(uint64_t)k >> sa_bit] = (uint64_t)(uint32_t)v;
+			if (v == 0) {
+#ifdef LIBSAIS_OPENMP
+				#pragma omp atomic write
+#endif
+				primary = k;
+			} else {
+				a32[k] = l2b_seqbase(l2b, is_meth, both_strand, n_fwd, v - 1);
+			}
+		}
+	} else {
+		int64_t *a64 = a;
+#ifdef LIBSAIS_OPENMP
+		#pragma omp parallel for num_threads(n_thread) schedule(static)
+#endif
+		for (int64_t k = 0; k <= len; ++k) {
+			int64_t v = a64[k];
+			if (((uint64_t)k & mask) == 0) ssa[(uint64_t)k >> sa_bit] = (uint64_t)v;
+			if (v == 0) {
+#ifdef LIBSAIS_OPENMP
+				#pragma omp atomic write
+#endif
+				primary = k;
+			} else {
+				a64[k] = l2b_seqbase(l2b, is_meth, both_strand, n_fwd, v - 1);
+			}
+		}
+	}
+	ssa[0] = (uint64_t)-1;
 	assert(primary != -1);
+
+	// Compact a[] (BWT chars now in the lower 2 bits) into one byte per char,
+	// then free the wide a[] before building the rank dict.
+	bwt_byte = kom_malloc(uint8_t, len + 1);
+	if (use_int32) {
+		int32_t *a32 = a;
+#ifdef LIBSAIS_OPENMP
+		#pragma omp parallel for num_threads(n_thread) schedule(static)
+#endif
+		for (int64_t k = 0; k <= len; ++k) bwt_byte[k] = (uint8_t)(a32[k] & 3);
+	} else {
+		int64_t *a64 = a;
+#ifdef LIBSAIS_OPENMP
+		#pragma omp parallel for num_threads(n_thread) schedule(static)
+#endif
+		for (int64_t k = 0; k <= len; ++k) bwt_byte[k] = (uint8_t)(a64[k] & 3);
+	}
 	free(a);
-	bwt = mb_bwt_init_from_raw(1, seq, len, primary);
+
+	bwt = mb_bwt_init_from_inverted_sa(bwt_byte, len, primary, n_thread);
 	bwt->sa_bit = sa_bit, bwt->n_sa = n_ssa, bwt->sa = ssa;
-	free(seq);
+	free(bwt_byte);
 	return bwt;
 }
 
@@ -290,6 +341,8 @@ int main_index(int argc, char *argv[])
 		else if (c == 902) is_meth = 1;
 	}
 	if (argc - o.ind == 0) return usage_index(stderr, seed, sa_bit, n_thread);
+	if (n_thread < 1) n_thread = 1;
+	kom_assert(sa_bit >= 0 && sa_bit < 32, "-u must be in [0, 31]");
 
 	prefix = o.ind + 1 < argc? argv[o.ind+1] : argv[o.ind];
 	fn_l2b = kom_calloc(char, strlen(prefix) + 10);