Eytzinger CDF, alias table is packed, 2 versions, CDF sampler now has 3 enumed types (tracking, YOLO, Eytzinger)

Author: karimsayedre

examples_tests

@@ -7,6 +7,7 @@
 
 #include <nbl/builtin/hlsl/cpp_compat.hlsl>
 #include <nbl/builtin/hlsl/bit.hlsl>
+#include <nbl/builtin/hlsl/cpp_compat/intrinsics.hlsl>
 #include <nbl/builtin/hlsl/concepts/core.hlsl>
 #include <nbl/builtin/hlsl/concepts/accessors/generic_shared_data.hlsl>
 
@@ -17,84 +18,187 @@ namespace hlsl
 namespace sampling
 {
 
-// Alias Method (Vose/Walker) discrete sampler.
-//
-// Samples a discrete index in [0, N) with probability proportional to
-// precomputed weights in O(1) time per sample, using a prebuilt alias table.
-//
-// Accessor template parameters must satisfy GenericReadAccessor:
-//   accessor.template get<V, I>(index, outVal)  // void, writes to outVal
-//
-// - ProbabilityAccessor: reads scalar_type threshold in [0, 1] for bin i
-// - AliasIndexAccessor:  reads uint32_t redirect index for bin i
-// - PdfAccessor:         reads scalar_type weight[i] / totalWeight
-//
-// Satisfies TractableSampler (not BackwardTractableSampler: the mapping is discrete).
-// The cache stores the PDF value looked up during generate, avoiding redundant
-// storage of the codomain (sampled index) which is already the return value.
-template<typename T, typename Domain, typename Codomain, typename ProbabilityAccessor, typename AliasIndexAccessor, typename PdfAccessor
+// Packed alias-entry bit layout shared by every packed variant. One 32-bit
+// word holds the redirect index in the low Log2N bits and the stay-
+// probability quantized as an unorm in the high (32 - Log2N) bits.
+//   u * N  = scaled;  bin = floor(scaled);  remainder = scaled - bin
+//   if (remainder < getStayProb(word))  -> result = bin
+//   else                                -> result = getTarget(word)
+// Quantizing the threshold to (32 - Log2N) bits is precision-neutral: `u`
+// already consumed Log2N bits of randomness producing `bin`, so `remainder`
+// carries exactly that many bits of discriminatory power.
+namespace impl
+{
+template<uint32_t Log2N>
+struct AliasBitDecoder
+{
+	static uint32_t getTarget(uint32_t word)
+	{
+		return word & ((1u << Log2N) - 1u);
+	}
+	template<typename T>
+	static T getStayProb(uint32_t word)
+	{
+		const uint32_t unormMax = (~0u) >> Log2N;
+		return T(word >> Log2N) / T(unormMax);
+	}
+};
+} // namespace impl
+
+// 8 B entry used by the NBig == true variant. Embeds the bin's own pdf
+// alongside the packed word so the common stay-case needs no extra tap.
+template<typename T>
+struct PackedAliasEntryB
+{
+	uint32_t packedWord;	// low Log2N: redirect target; high 32-Log2N: stayProb unorm
+	T ownPdf;				// pdf of this bin
+};
+
+
+// NBig == false: 4 B packed word per bin + separate pdf[] array. Per sample
+// = one 4 B word load + one unconditional 4 B pdf[] tap indexed by the
+// selected bin (either the current bin or its redirect). Total 8 B whether
+// the sample stays or aliases. Favours small N.
+template<typename T, typename Domain, typename Codomain, typename PackedWordAccessor, typename PdfAccessor, uint32_t Log2N
 	NBL_PRIMARY_REQUIRES(
 		concepts::UnsignedIntegralScalar<Codomain> &&
-		concepts::accessors::GenericReadAccessor<ProbabilityAccessor, T, Codomain> &&
-		concepts::accessors::GenericReadAccessor<AliasIndexAccessor, Codomain, Codomain> &&
+		concepts::accessors::GenericReadAccessor<PackedWordAccessor, uint32_t, Codomain> &&
 		concepts::accessors::GenericReadAccessor<PdfAccessor, T, Codomain>)
-struct AliasTable
+struct PackedAliasTableA
 {
 	using scalar_type = T;
-
 	using domain_type = Domain;
 	using codomain_type = Codomain;
 	using density_type = scalar_type;
 	using weight_type = density_type;
+	using decoder = impl::AliasBitDecoder<Log2N>;
+	NBL_CONSTEXPR_STATIC_INLINE bool NBig = false;
 
 	struct cache_type
 	{
 		density_type pdf;
 	};
 
-	static AliasTable create(NBL_CONST_REF_ARG(ProbabilityAccessor) _probAccessor, NBL_CONST_REF_ARG(AliasIndexAccessor) _aliasAccessor, NBL_CONST_REF_ARG(PdfAccessor) _pdfAccessor, codomain_type _size)
+	static PackedAliasTableA create(NBL_CONST_REF_ARG(PackedWordAccessor) _entryAcc, NBL_CONST_REF_ARG(PdfAccessor) _pdfAcc, codomain_type _size)
 	{
-		AliasTable retval;
-		retval.probAccessor = _probAccessor;
-		retval.aliasAccessor = _aliasAccessor;
-		retval.pdfAccessor = _pdfAccessor;
-		// Precompute tableSize as float minus 1 ULP so that u=1.0 maps to bin N-1
+		PackedAliasTableA retval;
+		retval.entryAcc = _entryAcc;
+		retval.pdfAcc = _pdfAcc;
 		const scalar_type exact = scalar_type(_size);
 		retval.tableSizeMinusUlp = nbl::hlsl::bit_cast<scalar_type>(nbl::hlsl::bit_cast<uint32_t>(exact) - 1u);
 		return retval;
 	}
 
-	// BasicSampler interface
 	codomain_type generate(const domain_type u) NBL_CONST_MEMBER_FUNC
 	{
 		const scalar_type scaled = u * tableSizeMinusUlp;
 		const codomain_type bin = _static_cast<codomain_type>(scaled);
 		const scalar_type remainder = scaled - scalar_type(bin);
 
-		scalar_type prob;
-		probAccessor.template get<scalar_type, codomain_type>(bin, prob);
-
-		// Use if-statement to avoid select: aliasIndex is a dependent read
-		codomain_type result;
-		if (remainder < prob)
-		{
-			result = bin;
-		}
-		else
-		{
-			codomain_type alias;
-			aliasAccessor.template get<codomain_type, codomain_type>(bin, alias);
-			result = alias;
-		}
+		uint32_t packedWord;
+		entryAcc.template get<uint32_t, codomain_type>(bin, packedWord);
+		return hlsl::select(remainder < decoder::template getStayProb<scalar_type>(packedWord), bin, codomain_type(decoder::getTarget(packedWord)));
+	}
 
+	codomain_type generate(const domain_type u, NBL_REF_ARG(cache_type) cache) NBL_CONST_MEMBER_FUNC
+	{
+		const codomain_type result = generate(u);
+		pdfAcc.template get<scalar_type, codomain_type>(result, cache.pdf);
 		return result;
 	}
 
-	// TractableSampler interface
+	density_type forwardPdf(const domain_type u, NBL_CONST_REF_ARG(cache_type) cache) NBL_CONST_MEMBER_FUNC
+	{
+		return cache.pdf;
+	}
+
+	weight_type forwardWeight(const domain_type u, NBL_CONST_REF_ARG(cache_type) cache) NBL_CONST_MEMBER_FUNC
+	{
+		return cache.pdf;
+	}
+
+	density_type backwardPdf(const codomain_type v) NBL_CONST_MEMBER_FUNC
+	{
+		scalar_type pdf;
+		pdfAcc.template get<scalar_type, codomain_type>(v, pdf);
+		return pdf;
+	}
+
+	weight_type backwardWeight(const codomain_type v) NBL_CONST_MEMBER_FUNC
+	{
+		return backwardPdf(v);
+	}
+
+	PackedWordAccessor entryAcc;
+	PdfAccessor pdfAcc;
+	scalar_type tableSizeMinusUlp;
+};
+
+// NBig == true: 8 B entry {packedWord, ownPdf} + separate pdf[] array. Per
+// sample = one 8 B entry load (covers the common stay case where cache
+// already has ownPdf). If the sample aliases, a conditional 4 B pdf[target]
+// tap fills the cache. Total 8 B stay, 12 B aliased. Favours large N.
+template<typename T, typename Domain, typename Codomain, typename EntryAccessor, typename PdfAccessor, uint32_t Log2N
+	NBL_PRIMARY_REQUIRES(
+		concepts::UnsignedIntegralScalar<Codomain> &&
+		concepts::accessors::GenericReadAccessor<EntryAccessor, PackedAliasEntryB<T>, Codomain> &&
+		concepts::accessors::GenericReadAccessor<PdfAccessor, T, Codomain>)
+struct PackedAliasTableB
+{
+	using scalar_type = T;
+	using domain_type = Domain;
+	using codomain_type = Codomain;
+	using density_type = scalar_type;
+	using weight_type = density_type;
+	using entry_type = PackedAliasEntryB<scalar_type>;
+	using decoder = impl::AliasBitDecoder<Log2N>;
+	NBL_CONSTEXPR_STATIC_INLINE bool NBig = true;
+
+	struct cache_type
+	{
+		density_type pdf;
+	};
+
+	static PackedAliasTableB create(NBL_CONST_REF_ARG(EntryAccessor) _entryAcc, NBL_CONST_REF_ARG(PdfAccessor) _pdfAcc, codomain_type _size)
+	{
+		PackedAliasTableB retval;
+		retval.entryAcc = _entryAcc;
+		retval.pdfAcc = _pdfAcc;
+		const scalar_type exact = scalar_type(_size);
+		retval.tableSizeMinusUlp = nbl::hlsl::bit_cast<scalar_type>(nbl::hlsl::bit_cast<uint32_t>(exact) - 1u);
+		return retval;
+	}
+
+	codomain_type generate(const domain_type u) NBL_CONST_MEMBER_FUNC
+	{
+		const scalar_type scaled = u * tableSizeMinusUlp;
+		const codomain_type bin = _static_cast<codomain_type>(scaled);
+		const scalar_type remainder = scaled - scalar_type(bin);
+
+		entry_type entry;
+		entryAcc.template get<entry_type, codomain_type>(bin, entry);
+		return hlsl::select(remainder < decoder::template getStayProb<scalar_type>(entry.packedWord), bin, codomain_type(decoder::getTarget(entry.packedWord)));
+	}
+
 	codomain_type generate(const domain_type u, NBL_REF_ARG(cache_type) cache) NBL_CONST_MEMBER_FUNC
 	{
-		const codomain_type result = generate(u);
-		pdfAccessor.template get<scalar_type, codomain_type>(result, cache.pdf);
+		const scalar_type scaled = u * tableSizeMinusUlp;
+		const codomain_type bin = _static_cast<codomain_type>(scaled);
+		const scalar_type remainder = scaled - scalar_type(bin);
+
+		entry_type entry;
+		entryAcc.template get<entry_type, codomain_type>(bin, entry);
+
+		const bool stay = remainder < decoder::template getStayProb<scalar_type>(entry.packedWord);
+		
+		cache.pdf = entry.ownPdf;
+		codomain_type result = bin;
+		if (!stay)
+		{
+			const codomain_type target = codomain_type(decoder::getTarget(entry.packedWord));
+			pdfAcc.template get<scalar_type, codomain_type>(target, cache.pdf);
+			result = target;
+		}
 		return result;
 	}
 
@@ -111,7 +215,7 @@ struct AliasTable
 	density_type backwardPdf(const codomain_type v) NBL_CONST_MEMBER_FUNC
 	{
 		scalar_type pdf;
-		pdfAccessor.template get<scalar_type, codomain_type>(v, pdf);
+		pdfAcc.template get<scalar_type, codomain_type>(v, pdf);
 		return pdf;
 	}
 
@@ -120,9 +224,8 @@ struct AliasTable
 		return backwardPdf(v);
 	}
 
-	ProbabilityAccessor probAccessor;
-	AliasIndexAccessor aliasAccessor;
-	PdfAccessor pdfAccessor;
+	EntryAccessor entryAcc;
+	PdfAccessor pdfAcc;
 	scalar_type tableSizeMinusUlp;
 };