1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

void merge_sequential(int* A, int* B, int* C, int m, int n) {
	int i = 0, j = 0, k = 0;  // Indices for A, B, and C
	while (i < m && j < n) {
		if (A[i] < B[j]) {
			C[k++] = A[i++];
		} else {
			C[k++] = B[j++];
		}
		if (i == m) {  // Done with A[], handling remaining B
			while (j < n) {
				C[k++] = B[j++];
			}
		}
		else {  // Done with B[], handling remaining A
			while (i < m) {
				C[k++] = A[i++];
			}
		}
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

int co_rank(int k, int* A, int m, int* B, int n) {  // C[k] comes from A[i] of B[j]
	// k = i + j
	int i = k < m ? k : m;  // max starting search value for A, i.e. A[k-1] < B[0]
	int i_low = 0 > (k - n) ? 0 : k - n;  // when B is done, min starting search value for A is k-n
	int j = k - i;
	int j_low = 0 > (k - m) ? 0 : (k - m);
	int delta;
	bool active = true;
	while (active) {  // Binary search for C[k]
		if (i > 0 && j < n && A[i - 1] > B[j]) {
			delta = (i - i_low + 1) >> 1;
			j_low = j;
			j += delta;
			i -= delta;
		} else if (j > 0 && i < m && B[j - 1] > A[i]) {
			delta = (j - j_low + 1) >> 1;
			i_low = i;
			i += delta;
			j -= delta;
		} else {  // Found the correct position for C[k]
			active = false;
		}
		return i;
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

__global__
void mergre_basic_kernel(int* A, int* B, int* C, int m, int n) {  // Each thread handles a section of C
	int tid = blockIdx.x * blockDim.x + threadIdx.x;
	int elementsPerThread = ceil(m + n) / (blockDim.x * gridDim.x);
	int start = tid * elementsPerThread;
	int end = std::min(start + elementsPerThread, m + n);
	// Determin the range of A and B to be merged for this thread
	int i_curr = co_rank(start, A, m, B, n);
	int i_next = co_rank(end, A, m, B, n);
	int j_curr = start - i_curr;
	int j_next = end - i_next;
	merge_sequential(A + i_curr, B + j_curr, C + start, i_next - i_curr, j_next - j_curr);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

__global__
void merge_tiled_kernel(int* A, int* B, int* C, int m, int n, int tile_size) {
	/* Part 1: Identify block-level output & input subarrays */
	// Use extern keywords to determine 
	// the shared memory size at runtime rather than compilation
	extern __shared__ int shared_AB[];
	int* A_s = &shared_AB[0];  // Start index of ShareA
	int* B_s = &shared_AB[tile_size];  // Start index of ShareB
	int C_curr = blockIdx.x * ceil((m + n) / gridDim.x);  // Start index of C for this block
	int C_next = std::min(C_curr + int(ceil((m + n) / gridDim.x)), m + n);  // End index of C for this block

	if (threadIdx.x == 0) {
		A_s[0] = co_rank(C_curr, A, m, B, n);  // Make block level co-rank values visible
		A_s[1] = co_rank(C_next, A, m, B, n);  // Next threads co-rank values in the block
	}
	__synctyhreads();
	int A_curr = A_s[0];
	int A_next = A_s[1];
	int B_curr = C_curr - A_curr;
	int B_next = C_next - A_next;

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

/* Part 2: Loading A & B elements into the shared memory */
	int counter = 0;
	int lenC = C_next - C_curr;
	int lenA = A_next - A_curr;
	int lenB = B_next - B_curr;
	int num_iterations = ceil(lenC / tile_size);
	// index of completed merge in 
	int C_completed = 0;
	int A_completed = 0;
	int B_completed = 0;
	while (counter < num_iterations) {  // Each iter threads in a block will generate tile_size C elements
		// Loading tile_size A and B elements into shared memory
		for (int i = 0; i < tile_size; i += blockDim.x) {  // Coalecsing loading from global memory
			if (i + threadIdx.x < lenA - A_completed) {
				A_s[i + threadIdx.x] = A[i + threadIdx.x + A_curr + A_completed];
			}
			if (i + threadIdx.x < lenB - B_completed) {
				B_s[i + threadIdx.x] = B[i + threadIdx.x + B_curr + B_completed];
			}
		}
		__syncthreads();

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

		/* Part 3: All threads merge their subarrays in prallel */
		int c_curr = threadIdx.x * (tile_size / blockDim.x);  // Output index in shared memory
		int c_next = c_curr + (tile_size / blockDim.x);
		c_curr = (c_curr <= lenC - C_completed) ? c_curr : lenC - C_completed;
		c_next = (c_next <= lenC - C_completed) ? c_next : lenC - C_completed;
		// find co-rank for c_curr and c_next
		int a_curr = co_rank(c_curr, A_s, std::min(tile_size, lenA - A_completed), 
								B_s, std::min(tile_size, lenB - B_completed));
		int b_curr = c_curr - a_curr;
		int a_next = co_rank(c_next, A_s, std::min(tile_size, lenA - A_completed), 
								B_s, std::min(tile_size, lenB - B_completed));
		int b_next = c_next - a_next;
		// merge the subarrays
		merge_sequential(A_s + a_curr, B_s + b_curr, C + C_curr + C_completed + c_curr, a_next - a_curr, b_next - b_curr);
		// Update completed indices
		C_completed += tile_size;
		A_completed += co_rank(tile_size, A_s, tile_size, B_s, tile_size);  // Idx of A_s to generate tile_size Idx of merged A_s and B_s
		B_completed += tile_size - A_completed;
	}
}