Parallel prefix sum/scan on trees

Question

Let $T$ be a rooted, finitely branching tree, with values from $A$ in its nodes; $T(j)$ is the value of a node called $j$. $A$ is a commutative monoid. The "bottom-up scan tree" of $T$ is a rooted tree $T^\ast$ isomorphic to $T$, such that if $j$ is a leaf of $T$, $T(j)=T^\ast(j)$, and otherwise $T^\ast(j) = \sum_{i\leq j}T(i)$, where $i\leq j$ indicates that $i$ is a descendant of $j$.

1. What algorithms exist to compute $T^\ast$ quickly on a parallel random access machine?

2. Is there an algorithm which performs in time strictly less than $O(N)$, even in the pathological case where $T$ is a linked list?

3. In the real world, how should $T^\ast$ and $T$ be implemented in order to get good performance on a GPU?

My motivation comes from studying the description of the network simplex algorithm in the book "Network flows" by Ahuja and Magnanti, where this problem appears as "compute-flows" in section 11.4. I am trying to solve optimal transport problems.

I think that it should be relatively straightforward to solve this problem in a loop indexed by the height of tree, where the time of each loop iteration is at most $\log(N)$. If the tree is balanced, then this algorithm suffices, and there is no need for sophisticated algorithms. However, I can see no reason why the trees in the network simplex algorithm for minimum flow cost should be reasonably balanced. (If there are adaptations of the network simplex algorithm where the tree is always balanced, I would be interested in hearing about this.)

Answer 1

Here is one approach that ought to run fast on real hardware. Store the structure of your tree in an array $P$ so that the parent of node $i$ is found at $P[i]$. For example, with $$ P = [-1, 0, 0, 2, 3], $$ $P[3] = 2$ so the parent of node 3 is node 2. -1 is a dummy value for the parent-less root node. Store all the intermediate sum of the nodes in an equally long array called $V$. Create another array $C$ with all elements set to 0 where you compute how many children each node has:

N = len(P)
C = [0] * N
for i in range(N):
    C[P[i]] += 1

Now process the tree in bottom-up order. Leaf nodes are childless, so we begin with those:

changed = False
parfor i in range(N):
    if C[i] == 0:
        C[i] -= 1
        atomic C[P[i]] -= 1
        atomic V[P[i]] += V[i]
        changed = True

So for every leaf node, we set its child count to -1 to indicate that it has been processed. Then we decrement the parent's child count.Effectively, this generates new leaf nodes and we can iterate the process until no work remain:

while True:
    changed = False
    parfor i in range(N):
        if C[i] == 0:
            C[i] -= 1
            atomic C[P[i]] -= 1
            atomic V[P[i]] += V[i]
            changed = True
    if not changed:
        break

I've invented the keyword "parfor" to denote loops whose iterations can run in parallel and the keyword "atomic" for operations that must run atomically as multiple threads may write to the same memory. There are many standard approaches for handling this effectively. I've also glossed over handling of the parent-less root note. But it's easy to add checks for. If memory is the bottleneck you could rearrange the tree so that nodes and their children are close. I.e, the $|i - P[i]|$ should preferably be small.