Speed up scalar Sobol with closed-form evaluation

[wantonsushi]

Hello OpenQMC maintainers!

This PR implements the Ahmed 2024 paper mentioned in issue #67 for the scalar path of sobolReversedIndex. I used the existing benchmark tool to check timings and generate to confirm the output is unchanged.

Changes

• include/oqmc/owen.h: the scalar path uses the closed form.
• src/tools/cli/matrices.cpp: derives the steps from the generator matrices and prints them, the same way it already prints directions[].
• src/tests/owen.cpp: OwenTest.SobolReversedIndex checks the result against the direction matrices for every 16-bit index and dimension.

I also tried to implement a SIMD version (sobol_simd_experiment.cpp). It's not in this PR since it doesn't beat the existing SIMD, but I'm sharing it so you can reproduce the SIMD timings below.

Timings

Using benchmark sobol samples, median of 9 runs, in microseconds:

target	old	new	speedup
CPU scalar	60034	56482	1.06x
GPU (RTX 4070)	128830	124873	1.03x
CPU SSE	47545	62700	0.76x
CPU AVX	47627	62856	0.76x

The speedup is small because generation is only part of a draw, along with scrambling and state hashing. The SIMD implementation fits one dimension per lane (four lanes). The existing SIMD path is faster because it packs all sixteen matrix columns into a wider register, so I left the SIMD paths alone.

Verification

• generate sobol output is byte-identical to before on scalar, SSE and AVX
• Tests pass (168/168), clang-format and clang-tidy clean

[joshbainbridge]

Hi @wantonsushi. I've just given the PR a scan and it is looking good. This is great work, thank you.

Very interesting to see the comparison. I'll follow up with a closer read of the code details sometime tomorrow.

[joshbainbridge]

These changes are stellar, thank you. I would be happy to add this into the project. I've added some minor comments to the PR.

There are also a few additional of files we should update:

• Adding a change log entry to CHANGELOG.md file.
• Referencing the technique and the author in both include/oqmc/sobol.h:103 and README.md:979. Just a short sentence would be totally fine.

Once we are happy with the PR, next steps will be rebase onto main, and squash the commits down to a single commit. You will need to digitally sign the commit. Info on how to digitally sign, and commit message format can be found here:

• Digital signing: https://docs.github.com/en/authentication/managing-commit-signature-verification/about-commit-signature-verification
• Commit message format: https://github.com/AcademySoftwareFoundation/openqmc/blob/main/CONTRIBUTING.md#developer-workflow

[joshbainbridge]

Everything is looking good to me. Thank you for adding those additional changes. I think we are now ready for next steps:

• Rebasing onto main
• Squashing commits
  • Authoring Git message (brief, why, what)
  • Make sure to sign the commit with -S

Once that is done, the pipeline should allow us to merge the PR 🚀

[wantonsushi]

Pushed. Should be good to go now, let me know if anything else is required. Thanks for the opportunity to contribute! Might do more in the future :)

[fpsunflower]

This is cool! I hadn't figured out how to extend Ahmed's construction to arbitrary matrices as you've done here.

If I am reading the code right though, aside from the Pascal matrix (dimension 1) which has a compact O(log(n)) evaluation, we should expect to get roughly O(n) steps for an arbitrary matrix? So if we wanted to extend the code to 32-bit indices we would expect roughly twice as many steps?

Ahmed's follow up paper on SZ sequences has the nice property that higher order terms also have a factored representation so they are all O(log(n)) to evaluate, which allows extending even to 64-bit sequences (potentially useful if you are sharing a single sequence across pixels). Also since all pairs of dimensions are themselves (0,2) sequences, you can benefit from the super fast pascal matrix evaluation for every other term.

[wantonsushi]

Hi Chris. Exciting to see you comment here. Is it weird to say I'm a big fan of your work at SPI?

Yes, you're right: dim 1 is O(log n), whereas dims 2/3 are O(n). I verified this, for dim 2:

num bits	num steps
16	12
32	22
64	44

so roughly doubling like you said.

The SZ construction seems to make sense: every dim is a Pascal matrix P(a), and the S-P-Z decomposition evaluates it as the shared Pascal plus precision-independent corrections, so all dims stay O(log n) even at 64-bit (with every other dim reducing to plain Pascal)? I guess the catch is it's a different sequence past dim 1, so for OpenQMC, this isn't a simple drop-in replacement. @joshbainbridge what do you think?

I think I'll try making an SZ prototype for OpenQMC and benchmark it against the current path when I get a chance, I will follow up. Thanks for the comment.

[fpsunflower]

I played a bit with the code and got the same results as you for the Kuo Sobol matrices which OpenQMC is using. The SZ matrices have a slightly less random structure, so more stuff cancels out and you get only 12 steps required for 32-bit input for dimensions 2,3 (which is fairly close to what you need with the factored implementation and is easier to reason about). I'll do some tests to see if it makes any big difference in practice.

One small optimization, if instead of:

  index ^= (index & mask) << shift;

You do this:

   index ^= (index & reversebits(mask)) >> shift;

you can skip the final call to reversebits16() which should be a bit faster (of course the reversing of the mask is precomputed, and therefore free).

[wantonsushi]

@fpsunflower Thanks for the suggestion. I applied it, but ran into an issue. The conjugated form measured 0.83x on GPU for me. To my understanding, right shifts contend with the logic pipe, while the original left shifts lower to multiplies on the multiply pipe. I tried right-shift-as-multiply-high, but __umulhi expands on Ada and I measured 0.59x. So I kept the old method for the GPU path. Maybe there is something I overlooked.

For CPU though, works great. With benchmark sobol samples, median of 9 runs: 1.14x speedup on CPU scalar, with byte-identical output.

I've left this as an independent commit so you and @joshbainbridge can easily verify the changes. I'll squash it to one commit if we're happy with it.

[fpsunflower]

Interesting. That difference may depend on the GPU architecture in practice, I hadn't considered the fact that left-shifts are potentially faster.

Still, its better if the reversebits() happens on the input since its then shared across all 4 dimensions. The code probably should be restructured a bit to make this more obvious, again probably better left for another PR.

I'll also point out that in OpenQMC, reverseAndShuffle() calls reversebits() on the same index each time. This can be pulled out of drawSample(). In the end, you should only have 1 + num_dim calls to reversebits() per drawSample call (one for the index, and one per output dimension).

Also, while the 16-bit implementation is sufficient for a randomized seed (all the index bits are used for the sample index), when using the Z-sampler construction, part of the index gets used for the pixel coordinates and you need a fully 32-bit evaluation. With the optimized method used here, I don't think it should be that much slower than the 16-bit version and would scale to higher sample counts.

[joshbainbridge]

Thanks @fpsunflower, really thoughtful input.

@wantonsushi on the subject of Ahmed's follow-up paper on SZ sequences, although I've not yet read through it in detail, the advantages would be a notable improvement. I would be happy for us to explore that as a follow-up PR and accept any changes to the resulting sequence. Thank you for offering. We would also need to re-generate the blue noise tables, but I'll be able to help out with that part.

As for the other note on reversing the masks. It makes sense that the mask reversal could be precomputed. But does that not then require us to invert the index prior to the XOR operations, which in-turn nullifies the compute saving?

[fpsunflower]

You are right that reversing the mask and shift direction just lets you changes where the reversebits() call needs to happen (either before doing the xor-shifts, or after). It doesn't change the fact that you need this call.

However, since before and after the sequence evaluation there is an owen-scramble call which itself contains reversebits() calls around the hash function, there is opportunity for cancellation either way.

To me it makes more sense to cancel the reversebits() call on the output side since that one is proportional to the number of dimensions. In the end, you should only have one call on the input side and one per dimension (after hashing). This is already what OpenQMC is doing, its just that the reversebits() call is hidden for dimensions 1,2,3 because the sequence is evaluated the other way around. Once you use Ahmed's xor-shift factorization, the reversebits() call is better to consider explicitly.

In pseudo-code, assuming a 4D sampler like in OpenQMC:

float4 draw_sample(uint index) {
    uint shuffled_index = reversebits(owen_hash(reversebits(index)));
#if 0
    // flip then xor-shift (allows cancellation with 1 `reversebits()` above)
    uint4 sobol_value = sobol_sequence(reversebits(shuffled_index));
#else
    // xor-shift then flip (allows cancellation with 4 `reversebits()` below)
    uint4 sobol_value = reversebits(sobol_sequence_inv(shuffled_index));
#endif

    uint4 result = reversebits(owen_hash(reversebits(sobol_value)));

    return bits_to_uniform01(result);
}

Like I mentioned above, the very first reversebits(index) call should happen when constructing the sampler, not on every draw_sample call since index generally stays constant unless you are splitting the path. That brings things down to 1+num_dimensions reversebits() calls per draw_sample.

On processors (x86 based) that don't have a native reversebits() instruction - its possible that Ahmed's factorization trick isn't worthwhile compared to 16-bit SIMD sequence evaluation, but for 32-bits and above I think it will definitely still pay-off (though I haven't tried).

[joshbainbridge]

@fpsunflower thank you. That now all makes sense to me. I think we can get there incrementally.

An initial change could be taking advantage of inverting the index just in the shuffledScrambledSobol() function. This wouldn't be that large a restructure, and self-contained to include/oqmc/owen.h.

A follow-up to that could be to then move this further up the call graph and cache the value in the state object. But this would need some thought as it would be a larger restructure. I'd agree though, worth doing.

As for this PR. @wantonsushi I think what you have done is great, including the additional commit, as this sets us up for further optimizations.

Looks like the format check just failed on the code comment. But if you want to squash, format and push that code; that should be all we need to get this merged.

[joshbainbridge]

All merged. Great contribution @wantonsushi!