SZ sequences prototype for OpenQMC #103

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):
sampler_2dpts_r_sobol_kuo 1024spp

SZ:
sampler_2dpts_r_sobol_sz_diag 1024spp

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.