Add alias sampling by wsttiger · Pull Request #48 · NVIDIA/cudaq-algorithms

wsttiger · 2026-09-08T17:15:23Z

Summary

Adds AliasSamplingPrepare, an implementation of the PREPARE-with-garbage construction described by Babbush et al. in arXiv:1805.03662, Sec. III.D.

Given nonnegative weights w_k, the primitive prepares a state whose index-register measurement probabilities exactly match an integer-discretized alias table. This construction is a standard state-preparation component for qubitized quantum simulation.

The implementation builds on the arithmetic primitives introduced earlier in this PR stack:

QROM loads the (alias, keep) pair associated with each bin.
A CDKM comparator compares the keep threshold with a uniform reference register.
Controlled swaps replace rejected bins with their aliases.
A hand-written adjoint reverses the complete preparation.

This PR is stacked on the reversible-arithmetic PR.

Exact, bounded discretization

The input weights are discretized to mu bits using an integer version of Vose's alias-table algorithm.

The resulting quantum circuit realizes the discretized distribution exactly. Tests compare against exact rational probabilities rather than relying only on floating-point tolerances.

For K bins, the error for each probability is bounded by

1.5 / (K * 2^mu).

The module docstring derives this bound, including why the residual-redistribution cursor cannot wrap around. Zero-weight entries and bins introduced through power-of-two padding receive exactly zero probability.

Resource costs and register layout

The implementation exposes and tests the following resource contract:

Toffoli count:

qrom.toffoli_count + 4 * (mu + 1)
Fredkin count:

num_index

The resource tests validate these formulas against the compiled kernel using cudaq.estimate_resources. They also cross-check the QROM object's reported costs against the compiled QROM circuit.

The garbage-register layout is documented field by field. Its total width is

num_garbage = num_index + 2 * mu + 4 + qrom.num_ladder.

Because qrom.num_ladder depends on the selected QROM implementation, the garbage width automatically reflects the chosen variant.

The internal QROM accepts variant= and block_size= options. The default "auto" mode evaluates the available constructions and selects the lowest-cost option. Tests also force the select_swap variant to verify that both PREPARE and its adjoint are independent of the selected QROM implementation, including clean-ladder reuse during uncomputation.

PREPARE-with-garbage contract

AliasSamplingPrepare guarantees the probability associated with each index value. It does not guarantee clean state preparation of the form

sum_k sqrt(p_k) |k>|0>.

Instead, the prepared state generally has the form

sum_k sqrt(p_k) |k>|G_k>,

where the conditional garbage states |G_k> need not be mutually orthogonal. Consequently, the reduced state of the index register is not necessarily diag(p_k).

This is sufficient for the qubitization identity

<0| PREP^dag SELECT PREP |0> = sum_k p_k U_k.

Because SELECT is diagonal in the index basis, index orthogonality eliminates cross terms regardless of overlaps between the garbage states. The class documentation states this contract explicitly so consumers do not assume clean amplitude preparation or a diagonal reduced index state.

Validation

The test suite includes:

Exact index-marginal comparisons against independently generated alias tables.
Brute-force enumeration of every (bin, reference) branch and comparison with circuit results.
PREPARE-followed-by-adjoint identity checks.
Resource-identity checks for both the default and forced QROM variants.
Exact checks that zero-weight and padded bins remain zero.
An independent adversarial sweep over approximately 4,000 weight tables, including:
- Dynamic ranges from exp(-18) through exp(4).
- Irrational weights.
- Tables containing many zeros.
- Precision values through mu = 8.

The adversarial sweep found no discretization-bound violations, and padded bins remained exactly zero.

At this branch tip:

26 tests are added by this PR.
The complete suite reports 560 passed and 3 skipped.

Deferred work

The following extensions are deliberately left for future changes:

Sub-bit-precision variants.
Measurement-assisted QROM uncomputation.

Both can be introduced behind the existing QROM interface when a consumer requires the associated resource tradeoff.

CDKM/Cuccaro ripple-carry family (register and constant add/subtract, >=-constant comparator) and the ancilla-free Draper QFT family (qft/iqft, Fourier-basis constant add/subtract, the comparator pair with its hand-written adjoint), all little-endian in-place kernels with hand-written inverses. Tests are exhaustive over all inputs for widths up to 5 via the superposition harness, pin every inverse by op-then-inverse identity, and hold the documented gate prices against the compiler with cudaq.estimate_resources: exactly 2n Toffolis for every CDKM operation (0 at K = 0), and zero Toffolis for the QFT family with exact controlled-r1 / r1 / h budgets (n(n-1) + n for the adder, n(n+1) + n + 1 per comparator side). Review fixes: document the 0 <= K <= 2^n precondition on cmp_ge_constant_qft (the borrow wraps mod 2^(n+1) for larger K) and the exact constant_bits/complement_bits length preconditions on add_constant and cmp_ge_constant; note in the module docstring that operands of any one op must be pairwise disjoint; version-scope the cudaq.adjoint limitation (broken as of CUDA-Q 0.15, still unresolved on the 0.16 pre-release) instead of stating it timelessly. Signed-off-by: Scott Thornton <wsttiger@gmail.com>

…amily cmp_ge_register(a, b, carry, out) XOR-loads (a >= b) (unsigned, equal-width little-endian registers) into the out qubit and restores a, b and the carry ancilla: complement b in place (2^n - b = ~b + 1, carry-in set to 1), MAJ sweep, copy the top ripple carry into out, reverse the sweep and undo the complement. Because b doubles as the constant-load register, no n-qubit work register is needed — only the one-qubit carry — and the price is the family contract of exactly 2n Toffolis, pinned against cudaq.estimate_resources with runtime-argument harnesses at widths 1-5 and 8. cmp_gt_register is the free strict variant (a > b <=> not (b >= a): one X plus the >= comparator with the roles swapped), same price. Both are self-inverse (compute-copy- uncompute operand action, XOR-accumulated flag), pinned by apply-twice-is-identity tests on random entangled states, alongside exhaustive truth tables (widths 1-4, flag initially 0 and 1) and superposed-register statevector checks against dense NumPy references. This is the comparator the alias-sampling conditional (keep-value vs alt-value test) and the first-quantized momentum comparisons need (colleague review feedback; GAP_first_quantized_lcu.md). Signed-off-by: Scott Thornton <wsttiger@gmail.com>

AliasSamplingPrepare: the Babbush-style (arXiv:1805.03662, Sec. III.D) PREPARE-with-garbage whose index-register marginal realizes the integer-Vose table distribution exactly (per-bin discretization bound 1.5 / (K 2^mu), derived not assumed), with a hand-written adjoint. The (alias, keep) lookup is the variants QROM under variant="auto": the ladder slice inside the garbage register is qrom.num_ladder wide, so num_garbage = m + 2mu + 4 + qrom.num_ladder (the historical 2m + 2mu + 4 whenever the plain select walk wins the pricing). No variant is pinned - every variant restores its ladder and self-inverts on the clean-ladder sector, which is what the adjoint's second lookup uses. The new qrom property exposes the priced lookup. Tests pin the index marginal against an independent brute-force branch enumeration, the table distribution against the ideal within the derived bound, prepare-then-adjoint identity, degenerate weights and validation messages, plus compiler-pinned resource contracts: each PREPARE side costs exactly qrom.toffoli_count + 4 (mu + 1) Toffolis (the lookup at the QROM's own reported price plus the two CDKM comparator adders) and num_index Fredkins, adjoint == forward gate-for-gate. Review fixes: pass variant/block_size through to the QROM constructor as keyword-only parameters (they shape num_garbage via qrom.num_ladder; QROM validation errors propagate as-is) and pin the variant-agnostic claim with a forced select_swap test asserting the exact marginal, the prepare-then-adjoint identity and the Toffoli cost identity; reject bool mu explicitly while keeping integral numpy ints; document that the per-bin garbage vectors are not mutually orthogonal (only the index amplitude magnitudes are guaranteed - exactly the qubitization <0|PREP' SELECT PREP|0> contract); say "the available QROM variants" instead of naming constructions this base does not ship; version-scope the cudaq.control nested-call rejection (through CUDA-Q 0.15, lifted in 0.16); mark the h-count assertion as incidental and the toffoli_count cross-check as a deliberate bookkeeping-vs-compiler pin. Signed-off-by: Scott Thornton <wsttiger@gmail.com>

Bare |0...0> kets and the |table - ideal| absolute value in the AliasSamplingPrepare docstrings parse as unterminated RST substitution references and fail the warnings-as-errors docs build now that the class is on the API reference page. Same fix as the unary-iteration module received on the QROM PR. Signed-off-by: Scott Thornton <wsttiger@gmail.com>

Implements the reviewer-suggested consolidation: the PREPARE's keep-vs-reference conditional now calls the arithmetic family's cmp_ge_register (flag <- (ref >= keep) on the bare mu-bit operands) instead of the bespoke two-adder weave (subtract keep on a (mu+1)-bit extension, copy the borrow, negate, add back). - Toffoli count per PREPARE side drops from 4(mu+1) to 2mu: the cost identity is now qrom.toffoli_count + 2 mu Toffolis + num_index Fredkins, re-derived in the docstrings and re-pinned by the resource tests against cudaq.estimate_resources. - The (mu+1)-bit extension's two pad qubits (keep_pad, ref_pad) are gone: cmp_ge_register complements its second operand in place and needs only the one-qubit carry the layout already had. num_garbage shrinks by 2 to num_index + 2 mu + 2 + qrom.num_ladder, and ladder_offset moves down accordingly. - The hand-written adjoint stays the literal gate-reversal: cmp_ge_register is palindromic around its self-inverse flag copy, so re-applying it is the reversed gate sequence. All marginal/bound/identity tests pass unchanged; only the layout and cost accounting assertions are re-derived. Signed-off-by: Scott Thornton <wsttiger@gmail.com>

smcardleQ · 2026-09-16T00:24:50Z

+        QROM's own ``ValueError``.
+
+    Garbage caveat: the per-bin garbage vectors ``|g_k>`` are NOT
+    mutually orthogonal (measured off-diagonal overlaps reach ~0.19), so


Remove "(measured off-diagonal overlaps reach ~0.19)" from docstring, as it is instance specific without details.

smcardleQ · 2026-09-16T00:36:26Z

+        # shapes the garbage layout below. QROM validation errors
+        # propagate as-is.
+        self._qrom = QROM(
+            [alias[k] | (keep[k] << num_index) for k in range(num_bins)],


Add a comment on the packing convention here, as the one-liner made it non-obvious for me.

smcardleQ · 2026-09-16T00:38:01Z

+        k0 = m  # keep (keep_pad at k0 + mu)
+        r0 = m + mu + 1  # ref (ref_pad at r0 + mu)
+        flag = m + 2 * mu + 2
+        l0 = m + 2 * mu + 3
+        c0 = l0 + num_ladder  # carry


k0, r0, l0, c0 are not sufficiently informative as variable names.

smcardleQ · 2026-09-16T00:41:51Z

+            # pass, untouched since), so re-applying the lookup XORs the
+            # table back out of garbage[0 : m + mu] — the documented
+            # self-inverse contract of every QROM variant.
+            qrom_kernel(index, garbage[l0:l0 + num_ladder], garbage[0:m + mu])


Does the inverse QROM assume the same ladder ancilla qubits as the forward QROM? The ladder ancillas will be left clean (in |0>) after the forward QROM - hence they might be used by some other subroutine in the algorithm. We may need to allocate new ladder ancillas to help the inverse QROM.
This isn't a deal-breaker; one could instead reserve the ladder ancillas, and allocate other ancillas in the meantime. But potentially something for us to be aware of.

Yes — the adjoint reuses the same ladder ancillas (garbage[ladder_off : ladder_off + num_ladder]), which the forward pass leaves clean in |0⟩, so the self-inverse lookup is valid. You're right that this reserves those qubits across the PREPARE/PREPARE† sandwich rather than freeing them for other subroutines in between; since coherent alias sampling is only ever used as that symmetric sandwich, holding them is intentional. If a consumer needs them freed mid-sandwich we'd allocate separate ancillas for the inverse — noted. (Also left a TODO at this site for the related qrom_inverse/measurement-based-uncompute point in your other comment.)

smcardleQ · 2026-09-16T00:43:31Z

+            # pass, untouched since), so re-applying the lookup XORs the
+            # table back out of garbage[0 : m + mu] — the documented
+            # self-inverse contract of every QROM variant.
+            qrom_kernel(index, garbage[l0:l0 + num_ladder], garbage[0:m + mu])


Current implementation is fine. However, in future we may want to explicitly make this a "qrom_inverse" call. The reason being is that once we allow measurement-based uncomputation, uncomputing the QROM can be significantly cheaper than forward pass computing it, making the inverse alias sampling much cheaper than the forward. Leave a "To Do" ?

smcardleQ

Minor comments on coherent alias sampling, otherwise looks good to merge. I have not checked the arithmetic primitives yet.

…M TODO Per smcardleQ's review of PR NVIDIA#48: - drop the instance-specific '~0.19' overlap figure from the garbage caveat - spell out the QROM-word packing convention (alias low bits, keep above) - rename terse offsets k0/r0/l0/c0 -> keep_off/ref_off/ladder_off/carry_off - add a TODO to expose the inverse QROM as an explicit qrom_inverse call, since measurement-based uncomputation makes the adjoint PREPARE cheaper No behavior change; alias-sampling suite 26/26. Signed-off-by: Scott Thornton <wsttiger@gmail.com>

wsttiger requested a review from smcardleQ September 11, 2026 16:12

wsttiger force-pushed the features/primitives_alias_sampling branch from 77b8e68 to 7d82a46 Compare September 11, 2026 23:10

wsttiger added 5 commits September 15, 2026 16:21

wsttiger force-pushed the features/primitives_alias_sampling branch from 7d82a46 to 334e160 Compare September 15, 2026 16:22

wsttiger mentioned this pull request Sep 15, 2026

Rebuild select_swap as composed kernels; promise the composition contract (cudaq >= 0.16) #52

Open

smcardleQ reviewed Sep 16, 2026

View reviewed changes

smcardleQ approved these changes Sep 16, 2026

View reviewed changes

wsttiger added the enhancement New feature or request label Sep 17, 2026

kvmto mentioned this pull request Sep 19, 2026

Add reversible integer arithmetic primitives #47

Open

Navigation Menu

Add alias sampling - #48

Add alias sampling#48
wsttiger wants to merge 6 commits into
NVIDIA:mainfrom
wsttiger:features/primitives_alias_sampling

wsttiger commented Sep 8, 2026

Uh oh!

smcardleQ Sep 16, 2026

Uh oh!

smcardleQ Sep 16, 2026

Uh oh!

smcardleQ Sep 16, 2026

Uh oh!

smcardleQ Sep 16, 2026

Uh oh!

wsttiger Sep 17, 2026

Uh oh!

smcardleQ Sep 16, 2026

Uh oh!

smcardleQ left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

wsttiger commented Sep 8, 2026

Summary

Exact, bounded discretization

Resource costs and register layout

PREPARE-with-garbage contract

Validation

Deferred work

Uh oh!

smcardleQ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

smcardleQ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

smcardleQ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

smcardleQ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

wsttiger Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

smcardleQ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

smcardleQ left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants