Nibble tabulation for the Sobol draw

[joshbainbridge]

Hello :) I've been thinking about how we evaluate our generator matrices, and I want to share an idea on how we might optimise the sobol draw using tabulation. I've often wondered about the tradeoff between sobol and pmj, the pros vs cons of tabulating the sequence, and how this is likely problematic when considering cache-misses.

Tabulation

All ways of returning a sequence based on a generator matrix are a trade-off between memory and work per query. Laid out on one axis:

All three compute the same thing. They differ only in how much of the work is done ahead of time versus per draw. In this discussion I'd like to propose that we investigate the middle option.

Current implementations

OpenQMC already implements the extremes.

owen.h computes Sobol values live. The scalar path is the shift chains (diagonal factoring, Ahmed 2024), and the SIMD paths are column accumulation with AVX, SSE and ARM variants, plus a separate left-shift chain for CUDA. Five implementations of one 16x16 bit matrix multiply.

pmj.h sits at the far right: stochasticPmjInit fills a samples[65536][4] table, 1 MB, and every draw is a lookup through lookup.h.

Idea

The reason why this middle option typically isn't feasible is that it only exists for linear transforms. Our Sobol draw is multiplication by a matrix over GF(2), which means it distributes over xor:

f(a ^ b) = f(a) ^ f(b)

Split the input index into its four nibbles and the whole transform becomes four tiny table lookups xored together. Each nibble position has only 16 possible values, so each table has 16 entries. Better still, all four output dimensions can be packed into one 64-bit entry, since xor acts independently on each 16-bit lane. And since reverseBits16 is also linear, the bit reversal folds into the tables for free.

Total: 4 tables x 16 entries x 8 bytes = 512 bytes, generated from the matrices at build time (the same matrices.cpp cli tool that emits the shift chains today could emit these instead).

This is a known trick with many names: the Method of Four Russians in GF(2) linear algebra, slicing-by-N in the CRC world, tabulation hashing elsewhere.

Why is this optimal?

The exhaustive table's problem is the memory footprint. 1 MB is 16,384 cache lines, and a random index lands on a different one every draw, so in a real render almost every access misses L1 and competes with actual work for L2. The nibble tables are 8 cache lines total. Every draw touches the same 8 lines, so after warmup they are pinned in L1 permanently and every load is a 4-5 cycle hit (approximately). The randomness of the index only selects within a line, never which lines are resident.

And unlike the shift chains, there's no serial dependency: the four loads issue in parallel and the xor reduction is two deep. The chain's ~26 dependent rounds become ~13 cycles flat, for all four dimensions at once.

Sketch

// Build once: push each lone nibble value through the existing scalar
// transform, packing the dimensions into 16-bit lanes of each entry.
// entry = dim0 | (dim1 << 16) | (dim2 << 32) | (dim3 << 48)
std::uint64_t nib[4][16];

// Precompute: table computed offline and embedded in source.

for(int k = 0; k < 4; ++k)
	for(int n = 0; n < 16; ++n)
	{
		const auto v = std::uint16_t(n << (4 * k));
		nib[k][n] = std::uint64_t(sobolReversedIndex(v, 0))
		          | std::uint64_t(sobolReversedIndex(v, 1)) << 16
		          | std::uint64_t(sobolReversedIndex(v, 2)) << 32
		          | std::uint64_t(sobolReversedIndex(v, 3)) << 48;
	}

// Draw: one call transforms the index for all four dimensions.

// Inverse + LK permutation here...

const std::uint64_t r = nib[0][x & 15]
                      ^ nib[1][(x >> 4) & 15]
                      ^ nib[2][(x >> 8) & 15]
                      ^ nib[3][x >> 12];

const auto d0 = std::uint16_t(r);
const auto d1 = std::uint16_t(r >> 16);
const auto d2 = std::uint16_t(r >> 32);
const auto d3 = std::uint16_t(r >> 48);

// LK permutation + inverse here...

That could replace the scalar chains and all three SIMD paths, and the standalone bit reversal in the draw.

Initial tests

I passed this onto Claude and asked it to test the idea against the current owen.h scalar path exhaustively (all 65536 indices x 4 dimensions). It benchmarked 100M serially dependent draws, one index at a time, all four dimensions per draw. All three methods produced identical checksums. Xeon 2.8 GHz, gcc -O2:

Method	ns per 4-dim draw	~cycles
Shift chains (current scalar path)	18.4	52
Nibble tables, 512 B	4.6	13
Exhaustive table, 1 MB (pmj style)*	12.5	35

*note the exhaustive number is likely unrealistic and in reality much worse performance. The benchmark touches nothing but the table, so much of the 1 MB stays cached. In production that isn't the case and we would expect it to degrade; while the 512 B used to store the nibble tables would not.

What about pmj?

Should both samplers be tabulated using this type of table lookup? The key here is that the Sobol table compresses because the transform is linear; the pmj table cannot compress because it's random by construction. The pmj values contain genuine entropy from the PCG stream, so the 1 MB footprint is the smallest form I've managed to find. And even though pmj shows good results in benchmarks, I've always considered its poor cache efficiency to likely mean that in real world use the performance is unfavourable.

The other advantage of pmj was licensing, although given the licensing granted for the OpenQMC project, this is likely a moot point. It could be that a switch to this approach would make the pmj sampler redundant.

Open questions

• GPU path. Divergent per-thread lookups into constant memory serialise, so the shift chain may still be the right call under CUDA, where warp scheduling hides its latency anyway. Needs a device-side benchmark before touching it.

If there's interest I'll put together a PR with the table generation added to the matrices cli tool, the new draw path behind the existing arch defines, and the verification test.

@wantonsushi @fpsunflower I'd be interested in getting your thoughts. It was your recent work @wantonsushi that got me thinking about this.

FYI @mr-matthew-jones this stuff is up your street, you might also find it interesting.

[fpsunflower]

Interesting idea!

It doesn't seem like it would scale well to 32-bits though whereas the shift-chain approach is just a few extra operations thanks to the log-scaling of the diagonal factoring on dimensions 1 and 3.

It might also be worth checking the cost when only generating a 1,2 or 3 dimensions instead of 4? Its less clear to me how much code can be pruned in those cases since you are taking advantage of the fact that a uint64_t value holds the 4 uint16_t outputs in one register. In the other method, the unused dimensions drop out more directly.

Even with the small table here, cache misses could be a concern when inside a real application that puts other pressure on the cache. I'm not sure of a good way to measure this though as profiling inside a full renderer likely looses the sampling cost in the noise.

Also curious to see results on GPU. It would interesting to look at the raw SASS instructions on nvidia because I am not sure how 16-bit and 64-bit operations are handled (and PTX is not what actually gets run on the device).

But assuming the benchmarks hold up, I'm for whatever is faster. As the old saying goes "profiling gives you a leg up over experts that don't need to" ;)

Given how many device types are out in the wild, it might be good to let the end-user pick a method at compile time. Something like:

#ifndef OQMC_SOBOL_EVAL_METHOD
  // user did not pick a method, default to what we think is best
  #ifdef __CUDA_ARCH__
     #define OQMC_SOBOL_EVAL_METHOD 1 // shift-masks for cuda device code
  // ... other heuristics here as we find them ...
  #else
    #define OQMC_SOBOL_EVAL_METHOD 0 // nibble-table for CPU code
  #endif
#endif

Then all anyone needs to compare methods is redefine that macro themselves before including the OpenQMC header and benchmark in their specific situation.

[wantonsushi]

I gave the method a try and managed to reproduce similar results.

I wired the OpenQMC samplers into a pbrt-v4 fork for my blue noise work: https://github.com/wantonsushi/pbrt-v4-openqmc. I thought it'd be useful here to try measuring the performance in "real world use".

From the fork, I timed 5 scenes, 5x each, and nibble was faster in 4 of 5, 1.01x to 1.02x. Exhaustive was faster on only 2 of 5 scenes, the rest had at least one rep go the other way so I'd call it noise there. Sobolbn vs pmjbn, pmjbn was faster in all 5, 1.02x to 1.03x. Nibble-sobol vs pmj, pmj was still faster in all 5 around 1.00x to 1.01x, which is not what I expected given the cache argument.

Also, in the pbrt fork, I wired up the samplers in a way so that, even though pbrt only asks for 1D and 2D samples, we draw all four, cache the unused dimensions, and serve the following requests from the cached results. If done in this style, pruning for 1/2/3 dimension samples as @fpsunflower mentioned doesn't really become a concern. But OpenQMC's trace.cpp example doesn't work this way, and I wasn't sure if this is the correct way to use the samplers?

Regarding GPU, I think the serialization is a property of constant memory rather than of the table (i.e., a warp reading 32 different nibbles out of __constant__ costs 32 accesses). The same 512 bytes as an ordinary __device__ array read through the read only cache, or copied into __shared__ once per block, wouldn't behave that way I don't think?

Also extra note: it seems drawSample<4> is only about a sixth of the cost, the rest is reverseAndShuffle once plus a scrambleAndReverse per dimension. So it made me wonder how to optimize the other two routines. In particular, the per-dimension scrambles don't depend on each other, so we have four independent hashes rather than a dependent chain. The obvious idea is maybe SIMD could be worth revisiting for this?

Will try to open a PR for SZ in the coming days.

[fpsunflower]

I think caching partial results is risky because it becomes less predictable what dimensions of the sampler you are getting. It depends on how the calls are sequenced in the renderer, which can make the behavior less predictable. For instance you might disable a feature for debugging and then the random numbers change because the path taken through the code changes. The nice thing about the padded replication approach that OpenQMC uses is that you have predictable quality for each dimension you use. If you use consistent values when creating a new domain, you can guarantee consistent random numbers, even if parts of the renderer are turned off which is really helpful when debugging a big production render.

scrambleAndReverse is definitely worth optimizing - I have some suggestions to make, but SIMD can definitely help on the CPU. Also the way the seed is changed between dimensions (currently a byte rotation) could be tuned as well (CPUs have dedicated rotate instructions, but I don't think GPUs do).

There's a delicate dance between introducing enough randomness in the seed and the exact shape of the scrambling hash. We should probably braintstorm this in another ticket.

[wantonsushi]

Just to quickly note: I fixed my pbrt fork to draw a fresh OpenQMC domain per sampler request (instead of caching values). To be able to test the performance of dimensions 3/4, since pbrt only ever requests 1d/2d samples, I extended the code to ask for 3d/4d samples wherever the sampling decision is truly that wide (e.g. motion blur). I haven't tested super thoroughly but it should be an improvement.

[joshbainbridge]

That's a great quote @fpsunflower! Haven’t heard it before. Will be using that.

I'd agree that it is very hard to reason about the performance of this without benchmarking it. Thank you @wantonsushi for doing all the legwork and testing that out with PBRT 🙏 If you do get those results back with taking the samples dynamically, it would be interesting to see the data. I could also try putting together a version of OpenQMC to A/B test this on a heavy production shot at Framestore. That would be another data point for CPU.

Taking a step back. I agree that we are definitely squeezing the very last optimization out of the sobol function. The reverseAndShuffle() and scrambleAndReverse() are now where all the real cost is. It would be interesting to see if these could be optimized. Sounds like Chris has some ideas.

PMJ sampler is still a lot faster, but only because it doesn't perform a scrambleAndReverse() on output of each dimension, and instead does a basic XOR to scramble (random-digit scrambling, Kollig and Keller). Both have the improved integration rate from Owen scrambling (random error cancellation) but PMJ gets this randomization from its initial stochastic construction, and so can rely on a cheap XOR at runtime. Although it is this same stochastic construction that makes it the approach here with the nibble tables infeasible, and the full table lookup necessary, which I suspect is not great for cache performance.

The quality is also inferior, as the basic XOR still adds structure that is visible in the spectral transform. And soon it won't have the SZ features. All that to say, I'd really like to sunset PMJ if we can make the Sobol always a better tradeoff.