Speed up scalar Sobol with closed-form evaluation

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 ### Changed
+
+- Sped up scalar paths of `oqmc::sobolReversedIndex` using the construction method from Ahmed 2024.
+
 ### Deprecated
 ### Removed
 ### Fixed

README.md

@@ -975,7 +975,8 @@ The implementation uses an elegant construction by Burley [^1] for an Owen
 scrambled Sobol sequence. This also includes performance improvements such as
 limiting the index to 16 bits, pre-inverting the input and output matrices, and
 making use of CPU vector intrinsics. You need to select a `OPENQMC_ARCH_TYPE` to
-make use of the performance from vector intrinsics for a given architecture.
+make use of the performance from vector intrinsics for a given architecture. The
+scalar path uses the closed-form 2D Sobol construction by Ahmed [^5].
 
 <picture>
   <source media="(prefers-color-scheme: light)" srcset="./images/diagrams/sobol-design-light.png">
@@ -1470,6 +1471,10 @@ https://dx.doi.org/10.1007/s11075-011-9482-5
 Lessons Learned and Improvements when Building Screen-Space Samplers with Blue-Noise Error Distribution. ACM SIGGRAPH 2021 Talks.  
 https://dx.doi.org/10.1145/3450623.3464645
 
+[^5]: Abdalla G. M. Ahmed. 2024.  
+An Implementation Algorithm of 2D Sobol Sequence: Fast, Elegant, and Compact. EGSR.  
+https://dx.doi.org/10.2312/sr.20241147
+
 [^6]: Melissa E. O'Neill. 2014.  
 PCG: A Family of Simple Fast Space-Efficient Statistically Good Algorithms for Random Number Generation.  
 https://www.pcg-random.org

include/oqmc/owen.h

@@ -48,6 +48,107 @@ OQMC_HOST_DEVICE inline std::uint16_t sobolReversedIndex(std::uint16_t index,
 	assert(dimension >= 0);
 	assert(dimension <= 3);
 
+#if defined(OQMC_ARCH_SCALAR)
+
+	// Each matrix factors into shift-mask-xor steps (Ahmed 2024, eq. 18),
+	// dimension 1 being the Pascal matrix (Listing 19). Steps emitted by the
+	// matrices cli tool in src/tools/cli/matrices.cpp.
+
+#if defined(__CUDA_ARCH__)
+
+	// Prefer left shifts on GPU since they are more efficient than right
+	// shifts. Reversal must be applied to the result.
+	switch(dimension)
+	{
+	case 1:
+		index ^= static_cast<std::uint16_t>((index & 0x00ff) << 8);
+		index ^= static_cast<std::uint16_t>((index & 0x0f0f) << 4);
+		index ^= static_cast<std::uint16_t>((index & 0x3333) << 2);
+		index ^= static_cast<std::uint16_t>((index & 0x5555) << 1);
+		break;
+	case 2:
+		index ^= static_cast<std::uint16_t>((index & 0x0003) << 14);
+		index ^= static_cast<std::uint16_t>((index & 0x0004) << 13);
+		index ^= static_cast<std::uint16_t>((index & 0x000f) << 12);
+		index ^= static_cast<std::uint16_t>((index & 0x0033) << 10);
+		index ^= static_cast<std::uint16_t>((index & 0x0055) << 9);
+		index ^= static_cast<std::uint16_t>((index & 0x0030) << 8);
+		index ^= static_cast<std::uint16_t>((index & 0x0303) << 6);
+		index ^= static_cast<std::uint16_t>((index & 0x0505) << 5);
+		index ^= static_cast<std::uint16_t>((index & 0x0cf3) << 4);
+		index ^= static_cast<std::uint16_t>((index & 0x1111) << 3);
+		index ^= static_cast<std::uint16_t>((index & 0x0f0f) << 2);
+		index ^= static_cast<std::uint16_t>((index & 0x6666) << 1);
+		break;
+	case 3:
+		index ^= static_cast<std::uint16_t>((index & 0x000f) << 12);
+		index ^= static_cast<std::uint16_t>((index & 0x0030) << 10);
+		index ^= static_cast<std::uint16_t>((index & 0x0050) << 9);
+		index ^= static_cast<std::uint16_t>((index & 0x007f) << 8);
+		index ^= static_cast<std::uint16_t>((index & 0x03f0) << 6);
+		index ^= static_cast<std::uint16_t>((index & 0x0410) << 5);
+		index ^= static_cast<std::uint16_t>((index & 0x07e0) << 4);
+		index ^= static_cast<std::uint16_t>((index & 0x1c71) << 3);
+		index ^= static_cast<std::uint16_t>((index & 0x1c71) << 2);
+		index ^= static_cast<std::uint16_t>((index & 0x4924) << 1);
+		break;
+	default:
+		break;
+	}
+
+	return reverseBits16(index);
+
+#else
+
+	// Reversed masks and right shifts put the reversal on the shared input.
+	// This optimization hoists the reversal out of a draw
+	// (Chris Kulla, PR #97).
+	index = reverseBits16(index);
+
+	switch(dimension)
+	{
+	case 1:
+		index ^= static_cast<std::uint16_t>((index & 0xff00) >> 8);
+		index ^= static_cast<std::uint16_t>((index & 0xf0f0) >> 4);
+		index ^= static_cast<std::uint16_t>((index & 0xcccc) >> 2);
+		index ^= static_cast<std::uint16_t>((index & 0xaaaa) >> 1);
+		break;
+	case 2:
+		index ^= static_cast<std::uint16_t>((index & 0xc000) >> 14);
+		index ^= static_cast<std::uint16_t>((index & 0x2000) >> 13);
+		index ^= static_cast<std::uint16_t>((index & 0xf000) >> 12);
+		index ^= static_cast<std::uint16_t>((index & 0xcc00) >> 10);
+		index ^= static_cast<std::uint16_t>((index & 0xaa00) >> 9);
+		index ^= static_cast<std::uint16_t>((index & 0x0c00) >> 8);
+		index ^= static_cast<std::uint16_t>((index & 0xc0c0) >> 6);
+		index ^= static_cast<std::uint16_t>((index & 0xa0a0) >> 5);
+		index ^= static_cast<std::uint16_t>((index & 0xcf30) >> 4);
+		index ^= static_cast<std::uint16_t>((index & 0x8888) >> 3);
+		index ^= static_cast<std::uint16_t>((index & 0xf0f0) >> 2);
+		index ^= static_cast<std::uint16_t>((index & 0x6666) >> 1);
+		break;
+	case 3:
+		index ^= static_cast<std::uint16_t>((index & 0xf000) >> 12);
+		index ^= static_cast<std::uint16_t>((index & 0x0c00) >> 10);
+		index ^= static_cast<std::uint16_t>((index & 0x0a00) >> 9);
+		index ^= static_cast<std::uint16_t>((index & 0xfe00) >> 8);
+		index ^= static_cast<std::uint16_t>((index & 0x0fc0) >> 6);
+		index ^= static_cast<std::uint16_t>((index & 0x0820) >> 5);
+		index ^= static_cast<std::uint16_t>((index & 0x07e0) >> 4);
+		index ^= static_cast<std::uint16_t>((index & 0x8e38) >> 3);
+		index ^= static_cast<std::uint16_t>((index & 0x8e38) >> 2);
+		index ^= static_cast<std::uint16_t>((index & 0x2492) >> 1);
+		break;
+	default:
+		break;
+	}
+
+	return index;
+
+#endif
+
+#else
+
 	if(dimension == 0)
 	{
 		return reverseBits16(index);
@@ -242,17 +343,6 @@ OQMC_HOST_DEVICE inline std::uint16_t sobolReversedIndex(std::uint16_t index,
 	return vgetq_lane_u16(bits, 0);
 #endif
 
-#if defined(OQMC_ARCH_SCALAR)
-	std::uint16_t sample = 0;
-	for(int i = 0; i < 16; ++i)
-	{
-		if((index & masks[i]) != 0)
-		{
-			sample ^= matrix[i];
-		}
-	}
-
-	return sample;
 #endif
 }
 

include/oqmc/sobol.h

@@ -99,7 +99,9 @@ void SobolImpl::drawRnd(std::uint32_t rnd[Size]) const
 /// includes performance improvements such as limiting the index to 16 bits,
 /// pre-inverting the input and output matrices, and making use of CPU vector
 /// intrinsics. You need to select an `OPENQMC_ARCH_TYPE` to make use of the
-/// performance from vector intrinsics for a given architecture.
+/// performance from vector intrinsics for a given architecture. The scalar path
+/// uses the construction method by Abdalla G. M. Ahmed in 'An Implementation
+/// Algorithm of 2D Sobol Sequence Fast, Elegant, and Compact' (EGSR 2024).
 ///
 /// This sampler has no cache initialisation cost, it generates all samples on
 /// the fly without touching memory. However the cost per draw sample call is

src/tests/owen.cpp

@@ -4,6 +4,7 @@
 #include "hypothesis.h"
 #include <oqmc/owen.h>
 #include <oqmc/pcg.h>
+#include <oqmc/reverse.h>
 
 #include <gtest/gtest.h>
 
@@ -142,4 +143,72 @@ TEST(OwenTest, ShirleyRemapping)
 	}
 }
 
+TEST(OwenTest, SobolReversedIndex)
+{
+	// clang-format off
+	constexpr std::uint16_t masks[16] = {
+		0b0000000000000001,
+		0b0000000000000010,
+		0b0000000000000100,
+		0b0000000000001000,
+		0b0000000000010000,
+		0b0000000000100000,
+		0b0000000001000000,
+		0b0000000010000000,
+		0b0000000100000000,
+		0b0000001000000000,
+		0b0000010000000000,
+		0b0000100000000000,
+		0b0001000000000000,
+		0b0010000000000000,
+		0b0100000000000000,
+		0b1000000000000000,
+	};
+
+	constexpr std::uint16_t directions[4][16] = {
+		{0x8000, 0x4000, 0x2000, 0x1000, 0x0800, 0x0400, 0x0200, 0x0100,
+		 0x0080, 0x0040, 0x0020, 0x0010, 0x0008, 0x0004, 0x0002, 0x0001},
+		{0xffff, 0x5555, 0x3333, 0x1111, 0x0f0f, 0x0505, 0x0303, 0x0101,
+		 0x00ff, 0x0055, 0x0033, 0x0011, 0x000f, 0x0005, 0x0003, 0x0001},
+		{0xaa09, 0x7706, 0x3903, 0x1601, 0x09aa, 0x0677, 0x0339, 0x0116,
+		 0x00a3, 0x0071, 0x003a, 0x0017, 0x0009, 0x0006, 0x0003, 0x0001},
+		{0xa0c3, 0x4041, 0x302d, 0x101e, 0x0b67, 0x079a, 0x02a4, 0x011b,
+		 0x00c9, 0x0045, 0x002e, 0x001f, 0x000a, 0x0004, 0x0003, 0x0001},
+	};
+	// clang-format on
+
+	// Reference code from the classic scalar implementation. Verifies that
+	// Ahmed 2024 closed form output matches (PR #97).
+	const auto reference = [&](std::uint16_t index, int dimension) {
+		if(dimension == 0)
+		{
+			return oqmc::reverseBits16(index);
+		}
+
+		const auto matrix = directions[dimension];
+
+		std::uint16_t sample = 0;
+		for(int i = 0; i < 16; ++i)
+		{
+			if((index & masks[i]) != 0)
+			{
+				sample ^= matrix[i];
+			}
+		}
+
+		return sample;
+	};
+
+	for(int dimension = 0; dimension < 4; ++dimension)
+	{
+		for(std::uint32_t index = 0; index < (1u << 16); ++index)
+		{
+			const auto value = static_cast<std::uint16_t>(index);
+
+			EXPECT_EQ(oqmc::sobolReversedIndex(value, dimension),
+			          reference(value, dimension));
+		}
+	}
+}
+
 } // namespace

src/tools/cli/matrices.cpp

@@ -32,10 +32,13 @@
 
 #include <oqmc/reverse.h>
 
+#include <array>
 #include <bitset>
 #include <cassert>
 #include <cstdint>
 #include <cstdio>
+#include <utility>
+#include <vector>
 
 constexpr auto numDimensions = 1024;
 constexpr auto size = 52;
@@ -13379,13 +13382,107 @@ void printMatrices(int dimensionSize, int indexSize)
 	}
 }
 
+// In the bit reversed basis, the generator matrix is unit lower-triangular, so
+// GF(2) elimination factors it into shift-mask-xor steps (Ahmed 2024, eq. 18),
+// one per sub-diagonal. Reversed, the steps form the program for that
+// dimension. Programs print twice: left shifts for the device, and the
+// conjugate form with reversed masks and right shifts for the host.
+
+constexpr int wordBits = 16;
+
+std::uint16_t reversedRow(int dimension, int row)
+{
+	assert(dimension >= 0 && dimension < numDimensions);
+	assert(row >= 0 && row < wordBits);
+
+	std::uint16_t bits = 0;
+	for(int column = 0; column < wordBits; ++column)
+	{
+		const auto reversedColumn = wordBits - 1 - column;
+		const auto reversedRowIndex = wordBits - 1 - row;
+
+		const auto value =
+		    oqmc::reverseBits32(matrices[dimension][reversedColumn]);
+		const auto reversed = std::bitset<wordBits>(value);
+		if(reversed[reversedRowIndex])
+		{
+			bits |= static_cast<std::uint16_t>(1u << column);
+		}
+	}
+	return bits;
+}
+
+void printPrograms(int dimensionSize, bool conjugate)
+{
+	assert(dimensionSize >= 0);
+
+	for(int dimension = 0; dimension < dimensionSize; ++dimension)
+	{
+		std::array<std::uint16_t, wordBits> rows;
+		for(int row = 0; row < wordBits; ++row)
+		{
+			rows[row] = reversedRow(dimension, row);
+		}
+
+		std::vector<std::pair<std::uint16_t, int>> program;
+		for(int shift = 1; shift < wordBits; ++shift)
+		{
+			std::uint16_t mask = 0;
+			for(int row = 0; row + shift < wordBits; ++row)
+			{
+				if(((rows[row + shift] >> row) & 1) != 0)
+				{
+					mask |= static_cast<std::uint16_t>(1u << row);
+				}
+			}
+			if(mask == 0)
+			{
+				continue;
+			}
+			for(int row = 0; row + shift < wordBits; ++row)
+			{
+				if(((mask >> row) & 1) != 0)
+				{
+					rows[row + shift] ^= rows[row];
+				}
+			}
+			program.emplace_back(mask, shift);
+		}
+
+		for(int row = 0; row < wordBits; ++row)
+		{
+			assert(rows[row] == static_cast<std::uint16_t>(1u << row));
+		}
+
+		std::printf("// dimension %d\n", dimension);
+		for(auto step = program.rbegin(); step != program.rend(); ++step)
+		{
+			if(conjugate)
+			{
+				std::printf("index ^= static_cast<std::uint16_t>((index & "
+				            "0x%04x) >> %d);\n",
+				            oqmc::reverseBits16(step->first), step->second);
+			}
+			else
+			{
+				std::printf("index ^= static_cast<std::uint16_t>((index & "
+				            "0x%04x) << %d);\n",
+				            step->first, step->second);
+			}
+		}
+		std::printf("\n");
+	}
+}
+
 int main()
 {
 	constexpr auto samplePrecision = 16;
 	constexpr auto dimensionSize = 4;
 	constexpr auto indexSize = 16;
 
 	printMatrices<samplePrecision>(dimensionSize, indexSize);
+	printPrograms(dimensionSize, false);
+	printPrograms(dimensionSize, true);
 
 	return 0;
 }