SZ sequences prototype for OpenQMC

[wantonsushi]

@fpsunflower @joshbainbridge

Following the discussion in #97, I wrote a standalone prototype for SZ sequences and a benchmarking tool to explore the timings vs Sobol at 16, 32, and 64 bits. It shows that SZ is slower than 16-bit SIMD Sobol, but it's increasingly faster at 32 and 64 bits (scalar only, no SIMD for 32/64 bit Sobol was implemented here).

Files

Files are not wired into the library. They compile against the installed headers (-I include) and produce an oqmc::SzSampler beside oqmc::SobolSampler.

• sz.h -- the SzSampler. Same Owen scramble / shuffle / rotate wrapper as SobolSampler; only the per-dimension evaluator differs. Includes the input reversal hoist Kulla suggested in the PR discussion (reverse the shared index once and reuse it across all dimensions).
• sz_directions.h -- the Sobol and SZ direction matrices and closed-form programs at 16/32/64 bits, plus the constructions used to verify them.
• sz_benchmark.cpp -- correctness, step counts, timings.

NB: I had to add a .txt postfix to the .h files because GitHub doesn't support uploading .h files, which seems silly...

To build:

g++ -std=c++17 -O3 -march=native -Wall -Wextra -I include sz_benchmark.cpp -o sz_benchmark

Step counts

Shift-mask-xor steps per dimension, re-derived live by factoring each direction matrix:

bits	Sobol d1/d2/d3	SZ d1/d2/d3
16	4 / 12 / 10	3 / 9 / 9
32	5 / 22 / 15	4 / 12 / 12
64	6 / 44 / 21	5 / 15 / 15

This shows Sobol's worst dimension grows linearly whereas SZ's grows logarithmically. At 64-bit it's 15 steps versus 44.

Timings

Full 4D draw, SZ vs Sobol, ns per draw. The 16-bit row uses OpenQMC's SIMD path. OpenQMC has no 32/64-bit path, so I implemented those in scalar. SIMD is untested at those widths and would move the numbers.

bits	sobol	sz	speedup
16 scalar	45.70	44.00	$1.04\times$
32	63.10	55.90	$1.13\times$
64	92.40	59.25	$1.56\times$
16 SIMD	35.45	41.58	$0.85\times$

Timings on a 13th-gen Intel i7-13700H

Correctness

• The generated points are perfectly stratified the way the SZ paper requires: every consecutive dimension pair is a base-4 (0, 2)-sequence and dimensions 0–3 form a (0, 4)-sequence [Ahmed et al. 2025, Fig. 1]. Verified on the sampler's output over the first $4^7$ points. This confirms it's genuinely the SZ sequence.
• All 18 closed-form programs (Sobol + SZ, dims 1-3, at 16/32/64 bits) reproduce their direction matrix bit-for-bit.
• oqmc::szReversedIndex matches the SZ construction on all $2^{16}$ indices $\times4$ dims, and dimension 0 matches oqmc::sobolReversedIndex.

Let me know what you think!

Edit: added a net-property check confirming the output is genuinely SZ. re-uploaded sz_benchmark.cpp

[fpsunflower]

Great stuff!

As a further easy optimization, you can look into eq (13) in the paper. Because successive pairs of matrices all form a (0,2) sequence, they obey the following equation:

    P * SZ[d] == SZ[d ^ 1];

Where SZ[d] is the matrix for the d'th dimension and P is the canonical Pascal matrix. SZ[0] = I and SZ[1] = P which obey the equation trivially because P is its own inverse: P * P = I. But more useful is that SZ[3] = P * SZ[2].

This means that once you evaluate SZ[2] for a given index, you can get SZ[3] for the same index more cheaply by just using the optimized Pascal matrix eval (3-xor-masked-shifts) on that result. This requires refactoring the code a bit more, so that dimension 3 can depend on the result of dimension 2.

A bit more advanced is equation (41) in the SZ paper. This saves only a handful of operations at 32-bits (haven't checked 16-bit or 64-bit). The exact speedup may differ based on the architecture. On my Nvidia GPU, eq (41) seems faster, on my laptop (Apple M2 GPU) the straightforward diagonal evaluation (as you wrote up here) is slightly better. I do like the simplicity of the current approach, so its definitely fine for a first pass.

Will be very curious to see what results you guys get with this construction in practice and if you notice any particular benefits convergence wise compared to Sobol. In my testing, the difference was fairly subtle and it really depends on what your renderer uses dimensions (2,3) of the pattern for in practice.

[fpsunflower]

Visualizing the 2d pairs of dimensions and their FFTs (here averaged over 256 random copies of 1024 points) shows the difference between Sobol and SZ is slight, but the FFT does gets nicely cleaned up for the higher pairs in the SZ case. This diagram also shows that pair (2,3) has the same power spectrum as (0,1), as advertised.

Sobol (using the Kuo matrices for dimensions 2 and 3):

SZ:

[wantonsushi]

Thanks for taking the time, the extra optimizations, and the FFT plots. I'm learning a lot from this.

I added the eq (13) chaining to the benchmark. Dim 3 drops from 9/12/15 steps to 3/4/5 at 16/32/64 bits, and I verified the chained draw is bit-exact with the direct one.

Interestingly, chaining only helps at 32-bit and 16-bit SIMD; it's about neutral at 64-bit and slightly worse at 16-bit scalar. I guess this is because the chain requires dim 3 to be performed sequentially after dim 2?

Here are the timings I got, ns per 4D draw, same machine as before:

bits	sobol	sz	sz chained
16 scalar	45.70	44.00 ($1.04\times$)	45.05 ($1.01\times$)
32	63.10	55.90 ($1.13\times$)	52.19 ($1.21\times$)
64	92.40	59.25 ($1.56\times$)	59.75 ($1.55\times$)
16 SIMD	35.45	41.58 ($0.85\times$)	38.04 ($0.93\times$)

I'll give the factored evaluation a try too. On convergence, when I put a PR together I'll run the library's provided error plot tools for sz vs sobol, extended to the (2,3) projections.

Once @joshbainbridge shares what he'd like the PR to look like I'll open it with these optimizations included. Thanks again!

[joshbainbridge]

Hi @wantonsushi, this is great work. Thanks for opening the discussion.

I think we should move forward with the PR. The improved quality of the projections is the big draw for me. And the performance scaling with precision is a fantastic addition.

As for the implementation options. Now that the gap between vectorization and non-vectorization has further closed, this might be a good opportunity to simplify the code. I would be tempted to drop the SSE and AVX implementations, and also go with just the szReversedIndexShared() variant (just the !defined(CUDA_ARCH) branch).

This would have a few benefits. Code becomes easier to maintain, and it leaves the door open to further caching of the index. Users who are using the library as plug-and-play, are likely not making use of SSE or AVX, as these options are off-by-default. We could even remove those build options and the variants on the CI to make the whole experience easier.

But open to other opinions. I'm still on the fence with that. Any thoughts @fpsunflower?

[fpsunflower]

Dropping SSE/AVX definitely makes sense if there is no major performance difference anymore.

[wantonsushi]

Sounds good. So for the PR I'll do scalar only, just the szReversedIndexShared() variant with the reversal hoisted out of the draw, plus the eq (13) chaining for dim 3?

A few clarifications:

• If I understand correctly, you want sz sampler to replace the current sobol path, and not be a new sampler beside sobol? To clarify, the timings I got was 16 bit SIMD sobol ~35 ns vs 16 bit scalar sz ~45ns.
• Do you want a 32-bit implementation in this first PR, or keep it 16-bit and revisit? It looks fairly self-contained (the scramble/shuffle utils are already 32-bit, and lattice already eats all 32 shuffled bits)?
• Regarding shared-only: CUDA then takes the right-shift form, which measured 0.83x on my GPU for Sobol in Speed up scalar Sobol with closed-form evaluation #97. The shared reversal saves 3 reversals per 4D draw, so I haven't measured exactly where the shared form lands on GPU, but I'd expect it's still a net loss there. Are you sure we don't want to keep the left-shifts on CUDA?

[joshbainbridge]

These are good questions, I’ve been thinking about this a little more today. Here is the current lay-of-the-land (if only for my own reasoning) and some thoughts.

Implementations on main branch:

• CPU SSE, Direct matrix-vector multiply
• CPU AVX, Direct matrix-vector multiply
• CPU ARM, Direct matrix-vector multiply
• CPU Scalar, shift chain + pre-inverse
• GPU, shift chain + post-inverse

Measured results for this change take maximum advantage of sharing a pre-inversed index across 4 dimensions. The benefit here though might be less in real-world scenarios due to sample calls sometimes requesting 1, 2, or 3 dimensions. Still a good change to make.

GPU might also not be benefiting much from sharing the pre-inversed index, as the inversion is very cheap on an NVIDIA GPU, just a single operation. And as you point out, the left shifts are causing a notable slowdown. So I’d agree, best to keep a dedicated variant for this.

Upon further consideration, I’m not sure if the index can be optimized beyond what you have shown in this draft. I’d originally considered that it might be possible to compute the inversion only when the index is incremented. But calling reverseAndShuffle() takes a seed value which is based on the domain pattern ID. So the index will be different for each domain. Draw calls are typically only done once for each domain instantiation, and so there is little to no optimization to be made beyond computing the index directly in shuffledScrambledSobol().

I was thinking that if we could further cache the index, it might make sense for the GPU to pre-inverse the index, but that now appears to be not a viable option.

As for the SIMD paths. This is where it might make sense for us to drop them now the gap has got closer on the timings. We mostly gain a reduction in complexity for doing this. That is a reduction in code to maintain, as well as fewer build options for a user to handle. On top of that, I’ve got some ideas on how we could cache the linear function, which could work well for the CPU (but not the GPU), and make the SIMD paths redundant. I can follow up on that in a separate thread.

Lastly, on the extension to 32-bit. There is a connection between the maximum sample count that we support, and the precision of the Sobol generator. We technically support just up to 16-bit indices (65k maximum samples extended to 4b with padding). This is then taken advantage of in how the state is packed. Without moving to 32-bit indices, I don't think there is any advantage to increasing the output precision, as we would then have more representable states on output than we do on input. Although that might change in the future, which we would then be in a good position to further leverage these changes.

TL;DR: Let’s move forward with just two paths, one CPU and one GPU, dropping the SIMD variants. We can keep the SIMD build options and infrastructure for now, just in case we need it (this can be removed in a future PR prior to release). I don’t think there is enough differentiation between an SZ sampler and a Sobol sampler to keep both, seeing as they are both 02 sequences. So I’d say that this should replace the current Sobol sampler implementation, keeping the Sobol name for the time being (as that will be a breaking change for users).

Thank you again @wantonsushi! Really enjoyed reading your discovery here.