How can I accelerate Quadtree traversal using GPU multithreading in a compute shader?

[ouzhou]

I plan to use a Quadtree to optimize the computation of data volume when implementing the force-directed algorithm.
I have a Quadtree array structure as follows:

@group(0) @binding(0) var<storage, read_wirte> quadTreeData: array<QuadTreeNode>;
struct QuadTreeNode {
    bounds: vec4<f32>, // x y width height
    childIndices: vec4<i32>, // leftTop rightTop leftBottom rightBottom
    data: array<i32, MAX_DATA_PER_NODE>, // some data 
};

Generally speaking, Quadtree traversal can be divided into two types: visit and visitAfter. It seems I can implement the traversal using a stack and a while loop (since recursion is not supported).

I'm stuck because I'm not sure how to leverage GPU workgroups for acceleration.

[greggman]

I'm sorry but I have never implemented quadtree traversal on a GPU. Maybe @toji has and has a demo? Or maybe someone else will have an answer.

I asked ChatGPT but I have no idea how much it hallucinated vs how much it's good info: https://chatgpt.com/share/67ac6c02-8654-8003-9db3-f506cfb803bc

[toji]

I have not implemented quad trees on the GPU either, but when it comes to making use of GPU workgroups we can maybe supply some general advice, depending on what you're trying to do with the data structure. (I'm only vaguely familiar with force-directed algorithms, so I'm not sure how to apply it here.)

The highest level advice I would give is that you're not likely to gain much if you are trying to structure your code such that each invocation of the shader navigate a single level of the tree, or anything really granular like that. It's going to require too much cooperation and waiting between thread. Instead, sometimes it's more beneficial to allow each thread to be more long-running and complete more of the workload if it can do so independently and you can have still have multiple threads all working in tandem.

[greggman]

To add, what specifically are you trying to do? Why do you want to traverse a quad tree? It might be better to know what problem you are trying to solve and then find out what solutions there are for that specific problem.

I don't know this is useful but I remember watching this video on optimizing particle lookups using a GPU friendly technique. It takes a while as the video works through algorithms on the CPU and then later adapting to the GPU

https://youtu.be/rSKMYc1CQHE?t=1323

[ouzhou]

My goal is to implement a WebGPU version of the d3-force algorithm. In this algorithm, to calculate the repulsive force between points, each point typically needs to calculate its repulsive force with every other point. To simplify this computation, d3-force implements a Quadtree.

d3-quadtree
d3-force

The steps are as follows:

1. Traverse all points and create a Quadtree.
2. Use visitAfter to traverse the entire tree and calculate the centroid.
3. Traverse all points to calculate repulsive forces (optimized by the Quadtree).

To utilize multi-threading, I have split the three steps into three compute passes:

1. Multi-threading is possible, but mutex locks may be needed when creating the tree.
2. Not sure if multi-threading is possible here.
3. Multi-threading is possible for this step.

js code

      (() => {
          const strengths = [];

          var distanceMin2 = 1,
            distanceMax2 = Infinity,
            theta2 = 0.81;

          nodes.forEach((node, i) => {
            strengths[i] = -30;
          });

          const fn = (alpha) => {
            const tree = d3
              .quadtree(
                nodes,
                (d) => d.x,
                (d) => d.y
              )
              .visitAfter((quad) => {
                let totalStrength = 0;
                let totalWeight = 0;
                let centerX = 0;
                let centerY = 0;

                // For internal nodes, accumulate forces from child quadrants.
                if (quad.length) {
                  for (let i = 0; i < 4; ++i) {
                    const childQuad = quad[i];
                    if (childQuad) {
                      const childStrength = Math.abs(childQuad.value);
                      totalStrength += childQuad.value;
                      totalWeight += childStrength;
                      centerX += childStrength * childQuad.x;
                      centerY += childStrength * childQuad.y;
                    }
                  }
                  quad.x = centerX / totalWeight;
                  quad.y = centerY / totalWeight;
                }

                // For leaf nodes, accumulate forces from coincident quadrants.
                else {
                  let currentQuad = quad;
                  currentQuad.x = currentQuad.data.x;
                  currentQuad.y = currentQuad.data.y;
                  do {
                    totalStrength += strengths[currentQuad.data.index];
                  } while ((currentQuad = currentQuad.next));
                }

                quad.value = totalStrength;
              });

            nodes.forEach((node, i) => {
              tree.visit((quad, x1, _, x2) => {
                if (!quad.value) return true;

                var deltaX = quad.x - node.x;
                var deltaY = quad.y - node.y;
                var width = x2 - x1;
                var distanceSquared = deltaX * deltaX + deltaY * deltaY;

                // Apply the Barnes-Hut approximation if possible
                if ((width * width) / theta2 < distanceSquared) {
                  if (distanceSquared < distanceMax2) {
                    deltaX = deltaX || jiggle(random);
                    deltaY = deltaY || jiggle(random);
                    distanceSquared = deltaX * deltaX + deltaY * deltaY;

                    if (distanceSquared < distanceMin2) {
                      distanceSquared = Math.sqrt(
                        distanceMin2 * distanceSquared
                      );
                    }

                    node.vx += (deltaX * quad.value * alpha) / distanceSquared;
                    node.vy += (deltaY * quad.value * alpha) / distanceSquared;
                  }
                  return true;
                }

                // Otherwise, process points directly
                else if (quad.length || distanceSquared >= distanceMax2) return;

                // Handle coincident or very close nodes
                if (quad.data !== node || quad.next) {
                  deltaX = deltaX || jiggle(random);
                  deltaY = deltaY || jiggle(random);
                  distanceSquared = deltaX * deltaX + deltaY * deltaY;

                  if (distanceSquared < distanceMin2) {
                    distanceSquared = Math.sqrt(distanceMin2 * distanceSquared);
                  }
                }

                let currentQuad = quad;
                do {
                  if (currentQuad.data !== node) {
                    const force =
                      (strengths[currentQuad.data.index] * alpha) /
                      distanceSquared;
                    node.vx += deltaX * force;
                    node.vy += deltaY * force;
                  }
                } while ((currentQuad = currentQuad.next));
              });
            });
          };

          return fn;
        })()

[greggman]

I quick search brings up this: https://developer.nvidia.com/blog/thinking-parallel-part-iii-tree-construction-gpu/ which shows a method of building a tree on the GPU. It's not a "quad tree" but it might suggest how to do it.

There's also this: https://adms-conf.org/2019-camera-ready/zhang_adms19.pdf