Implementing Blelloch Scan: Single-Thread Block Up-Sweep

Introduction

In our last lesson, we tackled the problem of contention in parallel algorithms by using atomic operations to build a correct and efficient histogram. This completed our initial tour of fundamental parallel primitives: reduction and histogramming. We now turn our attention to a third, exceptionally powerful primitive: the parallel scan, also known as the prefix sum.

The scan operation is a cornerstone of parallel computing. At first glance, it appears inherently sequential, yet as we'll see, an elegant and efficient parallel algorithm exists. Mastering this primitive is crucial because it unlocks parallel solutions for a vast range of problems, including sorting algorithms, stream compaction, and even accelerating the recurrence relations at the heart of modern AI architectures like Mamba.

This lesson focuses on the first half of this algorithm. Our goal is to implement the up-sweep, or reduction, phase of the work-efficient Blelloch scan algorithm. This phase builds a tree of partial sums, laying the groundwork for the final scan computation.

The Prefix Sum (Scan) Problem

The all-prefix-sums operation, or scan, takes an array of numbers and a binary associative operator (we'll use addition) and returns a new array where each element is the sum of all preceding elements in the input.

There are two main variants:

Inclusive Scan: The output at index i includes the input element at index i.
- Input: [3, 1, 7, 4]
- Output: [3, 4, 11, 15]
Exclusive Scan: The output at index i includes all elements before index i. The first element is the identity for the operator (0 for addition).
- Input: [3, 1, 7, 4]
- Output: [0, 3, 4, 11]

For the rest of this lesson, we will focus on implementing an exclusive scan. The two forms are easily converted from one to the other.

To understand the challenge in parallelizing this, let's look at the simple, sequential approach.

COMP526 3-7 §3.6 Parallel primitives, Prefix sum

This video from a lecture by Sebastian Wild provides an excellent definition of the prefix sum problem and highlights why its sequential form is difficult to parallelize.

Please watch from the beginning of the clip at the end of the section. Focus on the C-like code for the sequential algorithm and the lecturer's explanation of the data dependency: each step out[i] directly depends on the result of the previous step out[i-1]. This dependency chain is exactly what makes naive parallelization impossible.

As the video explained, the for loop out[i] = out[i-1] + in[i-1] creates a dependency chain that is antithetical to parallel execution. We can't compute all elements of the output simultaneously if each one requires the result from its neighbor. This suggests we need a completely different algorithmic approach.

A Work-Efficient Parallel Scan: The Blelloch Algorithm

In 1990, Guy Blelloch introduced a landmark parallel algorithm for scan that is both elegant and work-efficient. A parallel algorithm is work-efficient if it performs asymptotically the same number of total operations as the optimal sequential algorithm. For an array of size $N$ , the sequential scan performs $O (N)$ additions. A naive parallel approach might perform $O (N lo g N)$ additions, which is inefficient. Blelloch's algorithm achieves $O (N)$ work in $O (lo g N)$ parallel steps.

The algorithm works in two phases, conceptually building a binary tree over the data:

Up-Sweep (Reduction) Phase: Traverses from the leaves of the tree to the root (bottom-up), computing partial sums along the way. At the end of this phase, the last element of the array holds the sum of all elements, just like a standard parallel reduction. The key difference is that we preserve all the intermediate partial sums calculated at each level of the tree.
Down-Sweep Phase: Traverses from the root back to the leaves (top-down), using the partial sums from the up-sweep phase to construct the final exclusive scan result.

This lesson focuses exclusively on the up-sweep phase.

This diagram illustrates the two phases of the Blelloch scan for an array of 8 elements. The top half shows the up-sweep, where partial sums are computed with increasing strides. The bottom half, which we will cover in the next lesson, shows the down-sweep.

The Up-Sweep Phase in Detail

The up-sweep phase is essentially a parallel reduction, but instead of discarding intermediate sums, we store them in-place. It proceeds in $lo g_{2} N$ steps. In each step d, we add the value from an element $2^{d}$ positions away.

Let's trace the up-sweep from the diagram above for an array x of 8 elements:

Step 1 (offset = 1): In parallel, each element at an odd index i adds the value from index i-1.
- x[1] += x[0], x[3] += x[2], x[5] += x[4], x[7] += x[6]
- The array now contains partial sums of pairs. x[3] now holds the sum of the original x[0] through x[3].
Step 2 (offset = 2): In parallel, threads update elements at indices 3, 7, 11, ....
- x[3] += x[1], x[7] += x[5]
- x[3] now holds the sum of the original x0..x3. x[7] holds x4..x7.
Step 3 (offset = 4):
- x[7] += x[3]
- At the end, x[7] contains the sum of all 8 original elements. Notice how the other elements (x[1], x[3], x[5]) hold valuable partial sums that we will need later.

To formalize this, let's examine the pseudocode and a more detailed diagram from the canonical "GPU Gems" book.

Chapter 39. Parallel Prefix Sum (Scan) with CUDA

This section from the NVIDIA Developer site, taken from GPU Gems 3, introduces the work-efficient Blelloch algorithm and details its up-sweep phase.

Read the section titled "A Work-Efficient Parallel Scan". Focus on understanding the conceptual model of building a balanced binary tree over the input array. Pay close attention to Figure 39-3 and the accompanying pseudocode in "Example 3," as they perfectly illustrate the up-sweep phase. Notice how the algorithm works in log n steps, with a stride that doubles at each step.

Implementing the Up-Sweep in a CUDA Kernel

Now, let's translate this algorithm into a CUDA kernel that operates on a single thread block. We will assume the input data for the block has been loaded into a shared memory array, which we'll call temp. A single block implementation is a common building block for larger, multi-block scans.

The core of the implementation is a for loop that iterates $lo g_{2} (block_size)$ times. In each iteration, a subset of threads becomes active to perform the additions for that level of the tree. A __syncthreads() is crucial after each iteration to ensure all threads in the block have completed their additions before the next level begins, preventing race conditions.

Here is a code snippet for the up-sweep phase operating on an array in shared memory.

// Assumes 'n' is the number of elements in the block (a power of 2)
// and 'temp[]' is a __shared__ memory array of size n, pre-loaded with data.
// 'threadIdx.x' is the thread ID within the block.

int offset = 1;
// The loop iterates log2(n) times. 'd' is the number of active threads.
for (int d = n >> 1; d > 0; d >>= 1)
{
    // Synchronize to make sure all threads have finished the previous level.
    __syncthreads();

    // 'd' threads are active in this step.
    if (threadIdx.x < d)
    {
        // Determine the two elements to read.
        int ai = offset * (2 * threadIdx.x + 1) - 1;
        int bi = offset * (2 * threadIdx.x + 2) - 1;

        // The 'right' element 'bi' accumulates the sum from the 'left' element 'ai'.
        temp[bi] += temp[ai];
    }
    // Double the offset for the next level of the tree.
    offset *= 2;
}

Let's trace this code with n=8 for a thread with threadIdx.x = 0 to see how it builds the right side of the reduction tree:

Loop 1: d = 4, offset = 1. The condition threadIdx.x < 4 is true.
- ai = 1 * (2*0 + 1) - 1 = 0
- bi = 1 * (2*0 + 2) - 1 = 1
- temp[1] += temp[0]
Loop 2: d = 2, offset = 2. The condition threadIdx.x < 2 is true.
- ai = 2 * (2*0 + 1) - 1 = 1
- bi = 2 * (2*0 + 2) - 1 = 3
- temp[3] += temp[1]
Loop 3: d = 1, offset = 4. The condition threadIdx.x < 1 is true.
- ai = 4 * (2*0 + 1) - 1 = 3
- bi = 4 * (2*0 + 2) - 1 = 7
- temp[7] += temp[3]

This exactly matches the dependencies shown in the diagrams. The clever indexing ai and bi allows a contiguous block of threads (0 to d-1) to perform the sparse additions required at each level of the tree. After the final __syncthreads() (not shown, but would occur before the down-sweep), the temp array contains all the necessary partial sums for the next phase.

Conclusion

In this lesson, we dissected the prefix sum problem and began constructing one of the most fundamental algorithms in parallel computing. We saw how the inherent data dependencies of a sequential scan make it unsuitable for parallelization and introduced the work-efficient Blelloch scan as the solution.

Key Takeaways:

Prefix Sum (Scan): A powerful primitive that computes a running total over an array. It appears in both exclusive and inclusive forms.
Work-Efficiency: A critical concept for high-performance algorithms. A parallel algorithm should not perform asymptotically more total work than its sequential counterpart. The Blelloch scan is work-efficient, performing $O (N)$ work in $O (lo g N)$ time.
Up-Sweep Phase: The first part of the Blelloch scan is a modified parallel reduction. It computes partial sums in-place over $lo g_{2} N$ steps, preserving all intermediate results which are crucial for the next phase.
Implementation: The up-sweep can be implemented in a single CUDA block using shared memory, a for loop, and __syncthreads() to manage dependencies between levels of the computation.

We have successfully built the "up" ramp. The array in shared memory now holds a structured set of partial sums, with the total sum in the final element. In our next lesson, we will complete the algorithm by implementing the down-sweep phase, which traverses back down the conceptual tree to distribute these sums and construct the final, correct exclusive scan result.