TorchSparse: regular computation from irregular sparsity

Keep the small example from part 11 in view. Five input points P0 to P4, five output points Q0 to Q4, a 3 x 3 kernel, so nine weight offsets from W(-1,-1) to W(1,1). The map for this layer has eleven tuples. Two of them belong to W(-1,-1), one to W(-1,0), five to W(0,0), one to W(1,0), two to W(1,1), and the other four offsets have no entry at all in this tiny cloud.

The existing GPU implementation walks that list one weight at a time. It is weight-stationary: pick an offset, hold its C_in x C_out matrix still, and stream every map entry for that offset through it. Take W(-1,0). Its only entry is (P1, Q3), so the engine copies the single feature row f1 into a buffer of shape 1 x C_in, multiplies it by W(-1,0) to get one partial sum of shape 1 x C_out, and adds that row into Q3. Now take W(0,0). With stride 1 every input point sits on top of its own output, so the centre map contains all five entries. The buffer is the whole 5 x C_in input, the matmul is full height, and five partial sums land one to one in Q0 to Q4. Then W(1,0) with only (P3, Q1), then W(1,1) with (P1, Q0) and (P4, Q3), a 2 x C_in buffer whose two results go to two different output rows.

That is the whole rule, and it is the same three phases every time. The names are the ones the lecture and the paper use: gather, matmul, scatter. For each weight offset delta, gather the |M_delta| input rows named by the map into a contiguous buffer, run one dense matrix multiplication of the |M_delta| x C_in buffer by the C_in x C_out weight, and scatter-accumulate the |M_delta| x C_out partial sums into the output rows named by the map. Written as an update rule over the whole map:

In three dimensions the loop is longer. A 3 x 3 x 3 kernel has 27 offsets, and slide 105 notes that the centre (0,0,0) is usually computed separately because, at stride 1, its map is simply the identity over all points, leaving 26 irregular offsets. Because point-cloud occupancy is dynamic sparsity, the maps are built at run time for every scan, so the engine cannot precompute a fixed schedule the way it could for pruned weights.

Why this wastes most of the GPU

The baseline is correct and simple, and Tang et al. measured exactly where it loses. Two bottlenecks share the blame. First, the matmuls are non-uniform. The paper reports that map sizes for different weights can differ by an order of magnitude and that most map sizes are small, so each launch is a short, narrow matrix that cannot fill the device: MinkUNet 0.5x on an RTX 2080 Ti in FP16 reaches 8.1 TFLOP/s, only about 30 percent device utilization, with the matmul phase taking 20 to 50 percent of runtime (Tang et al., MLSys 2022, section 3). NVIDIA's own cuBLAS guidance says the same thing from the other side: when the matrices are small, the small matrix size prevents the GPU from being fully utilized, and launching them one after another produces many kernels launched in sequence (NVIDIA Developer Blog).

Second, and larger, gather and scatter are memory-bound. Each one is a random access pattern over the feature tensor, bottlenecked by GPU memory bandwidth rather than by computation resources, and together they take 40 to 50 percent of runtime. Worse, because every offset gathers and scatters on its own, the data movement is completely separated from one offset to the next: a row like f1 that appears in the maps of W(-1,0), W(0,0) and W(1,1) is fetched three times from off-chip memory with nothing reused (Tang et al., MLSys 2022, section 3).

Look at the same five features one more time, now the way TorchSparse runs them. The gather step still builds one buffer per offset, but adaptive grouping pads those buffers to a shared height, so the four small offsets become four buffers of exactly two rows each: {F0, F3} for W(-1,-1), {F1, pad} for W(-1,0), {F3, pad} for W(1,0) and {F1, F4} for W(1,1). Two of those rows are dashed: a pad, a row of zeros inserted so that every buffer has the same height. Because the four buffers now share a shape, one batched matrix multiplication (bmm, batch 4) runs all four at once. The centre weight W(0,0), whose buffer holds all five rows, runs as an ordinary matrix multiplication (mm). Partial sums PSUM 0 to PSUM 4 are scattered back with the same locality-aware access, and the two pad rows produce two blank partial-sum slots that nobody reads.

Those two dashed rows are the price. Two zero rows were multiplied for nothing so that five matmul launches became two. That is the entire idea of this concept, and the paper names it in the slide title: trading computation for regularity. The overview on slide 101 lists three techniques, locality-aware gather, adaptive grouping of the matmuls, and locality-aware scatter-accumulate. The first and last attack the memory-bound half of the runtime by reading and writing feature rows in an order that keeps neighbouring accesses close and by fusing and vectorizing them; the paper reports the cost of memory movement cut by 2.7x that way (Tang et al., MLSys 2022, abstract). The middle technique, grouping, is what the rest of this concept builds up.

The spectrum: separate, dense, grouped

Slides 102 to 104 draw seven weight offsets as seven columns of unequal height, one row per map entry, and put each execution strategy on two sliders: computation overhead and computation regularity. Separate computation, the baseline of the previous concept, launches one mm per column. Nothing is padded, so its overhead slider sits at best, but seven uneven kernels put its regularity at worst: many kernel calls, low device utilization. Dense convolution goes to the other end. Pad every column up to the tallest and run a single bmm with batch 7. Regularity is perfect, one launch, but the dashed rows now outnumber the real ones in the short columns, and overhead sits at worst. Computation with grouping refuses both extremes. The tallest column keeps its own mm. The four columns of similar height are padded by a total of two rows and run as one bmm with batch 4. The two shortest, equal already, form a bmm with batch 2. Slide 104 counts the cost: extra computation of 2 / 28, about 7 percent, for cutting seven launches to three.

How the groups are chosen

Grouping only pays when the columns in a group really are similar, so the question becomes how similar is similar enough. TorchSparse answers with two auto-tuned parameters. The first, epsilon, is the tolerance of redundant computation. Walk the offsets in kernel order and scan them once with two pointers, extending the current group while the redundant computation ratio stays at or below epsilon and starting a new group the moment adding the next offset would exceed it (Tang et al., MLSys 2022, section 4.2.3). The groups are therefore contiguous runs of weight indices, which is why the red boxes on slide 105 sit side by side along the weight axis; the tuner below and the slide 104 figure instead order the offsets by map size first, a simplification that makes the pad rows easier to see.

The second parameter, S, is a workload threshold. A group runs as bmm only if its largest map is below S; otherwise it runs as mm, because bmm improves device utilization for small workloads but has little benefit once a single matmul is already large enough to fill the GPU. The two knobs contain the whole spectrum as special cases: epsilon = 1 with S unbounded is dense convolution, S = 0 is separate computation, and epsilon = 0 with S unbounded is what the paper calls symmetric grouping. That last case is free: for an odd kernel size at stride 1 the map for offset (a, b, c) has exactly the same size as the map for (-a, -b, -c), since every pair of neighbours appears once in each direction, so 26 offsets collapse into 13 groups of two with no padding at all, worth up to about 1.2x (Tang et al., MLSys 2022, section 4.2.2).

Try the knobs yourself. The tuner below starts from exactly this example, and rerolling draws new map sizes so you can watch which groups form as epsilon and S move.

Measured: neither zero padding nor one group wins

Slide 105 is the experiment that settles the question. On the first sparse convolution layer of MinkUNet on SemanticKITTI, the paper sweeps the number of groups from 26 (no padding, the separate baseline at 1.0x) down to 1 (fully dense). Going from 26 to 13 groups gives about 1.2x, the symmetric pairing. At 6 groups the speedup peaks near 1.5x. Then it collapses: 3 groups fall below the baseline and 1 group runs at about 0.35x, worse than doing nothing. The slide annotates the two slopes as increasing regularity helps improve latency, then padding overhead hurts latency (Tang et al., MLSys 2022, Figure 7).

The two bar charts under the curve explain why the answer is adaptive rather than fixed. They plot map size against weight index 1 to 27 for a MinkUNet layer on two datasets, with red boxes marking the groups the tuner chose. SemanticKITTI maps are large, mostly in the thousands, and the boxes are narrow: 10 groups. nuScenes maps are an order of magnitude smaller, hundreds rather than thousands, so each launch is even less able to fill the GPU and the tuner groups more aggressively: 8 groups with wider boxes (Tang et al., MLSys 2022, Figure 12). The right amount of padding is a property of the data, which is why it is searched per model and per dataset instead of fixed in code.

Slides 106 and 107 isolate the matmul phase and ask a sharp question: on an RTX 2080 Ti in FP16, how much did grouping actually buy? Three strategies are compared, the separate baseline, a fixed grouping with three hand-made groups per layer type padded to their maximum, and adaptive grouping. Each is reported twice, as throughput in TFLOP/s and as normalized speedup, and the two columns disagree in an instructive way.

On SemanticKITTI (MinkUNet 0.5x) adaptive grouping lifts throughput from 8.1 to 11.9 TFLOP/s and finishes 1.39x faster. Fixed grouping is the trap: it posts a higher throughput than the baseline, 8.7 TFLOP/s, yet is slower, 0.87x. On nuScenes (MinkUNet, three frames) the trap is sharper still. Fixed grouping reaches the best throughput of the whole table, 21.1 TFLOP/s, and adaptive grouping only 16.9, yet adaptive is the faster of the two at 1.54x against 1.50x. Slide 107 gives the reason in one line: fixed grouping introduced a large amount of redundant computation. TFLOP/s counts every multiply the tensor cores performed, pad rows included. A strategy can therefore look excellent on throughput while spending that throughput on zeros. Latency counts only the clock, and the paper's caption puts it precisely: as we trade FLOPs for regularity, TFLOP/s and speedup are non-proportional (Tang et al., MLSys 2022, Table 2). Device utilization for the matmul phase rose from about 30 percent to 44.2 percent with adaptive grouping.

End to end, with locality-aware gather and scatter included, the paper reports 1.6x over MinkowskiEngine and 1.5x over SpConv across seven models and three datasets, measured on GTX 1080 Ti, RTX 2080 Ti and RTX 3090, with up to 2.16x on segmentation models over MinkowskiEngine on the RTX 3090 (Tang et al., MLSys 2022, abstract and section 5). These figures are not on the slides; the slides show only the matmul ablation, so quote them as the paper's, and note that the MLSys 2022 evaluation used desktop GPUs only. The edge numbers come with the successor.

TorchSparse++: overlap memory with computation

Grouping fixed the matmul phase, but the baseline still has a structural flaw that no grouping removes: the three phases of gather, matmul and scatter run one after another. Gather must finish before the matmul starts, and the matmul must finish before scatter starts, three separate CUDA kernel calls in each iteration of the host loop. While the memory system is busy gathering, the tensor cores wait; while the tensor cores multiply, the memory system idles. The MICRO 2023 paper calls gather-GEMM-scatter fundamentally inefficient due to the lack of overlap between computation and memory access (Tang, Yang et al., MICRO 2023, section 2.2). The remedy in TorchSparse++ is to stop treating the phases as separate kernels.

Two fused dataflows already existed in other engines, and TorchSparse++ builds on both. Fetch-on-demand, used by MinkowskiEngine, merges gather, multiply and scatter into a single kernel that loads only the input rows it needs straight into shared memory, multiplies on chip and scatters results from registers, with zero redundant computation but heavy output write traffic, 4x to 10x more than the theoretical minimum of one write per output row, since each point has 4 to 10 neighbours and every partial sum is written back individually. Implicit GEMM, used by SpConv v2, is output-stationary: each thread block owns a tile of output rows, walks the nine (or 27) weights, and fetches the matching input rows on demand, the way im2col does for dense convolution. It writes each output once and hides memory latency by pipelining, so the loads for the next tile are already in flight while the tensor cores work on the current tile (Tang, Yang et al., MICRO 2023, section 2.2 and Figure 3). That in-kernel double buffering is the overlap the slide title refers to.

The catch with implicit GEMM is a new kind of redundancy, and slide 109 shows it. Rows are output points B0 to B7, columns are the nine weights, and a grey cell is a real multiply because the input neighbour A_j exists. Threads in a warp execute in lockstep, so if any row in a thread block needs a weight, every row in that block runs it. The red cells are those wasted lockstep multiplies: 12 in the vanilla layout. Row reordering sorts the outputs by their neighbour bitmask so rows with similar patterns share a block, cutting the count to 10. Column splitting then divides the loop over the nine weights into three parts, each sorted on its own, so a row is grouped with different neighbours for different weight subsets: 8. The paper's own example runs 34 to 26 to 22 multiplies and reports that bitmask sorting can cut redundancy by up to 3x (Tang, Yang et al., MICRO 2023, Figures 5, 6 and 10). The fused kernel keeps its overlap; these two reorderings trim what it pays for it.

Writing such kernels by hand is expensive. SpConv v2 needed more than 40,000 lines re-implementing CUTLASS. TorchSparse++ therefore contributes two components, both drawn on slide 108. The Sparse Kernel Generator produces the fused, pipelined kernels automatically at under a tenth of that engineering cost, adapting dense GEMM templates to sparse gather (dense-to-sparse adaptation) and static shapes to run-time maps (static-to-dynamic adaptation). The Sparse Autotuner widens the design space (gather-GEMM-scatter with grouping, fetch-on-demand, implicit GEMM, and their tile and split parameters) and tunes a configuration per group of layers rather than per layer, so the search stays cheap (Tang, Yang et al., MICRO 2023, sections 3 and 4).

Reading the A100 column the other way round, TorchSparse++ is 2.9x faster than MinkowskiEngine, 3.3x faster than SpConv 1.2.1, 2.2x faster than TorchSparse and 1.7x faster than SpConv 2.3.5, the four numbers the paper's abstract reports. The Orin column is the one that matters for this course: on the Jetson Orin edge GPU, the strongest baseline, SpConv 2.3.5, sits at 0.80, so TorchSparse++ is 1.25x faster on average, with a consistent 1.3x to 1.4x on detection workloads (Tang, Yang et al., MICRO 2023, section 5). The chart also shows training columns and a TF32 column; the same engine covers all of them.

One more step remains beyond software. Every engine in this part still spends time building the maps themselves, a hashing or sorting problem the GPU is not shaped for. Part 13 moves that step into hardware with PointAcc, whose mapping unit produces the tuples with merge sort.