From df4599cd869d08a08ad2cb656034790e7088e20c Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Thu, 16 Jul 2026 22:32:24 -0700 Subject: [PATCH 1/7] First draft at unifiying the shaders --- .../src/gpu_full_state_simulator/unified.wgsl | 3587 +++++++++++++++++ 1 file changed, 3587 insertions(+) create mode 100644 source/simulators/src/gpu_full_state_simulator/unified.wgsl diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/unified.wgsl new file mode 100644 index 00000000000..9af9afcdb8c --- /dev/null +++ b/source/simulators/src/gpu_full_state_simulator/unified.wgsl @@ -0,0 +1,3587 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// See https://webgpufundamentals.org/webgpu/lessons/webgpu-wgsl.html for an overview +// See https://www.w3.org/TR/WGSL/ for the details +// See https://webgpu.github.io/webgpu-samples/ for examples + +//#region Compile time replaced constants + +// WGSL has pipeline overridables, but they're a pain and limited, so just string replace constants here +const QUBIT_COUNT: i32 = 8; // REPLACE +const RESULT_COUNT: u32 = 8; // REPLACE +const WORKGROUPS_PER_SHOT: i32 = 1; // REPLACE +const ENTRIES_PER_THREAD: i32 = 5; // REPLACE +const THREADS_PER_WORKGROUP: i32 = 32; // REPLACE +const MAX_QUBIT_COUNT: i32 = 27; // REPLACE +const MAX_QUBITS_PER_WORKGROUP: i32 = 5; // REPLACE +const NOISE_TABLE_COUNT: u32 = 1; // REPLACE +const NOISE_ENTRY_COUNT: u32 = 1; // REPLACE +const MAX_REGISTERS: u32 = 256; // REPLACE +const MAX_MEMORY: u32 = 256; // REPLACE +const INSTRUCTIONS_SIZE: u32 = 0; // REPLACE +const BLOCK_TABLE_SIZE: u32 = 0; // REPLACE +const FUNCTION_TABLE_SIZE: u32 = 0; // REPLACE +const PHI_TABLE_SIZE: u32 = 0; // REPLACE +const SWITCH_CASES_SIZE: u32 = 0; // REPLACE +const CALL_ARGS_SIZE: u32 = 0; // REPLACE +const CONSTANT_DATA_SIZE: u32 = 0; // REPLACE + +//#endregion + +//#region Error codes + +const ERR_INVALID_PROBS = 1u; +const ERR_INVALID_THREAD_TOTAL = 2u; +const ERR_CALL_STACK_OVERFLOW = 3u; +const ERR_CALL_STACK_UNDERFLOW = 4u; +const ERR_INVALID_INSTRUCTION = 5u; +const ERR_ALLOCA_OUT_OF_BOUNDS = 6u; +const ERR_MEMORY_OUT_OF_BOUNDS = 7u; +const ERR_UNSUPPORTED_LOSS_POLICY = 32u; + +//#endregion + +//#region Operation IDs +const OPID_ID = 0u; +const OPID_RESETZ = 1u; +const OPID_X = 2u; +const OPID_Y = 3u; +const OPID_Z = 4u; +const OPID_H = 5u; +const OPID_S = 6u; +const OPID_SAdj = 7u; +const OPID_T = 8u; +const OPID_TAdj = 9u; +const OPID_RX = 12u; +const OPID_RY = 13u; +const OPID_RZ = 14u; +const OPID_CX = 15u; +const OPID_CZ = 16u; +const OPID_RXX = 17u; +const OPID_RYY = 18u; +const OPID_RZZ = 19u; +const OPID_MZ = 21u; +const OPID_MRESETZ = 22u; +const OPID_SWAP = 24u; +const OPID_MAT1Q = 25u; +const OPID_MAT2Q = 26u; +const OPID_CY = 29u; + +const OPID_PAULI_NOISE_1Q = 128u; +const OPID_PAULI_NOISE_2Q = 129u; +const OPID_LOSS_NOISE = 130u; +const OPID_CORRELATED_NOISE = 131u; + +// If the application of noise results in a custom matrix, it will have been stored in the shot buffer +// These OPIDs indicate to use that matrix and for how many qubits. (The qubit ids are in the original Op) +const OPID_SHOT_BUFF_1Q = 256u; +const OPID_SHOT_BUFF_2Q = 257u; + +//#endregion + +//#region Misc constants + +// Tolerance for probabilities to sum to 1.0 +const PROB_THRESHOLD: f32 = 0.0001; + +// Always use 32 threads per workgroup for max concurrency on most current GPU hardware +const MAX_WORKGROUP_SUM_PARTITIONS: i32 = 1i << u32(MAX_QUBIT_COUNT - MAX_QUBITS_PER_WORKGROUP); + +// Loss policy values. These are stamped onto a gate op's `q3` field by the host +// (see `LossPolicy::as_u32` on the Rust side) and tell the shader how to handle +// the gate when one of its operands is lost. `0` means "no policy stamped", +// which the shader treats the same as SKIP. +const LOSS_POLICY_SKIP = 0u; +const LOSS_POLICY_PROPAGATE = 1u; +const LOSS_POLICY_DEGRADE = 2u; +const LOSS_POLICY_RESIDUAL_S_DAGGER = 3u; +const LOSS_POLICY_APPLY_ANYWAY = 4u; + +//#endregion + +//#region Adaptive interpreter constants + +const MAX_CLASSICAL_STEPS: u32 = 4096u; + +// Status codes +const STATUS_RUNNING: u32 = 0u; +const STATUS_QUANTUM_PENDING: u32 = 1u; +const STATUS_TERMINATED: u32 = 2u; +const STATUS_ERROR: u32 = 3u; +const STATUS_YIELD: u32 = 4u; + +// pending_op_type values: 0 = gate, 1 = measure, 2 = reset, 3 = loss commit. +// A loss-commit pending op carries the lost qubit in pending_op_idx (not an +// ops-pool index) and is produced while draining pending_loss_mask. Its value +// must not collide with the gate/measure/reset types resolved in prepare_op. +const PENDING_OP_LOSS_COMMIT: u32 = 3u; + +// ----------------------------------------------------------------------------- +// Adaptive interpreter — opcodes +// ----------------------------------------------------------------------------- + +// Shared opcode constants for the Adaptive Profile QIR bytecode interpreter. +// +// These constants define the bytecode encoding used by the Python AdaptiveProfilePass +// (emitter). Values must stay in sync with the Python ``_adaptive_opcodes.py`` file. +// +// Opcode word layout:: +// +// bits [7:0] = primary opcode +// bits [15:8] = sub-opcode / condition code +// bits [23:16] = flags +// +// Compose via bitwise OR: ``opcode | (sub << 8) | flag`` +// Example: ``OP_ICMP | (ICMP_SLE << 8) | FLAG_SRC1_IMM`` + +// -- Flags (pre-shifted to bit 16+) ------------------------------------------ +const FLAG_SRC0_IMM: u32 = 1 << 16; // src0 field is an immediate value, not a register +const FLAG_SRC1_IMM: u32 = 1 << 17; // src1 field is an immediate value, not a register +const FLAG_DST_IMM: u32 = 1 << 18; // dst field is an immediate value, not a register +const FLAG_AUX0_IMM: u32 = 1 << 19; // aux0 field is an immediate value, not a register +const FLAG_AUX1_IMM: u32 = 1 << 20; // aux1 field is an immediate value, not a register +const FLAG_AUX2_IMM: u32 = 1 << 21; // aux2 field is an immediate value, not a register +const FLAG_AUX3_IMM: u32 = 1 << 22; // aux3 field is an immediate value, not a register + +// -- Control Flow ------------------------------------------------------------- +const OP_NOP: u32 = 0x00; +const OP_RET: u32 = 0x02; +const OP_JUMP: u32 = 0x04; +const OP_BRANCH: u32 = 0x05; +const OP_SWITCH: u32 = 0x06; +const OP_CALL: u32 = 0x07; +const OP_CALL_RETURN: u32 = 0x08; + +// -- Quantum ------------------------------------------------------------------ +const OP_QUANTUM_GATE: u32 = 0x10; +const OP_MEASURE: u32 = 0x11; +const OP_RESET: u32 = 0x12; +const OP_READ_RESULT: u32 = 0x13; +const OP_RECORD_OUTPUT: u32 = 0x14; +const OP_READ_LOSS: u32 = 0x15; + +// -- Integer Arithmetic ------------------------------------------------------- +const OP_ADD: u32 = 0x20; +const OP_SUB: u32 = 0x21; +const OP_MUL: u32 = 0x22; +const OP_UDIV: u32 = 0x23; +const OP_SDIV: u32 = 0x24; +const OP_UREM: u32 = 0x25; +const OP_SREM: u32 = 0x26; + +// -- Bitwise / Shift --------------------------------------------------------- +const OP_AND: u32 = 0x28; +const OP_OR: u32 = 0x29; +const OP_XOR: u32 = 0x2A; +const OP_SHL: u32 = 0x2B; +const OP_LSHR: u32 = 0x2C; +const OP_ASHR: u32 = 0x2D; + +// -- Comparison --------------------------------------------------------------- +const OP_ICMP: u32 = 0x30; +const OP_FCMP: u32 = 0x31; + +// -- Float Arithmetic --------------------------------------------------------- +const OP_FADD: u32 = 0x38; +const OP_FSUB: u32 = 0x39; +const OP_FMUL: u32 = 0x3A; +const OP_FDIV: u32 = 0x3B; +const OP_FREM: u32 = 0x3C; + +// -- Type Conversion ---------------------------------------------------------- +const OP_ZEXT: u32 = 0x40; +const OP_SEXT: u32 = 0x41; +const OP_TRUNC: u32 = 0x42; +const OP_FPEXT: u32 = 0x43; +const OP_FPTRUNC: u32 = 0x44; +const OP_INTTOPTR: u32 = 0x45; +const OP_FPTOSI: u32 = 0x46; +const OP_SITOFP: u32 = 0x47; +const OP_FPTOUI: u32 = 0x48; +const OP_UITOFP: u32 = 0x49; + +// -- SSA / Data Movement ----------------------------------------------------- +const OP_PHI: u32 = 0x50; +const OP_SELECT: u32 = 0x51; +const OP_MOV: u32 = 0x52; +const OP_CONST: u32 = 0x53; + +// -- Memory Operations -------------------------------------------------------- +const OP_ALLOCA: u32 = 0x60; +const OP_LOAD: u32 = 0x61; +const OP_STORE: u32 = 0x62; +const OP_GEP: u32 = 0x63; + +// -- ICmp condition codes (sub-opcode, placed in bits[15:8] via << 8) --------- +// Reference: https://llvm.org/docs/LangRef.html#icmp-instruction +const ICMP_EQ: u32 = 0; +const ICMP_NE: u32 = 1; +const ICMP_SLT: u32 = 2; +const ICMP_SLE: u32 = 3; +const ICMP_SGT: u32 = 4; +const ICMP_SGE: u32 = 5; +const ICMP_ULT: u32 = 6; +const ICMP_ULE: u32 = 7; +const ICMP_UGT: u32 = 8; +const ICMP_UGE: u32 = 9; + +// -- FCmp condition codes ----------------------------------------------------- +// Reference: https://llvm.org/docs/LangRef.html#fcmp-instruction +const FCMP_FALSE: u32 = 0; +const FCMP_OEQ: u32 = 1; +const FCMP_OGT: u32 = 2; +const FCMP_OGE: u32 = 3; +const FCMP_OLT: u32 = 4; +const FCMP_OLE: u32 = 5; +const FCMP_ONE: u32 = 6; +const FCMP_ORD: u32 = 7; +const FCMP_UNO: u32 = 8; +const FCMP_UEQ: u32 = 9; +const FCMP_UGT: u32 = 10; +const FCMP_UGE: u32 = 11; +const FCMP_ULT: u32 = 12; +const FCMP_ULE: u32 = 13; +const FCMP_UNE: u32 = 14; +const FCMP_TRUE: u32 = 15; + +// -- Sentinel values ---------------------------------------------------------- +const VOID_RETURN: u32 = 0xFFFFFFFF; // Function does not have a return value. + +//#endregion + +//#region Data structures + +struct WorkgroupSums { + qubits: array, // Each vec2f holds (zero_probability, one_probability) +}; + +struct WorkgroupCollationBuffer { + sums: array, +}; + +struct QubitState { + zero_probability: f32, + one_probability: f32, + heat: f32, // -1.0 = lost + idle_since: f32, +} + +// Used to track state for the random number generator per shot. See `next_rand_f32` later for details. +struct xorwow_state { + counter: u32, + x: array +} + +/// GPU bytecode instruction. +/// +/// Layout: +/// - `opcode`: packed word — bits\[7:0\]=primary, bits\[15:8\]=sub/condition, bits\[23:16\]=flags +/// - `dst`: destination register or branch target +/// - `src0`, `src1`: source registers or immediates +/// - `aux0`-`aux3`: auxiliary fields (gate index, block ids, side-table offsets, etc.) +struct Instruction { + opcode: u32, + dst: u32, + src0: u32, + src1: u32, + aux0: u32, + aux1: u32, + aux2: u32, + aux3: u32, +} + +struct Block { + instr_offset: u32, + instr_count: u32, +} + +struct Function { + entry_block_id: u32, + param_count: u32, + param_base_reg: u32, + reserved: u32, +} + +struct PhiNodeEntry { + block_id: u32, + val_reg: u32, +} + +struct SwitchCase { + case_val: u32, + target_block: u32, +} + +struct Program { + /// Bytecode instructions. + instructions: array, + /// Block table: indexed by block ID. + block_table: array, + /// Function table. + function_table: array, + /// Phi entries table: `[predecessor_block_id, value_register]` entries. + phi_table: array, + /// Switch cases table: `[match_value, target_block]` entries. + switch_table: array, + /// Call argument register indices. + call_arg_table: array, + /// Constant data pool (flattened array constant values). + constant_data: array, +} + +struct CallStackFrame { + /// Resume on this block on return. + block_id: u32, + /// Instruction after the call. + return_pc: u32, + /// Where to write the return value. + return_reg: u32, + /// This is for alignment. + reserved: u32, +} + +/// Per-shot interpreter state. +struct InterpreterState { + /// Instruction index (absolute), PC stands for Program Counter. + pc: u32, + /// Current block ID. + current_block_id: u32, + ///Previous block ID (for phi resolution). + previous_block_id: u32, + /// 0=running, 1=quantum_pending, 2=terminated, 3=error, 4=yield. + status: u32, + /// Quantum op table index. + pending_op_idx: u32, + /// 0=gate, 1=measure, 2=reset. + pending_op_type: u32, + /// From ret instruction + exit_code: u32, + /// Call stack pointer. + call_sp: u32, + /// Call stack frames (4 u32 per frame × 14 frames = 56). + call_stack_frames: array, + /// Per-shot register file. + registers: array, + /// Per-shot memory (constant_data + alloca'd values). + memory: array, +} + +// Buffer containing the state for each shot to execute per kernel dispatch +// An instance of this is tracked on the GPU for every active shot +struct ShotData { + shot_id: u32, + next_op_idx: u32, + + // The below random numbers will be initialized from the RNG per operation in the 'prepare_op' stage + // Then the 'execute_op' stage will read these precomputed random numbers for noise modeling + rng_state: xorwow_state, // 6 x u32 + rand_pauli: f32, + rand_damping: f32, + rand_dephase: f32, + rand_measure: f32, + // Bitmask of qubits the most recent noise sampler chose to lose. A following + // loss-commit op consumes (and clears) its qubit's bit. + pending_loss_mask: u32, + + // The type of the next operation to execute. This will be OPID_SHOT_BUFF_* if it should use the unitary from the op buffer + op_type: u32, + op_idx: u32, + + duration: f32, // Total duration of the shot so far, used for time-dependent noise modeling and shot estimations + renormalize: f32, // Value to renormalize the state vector by on next execute (1.0 = no renormalization needed) + + // For quick testing during execution to enable skipping blocks of entries + // TODO: Actually use these masks during execution to skip unneeded work + qubit_is_0_mask: u32, // Bitmask for which qubits are currently in |0> state + qubit_is_1_mask: u32, // Bitmask for which qubits are currently in |1> state + + // Track which qubit probabilities were updated in the last operation (to collate on next prepare_op) + qubits_updated_last_op_mask: u32, + // 20 x 4 bytes to this point = 80 bytes + + // Track the per-qubit probabilities for optimization of measurement sampling and noise modeling + qubit_state: array, // 27 x 16 bytes = 432 bytes + // 512 bytes to this point + + // Map this to the Op structure for ease of use + unitary: array, // For MAT1Q and MAT2Q ops. + + // Adaptive interpreter state (embedded to reduce storage buffer count). + // This is initialized by the host after the GPU init kernel runs. + interp: InterpreterState, +} +// See https://www.w3.org/TR/WGSL/#structure-member-layout for alignment rules + +// Buffer containing the list of operations (gates and noise) that make up the program to simulate +struct Op { + id: u32, + q1: u32, + q2: u32, + q3: u32, + policy: u32, + pad0: u32, + pad1: u32, + pad2: u32, + // Entries in the unitary are: 00, 01, 02, 03, 10, 11, 12, 13, 20, ..., 32, 33 + // 1q matrix elements are stored in: 00, 01, 10, 11 (i.e., indices 0, 1, 4, and 5) + unitary: array, +} // Struct size: 4 * 4 + 16 * 8 = 160 bytes (which is aligned to 16 bytes) + +struct ShotParams { + shot_idx: i32, + shot_state_vector_start: i32, + workgroup_collation_idx: i32, + workgroup_idx_in_shot: i32, + thread_idx_in_shot: i32, + total_threads_per_shot: i32, + zero_entry_count: i32, + op_iterations: i32, +} + +struct NoiseTableMetadata { + /// The total probability of any noise (i.e. sum of all noise entries) in `Q1.63` format + noise_probability_lo: u32, + noise_probability_hi: u32, + /// The start offset of this table's entries in the global `NoiseTableEntry` array + start_offset: u32, + /// The number of entries in this noise table + entry_count: u32, +} + +struct NoiseTableEntry { + /// The correlated pauli string as bits (2 bits per qubit). If bit 0 is set, then it has bit-flip + /// noise, and if bit 1 is set then it has phase-flip noise. e.g., `110001 == "YIX"` + paulis_lo: u32, + paulis_hi: u32, + /// The probability of the noise occurring in `Q1_63` format. This is a float format where the high + /// order bit (bit 63) has the value 1.0 (`2^0 / 1`), bit 62 has the value 0.5 (`2^1 / 1`), etc. + /// all the way to bit 63 with a value of approx 1.0842e-19 (`2^63 / 1`). This gives a range of + /// values from [0..2) with equal spacing of 1.0842e-19 between values (unlike float or double), + /// which makes it more suitable for random numbers used to select between a large number of small + /// probability entries. + probability_lo: u32, + probability_hi: u32, +} + +// BatchData holds all the read-only data shared across all shots in a batch. +struct BatchData { + correlated_noise_tables: array, + correlated_noise_entries: array, + program: Program, +} + +// Result of sampling which correlated noise entry (if any) to apply. +struct CorrelatedNoiseSample { + should_apply: u32, // 0 = no noise, 1 = apply noise + paulis_lo: u32, + paulis_hi: u32, +} + +// For every qubit, each 'execute' kernel thread will update its own workgroup storage location for accumulating probabilities +// The final probabilities will be reduced and written back to the shot state after the parallel execution completes. +struct QubitProbabilityPerThread { + zero: array, + one: array, +}; // size: 216 bytes + +// When an error occurs, the below diagnostic data structure is used to store information about the error +struct DiagnosticData { + error_code: atomic, + termination_count: atomic, + extra1: u32, + extra2: f32, + extra3: f32, + _padding: u32, + shot: ShotData, // 640 bytes + op: Op, // 144 bytes + // Below is usually 6,912 bytes (size = THREADS_PER_WORKGROUP (32) * (8 * MAX_QUBIT_COUNT (27)) + workgroup_probabilities: array, + // Below is usually 27,648 bytes (1 << u32(MAX_QUBIT_COUNT - MAX_QUBITS_PER_WORKGROUP)) * (8 * MAX_QUBIT_COUNT) bytes + collation_buffer: WorkgroupCollationBuffer, +}; + +struct Uniforms { + batch_start_shot_id: i32, + rng_seed: u32, +} + +//#endregion + +//#region Buffers and workgroup memory + +@group(0) @binding(0) +var workgroup_collation: WorkgroupCollationBuffer; +// Around 128 max partitions times 27 qubits times 8 bytes = 27 KB max size + +@group(0) @binding(1) +var shots: array; + +@group(0) @binding(2) +var ops: array; + +// The one large buffer of state vector amplitudes. (Partitioned into multiple shots) +@group(0) @binding(3) +var stateVector: array; + +// Buffer for storing measurement results per shot +@group(0) @binding(4) +var results: array>; + +@group(0) @binding(5) +var diagnostics: DiagnosticData; + +@group(0) @binding(6) +var uniforms: Uniforms; + +@group(0) @binding(7) +var batch_data: BatchData; + +var qubitProbabilities: array; +// Workgroup memory size: THREADS_PER_WORKGROUP (32) * 216 = 6,912 bytes. + +//#endregion + +//#region Math utility functions + +// Get the magnitude squared of a complex number +fn cplxMag2(a: vec2f) -> f32 { + return (a.x * a.x + a.y * a.y); +} + +// Complex multiplication +fn cplxMul(a: vec2f, b: vec2f) -> vec2f { + return vec2f( + a.x * b.x - a.y * b.y, + a.x * b.y + a.y * b.x + ); +} + +// Complex negation +fn cplxNeg(a: vec2f) -> vec2f { + return vec2f(-a.x, -a.y); +} + +// Negate all elements in a 4-element row of complex numbers +fn rowNeg(a: array) -> array { + return array( + cplxNeg(a[0]), + cplxNeg(a[1]), + cplxNeg(a[2]), + cplxNeg(a[3])); +} + +// Compute the inner product of two 4-element rows of complex numbers +fn innerProduct(a: array, b: array) -> vec2f { + var result: vec2f = vec2f(0.0, 0.0); + for (var i: u32 = 0u; i < 4u; i++) { + result += cplxMul(a[i], b[i]); + } + return result; +} + +fn getOpRow(op_idx: u32, row: u32) -> array { + let op = &ops[op_idx]; + return array( + op.unitary[row * 4 + 0], + op.unitary[row * 4 + 1], + op.unitary[row * 4 + 2], + op.unitary[row * 4 + 3]); +} + +fn getUnitaryRow(shot_idx: i32, row: u32) -> array { + let shot = &shots[shot_idx]; + return array( + shot.unitary[row * 4 + 0], + shot.unitary[row * 4 + 1], + shot.unitary[row * 4 + 2], + shot.unitary[row * 4 + 3]); +} + +fn setUnitaryRow(shot_idx: u32, row: u32, newRow: array) { + let shot = &shots[shot_idx]; + shot.unitary[row * 4 + 0] = newRow[0]; + shot.unitary[row * 4 + 1] = newRow[1]; + shot.unitary[row * 4 + 2] = newRow[2]; + shot.unitary[row * 4 + 3] = newRow[3]; +} + +//#endregion + +//#region Hash and random number generation + +// See https://www.reedbeta.com/blog/hash-functions-for-gpu-rendering/ +// Use PCG hash function to generate a well-distributed hash from a simple integer input (e.g., shot id) +fn hash_pcg(input: u32) -> u32 { + var state = input * 747796405u + 2891336453u; + var word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u; + return (word >> 22u) ^ word; +} + +// Returns a random u32 value based on the xorwow algorithm +fn next_rand_u32(shot_idx: u32) -> u32 { + // Based on https://en.wikipedia.org/wiki/Xorshift + let rng_state = &shots[shot_idx].rng_state; + + let t: u32 = rng_state.x[4]; + let s: u32 = rng_state.x[0]; + rng_state.x[4] = rng_state.x[3]; + rng_state.x[3] = rng_state.x[2]; + rng_state.x[2] = rng_state.x[1]; + rng_state.x[1] = s; + + // TODO: Simplify with a `var` once https://github.com/wgsl-analyzer/wgsl-analyzer/issues/1317 is fixed + let t2 = t ^ (t >> 2u); + let t3 = t2 ^ (t2 << 1u); + let t4 = t3 ^ s ^ (s << 4u); + rng_state.x[0] = t4; + rng_state.counter = rng_state.counter + 362437u; + return t4 + rng_state.counter; +} + +fn next_rand_f32(shot_idx: u32) -> f32 { + let rand_u32: u32 = next_rand_u32(shot_idx); + + // Convert the 32 random bits to a float in the [0.0, 1.0) range + + // Keep only the lower 23 bits (the fraction portion of a float) with a 0 exponent biased to 127 + let rand_f32_bits = (rand_u32 & 0x7FFFFF) | (127 << 23); + // Bitcast to an f32 in the [1.0, 2.0) range + let f: f32 = bitcast(rand_f32_bits); + // And decrement by 1 to return values from [0..1) + return f - 1.0; +} + +//#endregion + +fn is_1q_phase_gate(op_id: u32) -> bool { + return (op_id == OPID_S || op_id == OPID_SAdj || op_id == OPID_T || op_id == OPID_TAdj || op_id == OPID_RZ); +} + +fn is_1q_op(op_id: u32) -> bool { + return ((op_id >= OPID_ID && op_id <= OPID_RZ) || + op_id == OPID_MZ || op_id == OPID_MRESETZ || + op_id == OPID_MAT1Q || op_id == OPID_SHOT_BUFF_1Q); +} + +fn shot_init_per_op(shot_idx: u32) { + let shot = &shots[shot_idx]; + + // Default to 1.0 renormalization (i.e., no renormalization needed). MResetZ or noise affecting the + // overall probability distribution (e.g. loss or amplitude damping) will update this if needed. + shot.renormalize = 1.0; + shot.qubits_updated_last_op_mask = 0u; + + // Generate the next set of random numbers to use for noise and measurement + shot.rand_pauli = next_rand_f32(shot_idx); + shot.rand_damping = next_rand_f32(shot_idx); + shot.rand_dephase = next_rand_f32(shot_idx); + shot.rand_measure = next_rand_f32(shot_idx); + // Reserved draw: qubit loss is now sampled from the combined `rand_pauli` + // distribution rather than its own value, but we still advance the RNG by + // one draw here to keep the per-op random stream (and thus seeded results) + // identical to the previous loss model. + next_rand_f32(shot_idx); +} + +// Resets the entire shot state, including RNG, probabilities, and per-qubit tracking. +fn reset_all(shot_idx: i32) { + let shot = &shots[shot_idx]; + + // One of the main goals of the shot_id is to seed the RNG state uniquely per shot + let rng_seed = uniforms.rng_seed; + let shot_id = u32(uniforms.batch_start_shot_id + shot_idx); + + // Due to DX12 backend issues, we can't just assign a zeroed struct, so manually reset all fields + // DX12-start-strip + *shot = ShotData(); + // DX12-end-strip + shot.shot_id = shot_id; + + // After init, start execution from the first op + shot.next_op_idx = 0u; + + shot.rng_state.x[0] = rng_seed ^ hash_pcg(shot_id); + shot.rng_state.x[1] = rng_seed ^ hash_pcg(shot_id + 1); + shot.rng_state.x[2] = rng_seed ^ hash_pcg(shot_id + 2); + shot.rng_state.x[3] = rng_seed ^ hash_pcg(shot_id + 3); + shot.rng_state.x[4] = rng_seed ^ hash_pcg(shot_id + 4); + + shot.op_type = 0; + shot.op_idx = 0; + + // rand_* will be initialized in shot_init_per_op when preparing the first op + shot.duration = 0.0; + shot.renormalize = 1.0; + + shot.qubit_is_0_mask = (1u << u32(QUBIT_COUNT)) - 1u; // All qubits are |0> + shot.qubit_is_1_mask = 0u; + shot.qubits_updated_last_op_mask = 0; + shot.pending_loss_mask = 0u; + + // Initialize all qubit probabilities to 100% |0> + for (var i: i32 = 0; i < QUBIT_COUNT; i++) { + shot.qubit_state[i].zero_probability = 1.0; + shot.qubit_state[i].one_probability = 0.0; + shot.qubit_state[i].heat = 0.0; + shot.qubit_state[i].idle_since = 0.0; + } + + // unitary will be set in prepare_op +} + +fn update_qubit_state(shot_idx: u32) { + let shot = &shots[shot_idx]; + + // If any qubits were updated in the last op, we may need to sum workgroup probabilities into the shot state + // This is only needed if multiple workgroups were used for the shot execution. If not, then the + // single workgroup for the shot would have written directly to the shot state already. + + // For each qubit that was updated in the last op + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + let qubit_mask: u32 = 1u << q; + if ((shot.qubits_updated_last_op_mask & qubit_mask) != 0u) { + // Sum the workgroup collation entries for this qubit into the shot state + // Note: We ignore the fact a qubit may be 'lost' here. It should already be + // in the |0> state if lost, so summing the probabilities is still valid. + var total_zero: f32 = 0.0; + var total_one: f32 = 0.0; + + if (WORKGROUPS_PER_SHOT > 1) { + // Offset into workgroup collation buffer based on shot index + let offset = shot_idx * u32(WORKGROUPS_PER_SHOT); + for (var wkg_idx: u32 = 0u; wkg_idx < u32(WORKGROUPS_PER_SHOT); wkg_idx++) { + let sums = workgroup_collation.sums[wkg_idx + offset]; + total_zero = total_zero + sums.qubits[q].x; + total_one = total_one + sums.qubits[q].y; + } + } else { + // Single workgroup per shot case - just read directly from the shot + total_zero = shot.qubit_state[q].zero_probability; + total_one = shot.qubit_state[q].one_probability; + } + + // Update the shot state with the summed probabilities + // Round to 0 or 1 if extremely close to mitigate minor floating point errors + // TODO: Use PROB_THRESHOLD constant here? + if (total_zero < 0.000001) { total_zero = 0.0; } + if (total_one < 0.000001) { total_one = 0.0; } + if (total_zero > 0.999999) { total_zero = 1.0; } + if (total_one > 0.999999) { total_one = 1.0; } + + shot.qubit_state[q].zero_probability = total_zero; + shot.qubit_state[q].one_probability = total_one; + + // NOTE: Any kind of operation with a NaN float value results in a NaN, or false for logical comparisons + // So beware of conditions that may not behave as expected if NaN values are possible. + let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; + if !within_threshold { + // Populate the diagnostics buffer, if not already set + let old_value = atomicCompareExchangeWeak( + &diagnostics.error_code, + 0u, + ERR_INVALID_PROBS); + if old_value.exchanged { + // This is the first error - fill in the details + diagnostics.extra1 = q; + diagnostics.extra2 = total_zero; + diagnostics.extra3 = total_one; + // DX12 backend has issues assigning structs. See https://github.com/gfx-rs/wgpu/issues/8552 + // DX12-start-strip + diagnostics.shot = *shot; + diagnostics.op = ops[shot.op_idx]; + // DX12-end-strip + } + // Store the error value (if none set already) + let err_index = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak( + &results[err_index], + 0u, + ERR_INVALID_PROBS); + } + + // Update the masks for definite states + shot.qubit_is_0_mask = select( + shot.qubit_is_0_mask & ~qubit_mask, + shot.qubit_is_0_mask | qubit_mask, + total_zero == 1.0); + shot.qubit_is_1_mask = select( + shot.qubit_is_1_mask & ~qubit_mask, + shot.qubit_is_1_mask | qubit_mask, + total_one == 1.0); + } + } +} + +// Build a measure-and-reset (or measure-only) instrument for `qubit` given a +// measured `result`, store it in the shot buffer, set up renormalization, and +// mark the qubit as no longer in a definite basis state so the execute stage +// recomputes its probabilities. Shared by `prep_measure_reset` and +// `prep_loss_commit`; the caller sets `shot.op_idx` and `shot.op_type`. +fn prep_measure_reset_instrument(shot_idx: u32, qubit: u32, result: u32, resets_to_zero: bool) { + let shot = &shots[shot_idx]; + + // Construct the measurement/reset instrument based on the measured result + // Put the instrument into the shot buffer for the execute_op stage to apply + if resets_to_zero { + // Reset variants (MResetZ, ResetZ): + // Result=0: [[1,0],[0,0]] - project onto |0⟩ (already there) + // Result=1: [[0,1],[0,0]] - swap |1⟩ into |0⟩ slot (reset) + shot.unitary[0] = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); + shot.unitary[1] = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); + shot.unitary[4] = vec2f(); + shot.unitary[5] = vec2f(); + } else { + // Measure-only (MZ): + // Result=0: [[1,0],[0,0]] - project onto |0⟩ + // Result=1: [[0,0],[0,1]] - project onto |1⟩ (keep in place) + shot.unitary[0] = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); + shot.unitary[1] = vec2f(); + shot.unitary[4] = vec2f(); + shot.unitary[5] = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); + } + + shot.renormalize = select( + 1.0 / sqrt(shot.qubit_state[qubit].zero_probability), + 1.0 / sqrt(shot.qubit_state[qubit].one_probability), + result == 1u); + + // We don't want the measurement pass to skip over this qubit, so ensure it's marked as not in a definite state + shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << qubit); + shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << qubit); + + // Set the qubits_updated_last_op_mask to all except those that were already in a definite + // state (so we don't waste time updating probabilities that are already known). Note that + // next 'prepare_op' should set the just measured qubit into a definite 0 or 1 state. + shot.qubits_updated_last_op_mask = + // A mask with all qubits set + ((1u << u32(QUBIT_COUNT)) - 1u) + // Exclude qubits already in definite states + & ~(shot.qubit_is_0_mask | shot.qubit_is_1_mask); +} + +fn prep_measure_reset(shot_idx: u32, op_idx: u32, is_loss: bool, stores_result: bool, resets_to_zero: bool) { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + let qubit = get_measure_qubit(shot_idx, op_idx); + + // Choose measurement result based on qubit probabilities and random number + let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); + + // If this is being called due to loss noise, we don't write the result back to the results buffer + // Instead, mark the qubit as lost by setting the heat to -1.0 + if !is_loss { + if stores_result { + let result_id = get_measure_result(shot_idx, op_idx); // Result id to store the measurement result in is stored in q2 + + // If the qubit is already marked as lost, just report that and exit. It's already in the zero + // state so nothing to update or renormalize. The execute op should be a no-op (ID) + if shot.qubit_state[qubit].heat == -1.0 { + atomicStore(&results[(shot_idx * RESULT_COUNT) + result_id], 2u); + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + // Qubit get reloaded after a Measurement, so set the heat back to 0.0 + shot.qubit_state[qubit].heat = 0.0; + return; + } else { + atomicStore(&results[(shot_idx * RESULT_COUNT) + result_id], result); + } + } else { + // No result to store (e.g. ResetZ). If the qubit is lost, it's already in the zero + // state so nothing to update. Just set to ID and return. + if shot.qubit_state[qubit].heat == -1.0 { + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + return; + } + } + } else { + shot.qubit_state[qubit].heat = -1.0; + } + + prep_measure_reset_instrument(shot_idx, qubit, result, resets_to_zero); + + shot.op_idx = op_idx; + // Use OPID_MRESETZ as the op_type for all three variants in execute stage + // (they all use the same matrix-apply + update_all_qubit_probs path) + shot.op_type = OPID_MRESETZ; +} + +// Starting from the given index, return the next index if pauli noise, else 0 +fn get_pauli_noise_idx(op_idx: u32) -> u32 { + if (arrayLength(&ops) > (op_idx + 1)) { + let op = &ops[op_idx + 1]; + if (op.id == OPID_PAULI_NOISE_1Q || op.id == OPID_PAULI_NOISE_2Q) { + return op_idx + 1u; + } + } + return 0u; +} + +// From the starting index given, return the next index if loss noise, else 0 +fn get_loss_idx(op_idx: u32) -> u32 { + if (arrayLength(&ops) > (op_idx + 1)) { + let op = &ops[op_idx + 1]; + if (op.id == OPID_LOSS_NOISE) { + return op_idx + 1u; + } + } + return 0u; +} + +// Returns true if the gate at `op_idx` touches at least one lost qubit. +// `q1`/`q2` are the (resolved) operands of the gate. +fn gate_has_lost_operand(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) -> bool { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + if (shot.qubit_state[q1].heat == -1.0) { + return true; + } + let is_2q = !is_1q_op(op.id); + return is_2q && (shot.qubit_state[q2].heat == -1.0); +} + +// Builds a 4x4 (in shot.unitary) that applies the 1-qubit matrix `m` (given as +// m00,m01,m10,m11) to `target_is_q2 ? q2 : q1` and identity to the other qubit +// of the pair. The lost qubit is in the |0> state, so the identity factor keeps +// it there. The 2-qubit basis is |q1 q2>, so the row/col index is +// (2 * q1_bit + q2_bit). +fn set_1q_on_pair_unitary(shot_idx: u32, target_is_q2: bool, + m00: vec2f, m01: vec2f, m10: vec2f, m11: vec2f) { + let shot = &shots[shot_idx]; + // Zero the whole 4x4 first. + for (var i = 0u; i < 16u; i++) { + shot.unitary[i] = vec2f(0.0, 0.0); + } + if target_is_q2 { + // Acts on q2 (low bit): block-diagonal diag(M, M). + // Top-left block (q1 = 0): + shot.unitary[0] = m00; shot.unitary[1] = m01; + shot.unitary[4] = m10; shot.unitary[5] = m11; + // Bottom-right block (q1 = 1): + shot.unitary[10] = m00; shot.unitary[11] = m01; + shot.unitary[14] = m10; shot.unitary[15] = m11; + } else { + // Acts on q1 (high bit): M (x) I. + shot.unitary[0] = m00; shot.unitary[2] = m01; + shot.unitary[8] = m10; shot.unitary[10] = m11; + shot.unitary[5] = m00; shot.unitary[7] = m01; + shot.unitary[13] = m10; shot.unitary[15] = m11; + } +} + +// Multiplies one row of the 4x4 pair unitary (in shot.unitary) by -i, in place. +// Folding a diag(1, -i) = S-dagger factor on one qubit into a 2-qubit matrix +// scales the rows whose target-qubit bit is 1 by -i. For a complex entry +// (x + y i), (x + y i) * -i = y - x i. +fn scale_pair_unitary_row_by_neg_i(shot_idx: u32, row: u32) { + let shot = &shots[shot_idx]; + for (var c = 0u; c < 4u; c++) { + let e = shot.unitary[row * 4u + c]; + shot.unitary[row * 4u + c] = vec2f(e.y, -e.x); + } +} + +// Sets up the shot to execute a 2-qubit shot-buffer op on the gate's operands. +fn finish_2q_shot_buffer(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { + let shot = &shots[shot_idx]; + shot.op_idx = op_idx; + shot.op_type = OPID_SHOT_BUFF_2Q; + shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); +} + +// Loses a single surviving `qubit` for the PROPAGATE policy: samples a +// measurement outcome, collapses the qubit to that outcome and resets it to +// |0>, and marks it lost (heat = -1.0). The collapse is expressed as a 2-qubit +// tensor on the gate's operands (reset on `qubit`, identity on the lost +// partner, which is already in |0>), reusing the standard shot-buffer execute +// path. `qubit` must be one of the gate's two operands `q1`/`q2`. +fn propagate_loss_to_qubit(shot_idx: u32, op_idx: u32, q1: u32, q2: u32, qubit: u32) { + let shot = &shots[shot_idx]; + + let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); + + // Reset instrument (project + move |1> into |0> slot), same as MResetZ: + // result==0: [[1,0],[0,0]] + // result==1: [[0,1],[0,0]] + let m00 = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); + let m01 = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); + let m10 = vec2f(0.0, 0.0); + let m11 = vec2f(0.0, 0.0); + + let target_is_q2 = (qubit == q2); + set_1q_on_pair_unitary(shot_idx, target_is_q2, m00, m01, m10, m11); + + // Renormalize by the measured branch probability. + shot.renormalize = select( + 1.0 / sqrt(shot.qubit_state[qubit].zero_probability), + 1.0 / sqrt(shot.qubit_state[qubit].one_probability), + result == 1u); + + // Mark the qubit lost and clear its definite-state bits so the probability + // pass recomputes it. + shot.qubit_state[qubit].heat = -1.0; + shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << qubit); + shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << qubit); + + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); +} + +// Records an error `code` for `shot_idx` in both the diagnostics buffer and the +// shot's result-code slot, mirroring the reporting done elsewhere in this file. +// Used for conditions the host guarantees never occur (e.g. a loss policy that +// is not valid for a given gate). +fn report_shot_error(shot_idx: u32, code: u32) { + atomicCompareExchangeWeak(&diagnostics.error_code, 0u, code); + let err_index = (shot_idx + 1u) * RESULT_COUNT - 1u; + atomicCompareExchangeWeak(&results[err_index], 0u, code); +} + +// Handles a gate whose operand(s) include at least one lost qubit, according to +// the loss policy stamped on the op's `policy` field. `q1`/`q2` are the +// (resolved) operands. The gate body is fully handled here (degraded unitary, +// loss propagation, or turned into Id); the caller must not run the original +// gate afterwards. Any attached Pauli noise is applied separately to the +// surviving operand via `apply_2q_pauli_noise_on_survivor`. +fn handle_lost_operand_policy(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + let is_1q = is_1q_op(op.id); + let is_2q = !is_1q; + let policy = op.policy; + + // Loss policies only make sense for multi-qubit gates. + // If this is a single-qubit gate, skip it entirely. + if (is_1q) { + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + return; + } + + let q1_lost = shot.qubit_state[q1].heat == -1.0; + let q2_lost = is_2q && (shot.qubit_state[q2].heat == -1.0); + let has_survivor = is_2q && !(q1_lost && q2_lost); + // The surviving operand (only meaningful when has_survivor is true). + let survivor = select(q1, q2, q1_lost); + let survivor_is_q2 = q1_lost; + + // SWAP is special: it physically relocates the two qubits, so their loss + // state is always exchanged regardless of the policy (the policy only + // governs whether the unitary runs). Handle it explicitly here. + if (op.id == OPID_SWAP) { + switch policy { + case LOSS_POLICY_PROPAGATE { + propagate_loss_to_qubit(shot_idx, op_idx, q1, q2, survivor); + return; + } + case LOSS_POLICY_RESIDUAL_S_DAGGER { + // Match the CPU/stabilizer SWAP + residual S-dagger semantics: + // 1. Apply the full SWAP (shot.unitary already holds it). + // 2. Apply S-dagger = diag(1, -i) to the (originally) lost + // operand's position, which after the SWAP holds the + // survivor's amplitudes. + // 3. Exchange the per-qubit loss flag (heat) of the operands. + + // Fold the S-dagger into the SWAP matrix by scaling, by -i, the + // two rows of the |q1 q2> pair matrix whose lost-qubit bit is 1. + // q1 is the high bit (rows 2, 3); q2 is the low bit (rows 1, 3). + let lost_row = select(1u, 2u, q1_lost); + scale_pair_unitary_row_by_neg_i(shot_idx, lost_row); + scale_pair_unitary_row_by_neg_i(shot_idx, 3u); + // Exchange the per-qubit loss flag (heat) of the two operands. + let heat1 = shot.qubit_state[q1].heat; + shot.qubit_state[q1].heat = shot.qubit_state[q2].heat; + shot.qubit_state[q2].heat = heat1; + // The 2-qubit execute path skips amplitudes for qubits known to be + // in a definite state, which would skip the amplitudes SWAP needs to move. + // Clear those bits for both operands so the swap is actually applied. + shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~((1u << q1) | (1u << q2)); + shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~((1u << q1) | (1u << q2)); + // shot.unitary now holds (S-dagger on lost) * SWAP. + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); + return; + } + case LOSS_POLICY_APPLY_ANYWAY { + // Exchange the per-qubit loss flag (heat) of the two operands. + let heat1 = shot.qubit_state[q1].heat; + shot.qubit_state[q1].heat = shot.qubit_state[q2].heat; + shot.qubit_state[q2].heat = heat1; + // The 2-qubit execute path skips amplitudes for qubits known to be + // in a definite state, which would skip the amplitudes SWAP needs to move. + // Clear those bits for both operands so the swap is actually applied. + shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~((1u << q1) | (1u << q2)); + shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~((1u << q1) | (1u << q2)); + // shot.unitary already holds the SWAP matrix (set by the caller). + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); + return; + } + case LOSS_POLICY_SKIP { + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + return; + } + default { + // SWAP only supports SKIP, PROPAGATE, RESIDUAL_S_DAGGER, and + // APPLY_ANYWAY. Any other policy (e.g. DEGRADE) is rejected by + // the host, so reaching here indicates a bug. + report_shot_error(shot_idx, ERR_UNSUPPORTED_LOSS_POLICY); + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + return; + } + } + } + + // APPLY_ANYWAY is only valid for SWAP, which is handled above. Reaching here + // with it on any other gate is rejected by the host, so it indicates a bug. + if (policy == LOSS_POLICY_APPLY_ANYWAY) { + report_shot_error(shot_idx, ERR_UNSUPPORTED_LOSS_POLICY); + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + return; + } + + if (policy == LOSS_POLICY_PROPAGATE && has_survivor) { + propagate_loss_to_qubit(shot_idx, op_idx, q1, q2, survivor); + return; + } + + if (policy == LOSS_POLICY_RESIDUAL_S_DAGGER && has_survivor) { + // Apply S-dagger = diag(1, -i) to the surviving operand. + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(1.0, 0.0), vec2f(0.0, 0.0), + vec2f(0.0, 0.0), vec2f(0.0, -1.0)); + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); + return; + } + + // DEGRADE is only valid for the two-qubit rotations (Rxx/Ryy/Rzz), so the + // op is guaranteed to be one of them when a survivor exists. + if (policy == LOSS_POLICY_DEGRADE && has_survivor) { + // Degrade the two-qubit rotation to its single-qubit version on the + // survivor. The op's unitary[0] holds cos(θ/2) for Rxx/Ryy; we recover + // the angle to build the 1-qubit rotation matrix. + let cos_half = op.unitary[0].x; + if (op.id == OPID_RXX) { + // Rx(θ) = [[c, -i s], [-i s, c]], where s = sin(θ/2). + let s = op.unitary[3].y * -1.0; // unitary[3] = (0, -sin(θ/2)) + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(cos_half, 0.0), vec2f(0.0, -s), + vec2f(0.0, -s), vec2f(cos_half, 0.0)); + } else if (op.id == OPID_RYY) { + // Ry(θ) = [[c, -s], [s, c]], where s = sin(θ/2). + let s = op.unitary[3].y; // unitary[3] = (0, sin(θ/2)) for Ryy + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(cos_half, 0.0), vec2f(-s, 0.0), + vec2f(s, 0.0), vec2f(cos_half, 0.0)); + } else { + // Rzz -> Rz(θ). The GPU Rz convention is [[1, 0], [0, e^{iθ}]], + // and unitary[5] = e^{iθ} holds the full-angle phase. + let phase = op.unitary[5]; + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(1.0, 0.0), vec2f(0.0, 0.0), + vec2f(0.0, 0.0), phase); + } + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); + return; + } + + // SKIP, or any policy when both operands are lost (no survivor to act on): + // skip the gate entirely. + shot.op_type = OPID_ID; + shot.op_idx = op_idx; +} + +fn apply_1q_pauli_noise(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32) { + // NOTE: Assumes that whatever prepared the program ensured that noise_op.q1 matches op.q1 and + // that op is a 1-qubit gate. `q1` is the resolved target qubit (may be + // dynamic for the adaptive interpreter, where op.q1 is only a placeholder). + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + let noise_op = &ops[noise_idx]; + + // Categorical outcome probabilities by 3-bit term (X=1, Z=2, Y=3, L=4), + // stored at flat slot k = term in `unitary[k / 2][k % 2]`. The identity + // outcome (slot 0) is implicit. + let p_x = noise_op.unitary[0].y; + let p_z = noise_op.unitary[1].x; + let p_y = noise_op.unitary[1].y; + let p_loss = noise_op.unitary[2].x; + + shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer + + let rand = shot.rand_pauli; + if (rand < p_x) { + // Apply the X permutation (basically swap the rows) + shot.unitary[0] = op.unitary[4]; + shot.unitary[1] = op.unitary[5]; + shot.unitary[4] = op.unitary[0]; + shot.unitary[5] = op.unitary[1]; + } else if (rand < (p_x + p_y)) { + // Apply the Y permutation (swap rows with negated |0> state) + shot.unitary[0] = cplxNeg(op.unitary[4]); + shot.unitary[1] = cplxNeg(op.unitary[5]); + shot.unitary[4] = op.unitary[0]; + shot.unitary[5] = op.unitary[1]; + } else if (rand < (p_x + p_y + p_z)) { + // Apply Z error (negate |1> state) + shot.unitary[0] = op.unitary[0]; + shot.unitary[1] = op.unitary[1]; + shot.unitary[4] = cplxNeg(op.unitary[4]); + shot.unitary[5] = cplxNeg(op.unitary[5]); + } else { + // Either loss or no noise: the gate executes unmodified. If loss was + // sampled, schedule a loss commit for this qubit; a following + // loss-commit op performs the measure + reset. + if (rand < (p_x + p_z + p_y + p_loss)) { + shot.pending_loss_mask |= (1u << q1); + } + // No noise. Set the op_type back to the op.id value if it's Id, MResetZ, MZ, or ResetZ, as they get handled specially in execute_op + if (op.id == OPID_ID || op.id == OPID_MRESETZ || op.id == OPID_MZ || op.id == OPID_RESETZ) { + shot.op_type = op.id; + } + if (is_1q_phase_gate(op.id)) { + // For phase gates, treat everything as RZ for execution purposes + shot.op_type = OPID_RZ; + } + } + + shot.op_idx = op_idx; + if (shot.op_type == OPID_ID || shot.op_type == OPID_RZ) { + shot.qubits_updated_last_op_mask = 0u; + } else { + shot.qubits_updated_last_op_mask = 1u << q1; + }; +} + +fn apply_2q_pauli_noise(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32, q2: u32) { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + let noise_op = &ops[noise_idx]; + + // The categorical distribution over the 25 (q1_term, q2_term) outcomes is + // stored at flat slot k = q1_term * 5 + q2_term in `unitary[k / 2][k % 2]`. + // Terms use the 3-bit encoding: I=0, X=1, Z=2, Y=3, L=4. The II slot (0) is + // implicit and carries the remaining probability. + var rand = shot.rand_pauli; + var q1_term = 0; + var q2_term = 0; + + // Find the terms to apply based on the random number and the probabilities + for (var a = 0; a < 5; a = a + 1) { + for (var b = 0; b < 5; b = b + 1) { + let k = a * 5 + b; + if (k == 0) { continue; } // II carries no stored probability + let slot = noise_op.unitary[k / 2]; + let p_ab = select(slot.x, slot.y, (k & 1) == 1); + if (rand < p_ab) { + q1_term = a; + q2_term = b; + // Break out of both loops + a = 5; + b = 5; + } else { + rand = rand - p_ab; + } + } + } + + // Schedule loss commits for any qubit whose sampled term is loss (L = 4). + // A following loss-commit op performs the measure + reset. + if (q1_term == 4) { shot.pending_loss_mask |= (1u << q1); } + if (q2_term == 4) { shot.pending_loss_mask |= (1u << q2); } + + // A Pauli fault (X, Z, Y = 1, 2, 3) is fused into the gate by permuting its + // rows. Loss (4) and identity (0) leave the gate unmodified for that qubit. + let q1_pauli = q1_term >= 1 && q1_term <= 3; + let q2_pauli = q2_term >= 1 && q2_term <= 3; + + if (q1_pauli || q2_pauli) { + // Get the rows of the 2 qubit unitary + var op_row_0 = getOpRow(op_idx, 0); + var op_row_1 = getOpRow(op_idx, 1); + var op_row_2 = getOpRow(op_idx, 2); + var op_row_3 = getOpRow(op_idx, 3); + + // Apply the Paulis to the matrices. Note this is just permuting the rows, and appliction + // commutes, so we can apply them in any order. High order bit is q1. Low order bit is q2. + // X on q1 is rows 2<>0 and 3<>1, X on q2 is rows 1<>0 and 3<>2, etc. + // Y on q1 is rows -2<>0 and -3<>1, Y on q2 is rows -1<>0 and -3<>2 + // Z on q1 is -2 and -3, Z on q2 is -1 and -3 + + // Apply the q1 permutations as needed + if (q1_term == 1) { + // Apply the X permutation + let old_row_0 = op_row_0; + let old_row_1 = op_row_1; + op_row_0 = op_row_2; + op_row_1 = op_row_3; + op_row_2 = old_row_0; + op_row_3 = old_row_1; + } else if (q1_term == 3) { + // Apply the Y permutation + let old_row_0 = op_row_0; + let old_row_1 = op_row_1; + op_row_0 = rowNeg(op_row_2); + op_row_1 = rowNeg(op_row_3); + op_row_2 = old_row_0; + op_row_3 = old_row_1; + } else if (q1_term == 2) { + // Apply Z permutation + op_row_2 = rowNeg(op_row_2); + op_row_3 = rowNeg(op_row_3); + } + // Apply the q2 permutations as needed + if (q2_term == 1) { + // Apply the X permutation + let old_row_0 = op_row_0; + let old_row_2 = op_row_2; + op_row_0 = op_row_1; + op_row_2 = op_row_3; + op_row_1 = old_row_0; + op_row_3 = old_row_2; + } else if (q2_term == 3) { + // Apply the Y permutation + let old_row_0 = op_row_0; + let old_row_2 = op_row_2; + op_row_0 = rowNeg(op_row_1); + op_row_2 = rowNeg(op_row_3); + op_row_1 = old_row_0; + op_row_3 = old_row_2; + } else if (q2_term == 2) { + // Apply Z permutation + op_row_1 = rowNeg(op_row_1); + op_row_3 = rowNeg(op_row_3); + } + // Write the rows back to the shot buffer unitary + setUnitaryRow(shot_idx, 0u, op_row_0); + setUnitaryRow(shot_idx, 1u, op_row_1); + setUnitaryRow(shot_idx, 2u, op_row_2); + setUnitaryRow(shot_idx, 3u, op_row_3); + shot.op_type = OPID_SHOT_BUFF_2Q; + } else { + // No Pauli fault to fuse (identity or loss only). Leave if CX, CY, CZ, or RZZ as they get handled specially in execute_op + if (op.id == OPID_CX || op.id == OPID_CY || op.id == OPID_CZ || op.id == OPID_RZZ) { + shot.op_type = op.id; + } else { + shot.op_type = OPID_SHOT_BUFF_2Q; + } + } + shot.op_idx = op_idx; + if (shot.op_type == OPID_CZ || shot.op_type == OPID_RZZ) { + shot.qubits_updated_last_op_mask = 0u; + } else { + shot.qubits_updated_last_op_mask = (1u << q1 ) | (1u << q2); + } +} + +// Left-multiplies the 4x4 pair unitary already in `shot.unitary` by a single +// Pauli (term: X=1, Z=2, Y=3) acting on `target_is_q2 ? q2 : q1`. This is the +// same row permutation/negation that `apply_2q_pauli_noise` fuses, just applied +// to the policy-degraded gate rather than the original op. Note the Y branch +// uses real signs (i.e. -i*Y), matching `apply_2q_pauli_noise`; the resulting +// global phase is unobservable for a Pauli noise channel. +fn fuse_1q_pauli_on_pair_unitary(shot_idx: u32, target_is_q2: bool, term: u32) { + let si = i32(shot_idx); + var row_0 = getUnitaryRow(si, 0u); + var row_1 = getUnitaryRow(si, 1u); + var row_2 = getUnitaryRow(si, 2u); + var row_3 = getUnitaryRow(si, 3u); + + if (!target_is_q2) { + // Acting on q1 (high bit): rows {0,1} <-> {2,3}. + if (term == 1u) { // X + let o0 = row_0; let o1 = row_1; + row_0 = row_2; row_1 = row_3; + row_2 = o0; row_3 = o1; + } else if (term == 3u) { // Y + let o0 = row_0; let o1 = row_1; + row_0 = rowNeg(row_2); row_1 = rowNeg(row_3); + row_2 = o0; row_3 = o1; + } else { // Z + row_2 = rowNeg(row_2); row_3 = rowNeg(row_3); + } + } else { + // Acting on q2 (low bit): rows {0,2} <-> {1,3}. + if (term == 1u) { // X + let o0 = row_0; let o2 = row_2; + row_0 = row_1; row_2 = row_3; + row_1 = o0; row_3 = o2; + } else if (term == 3u) { // Y + let o0 = row_0; let o2 = row_2; + row_0 = rowNeg(row_1); row_2 = rowNeg(row_3); + row_1 = o0; row_3 = o2; + } else { // Z + row_1 = rowNeg(row_1); row_3 = rowNeg(row_3); + } + } + + setUnitaryRow(shot_idx, 0u, row_0); + setUnitaryRow(shot_idx, 1u, row_1); + setUnitaryRow(shot_idx, 2u, row_2); + setUnitaryRow(shot_idx, 3u, row_3); +} + +// Applies the Pauli noise attached to a 2-qubit gate that had a lost operand. +// The gate body itself was already handled by `handle_lost_operand_policy` +// (which may have left a degraded 4x4 in `shot.unitary`, or turned the gate +// into Id for SKIP). This mirrors the CPU `apply_fault`: the joint (q1, q2) +// term is sampled, but only the operand still alive *after* the policy ran +// receives its term; a lost operand gets nothing. +// +// Because this is only reached when the gate has at least one lost operand, +// there is at most one surviving operand, so at most one single-qubit Pauli is +// fused. +fn apply_2q_pauli_noise_on_survivor(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32, q2: u32) { + let shot = &shots[shot_idx]; + let noise_op = &ops[noise_idx]; + + // Surviving operand(s) after the policy ran (alive => heat != -1.0). + let q1_alive = shot.qubit_state[q1].heat != -1.0; + let q2_alive = shot.qubit_state[q2].heat != -1.0; + // Both lost (e.g. PROPAGATE collapsed the survivor): nothing to apply. + if (!q1_alive && !q2_alive) { + return; + } + + // Sample the joint (q1_term, q2_term) outcome (same encoding/layout as + // apply_2q_pauli_noise: I=0, X=1, Z=2, Y=3, L=4). + var rand = shot.rand_pauli; + var q1_term = 0; + var q2_term = 0; + for (var a = 0; a < 5; a = a + 1) { + for (var b = 0; b < 5; b = b + 1) { + let k = a * 5 + b; + if (k == 0) { continue; } + let slot = noise_op.unitary[k / 2]; + let p_ab = select(slot.x, slot.y, (k & 1) == 1); + if (rand < p_ab) { + q1_term = a; + q2_term = b; + a = 5; + b = 5; + } else { + rand = rand - p_ab; + } + } + } + + // The survivor's own term. (At most one operand is alive here.) + let survivor_is_q2 = !q1_alive; + let survivor = select(q1, q2, survivor_is_q2); + let term = select(q1_term, q2_term, survivor_is_q2); + + // Loss (4): schedule a loss commit for the survivor; a later loss-commit op + // performs the measure + reset. The gate set up by the policy still runs. + if (term == 4) { + shot.pending_loss_mask |= (1u << survivor); + return; + } + + // Identity (0): nothing to fuse; leave the policy's setup untouched. + if (term == 0) { + return; + } + + // Pauli (X=1, Z=2, Y=3): fuse onto the survivor. + if (shot.op_type == OPID_SHOT_BUFF_2Q) { + // The policy left a degraded 4x4 in shot.unitary; left-multiply it by + // the survivor Pauli. + fuse_1q_pauli_on_pair_unitary(shot_idx, survivor_is_q2, u32(term)); + } else { + // The policy turned the gate into Id (SKIP). Build a pair unitary that + // applies just the Pauli to the survivor and identity to the lost + // partner (which is in |0>). Real-sign Y matches the fuse path above. + if (term == 1) { // X + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(0.0, 0.0), vec2f(1.0, 0.0), + vec2f(1.0, 0.0), vec2f(0.0, 0.0)); + } else if (term == 3) { // Y (real-sign, i.e. -i*Y) + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(0.0, 0.0), vec2f(-1.0, 0.0), + vec2f(1.0, 0.0), vec2f(0.0, 0.0)); + } else { // Z + set_1q_on_pair_unitary(shot_idx, survivor_is_q2, + vec2f(1.0, 0.0), vec2f(0.0, 0.0), + vec2f(0.0, 0.0), vec2f(-1.0, 0.0)); + } + finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); + } + + // The survivor's amplitudes may have been in a definite computational-basis + // state; clear its definite-state bits so the execute pass recomputes them + // after the Pauli (mirrors the SWAP handling in handle_lost_operand_policy). + shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << survivor); + shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << survivor); +} + +fn get_shot_params( + workgroupId: u32, + tid: u32, + op_qubit_count: i32) -> ShotParams { + // Workgroups are per shot if 22 or less qubits, else 2 workgroups for 23 qubits, 4 for 24, etc.. + let shot_idx: i32 = i32(workgroupId) / WORKGROUPS_PER_SHOT; + let shot_state_vector_start: i32 = shot_idx * (1i << u32(QUBIT_COUNT)); + let workgroup_idx_in_shot: i32 = i32(workgroupId) % WORKGROUPS_PER_SHOT; + let thread_idx_in_shot: i32 = workgroup_idx_in_shot * THREADS_PER_WORKGROUP + i32(tid); + let total_threads_per_shot: i32 = WORKGROUPS_PER_SHOT * THREADS_PER_WORKGROUP; + + // If using multiple workgroups per shot, each workgroup will write its partial sums to the collation + // buffer for later summing by the prepare_op stage. If single workgroup per shot, no collation needed. + // Use -1 as a marker for single workgroup per shot case (in which case we should write directly to the shot). + let workgroup_collation_idx: i32 = select(-1, i32(workgroupId), WORKGROUPS_PER_SHOT > 1); + + let zero_entry_count: i32 = (1i << u32(QUBIT_COUNT)) >> u32(op_qubit_count); + let op_iterations: i32 = zero_entry_count / total_threads_per_shot; + + return ShotParams( + shot_idx, + shot_state_vector_start, + workgroup_collation_idx, + workgroup_idx_in_shot, + thread_idx_in_shot, + total_threads_per_shot, + zero_entry_count, + op_iterations + ); +} + +fn apply_1q_op(workgroupId: u32, tid: u32, q1: u32) { + let params = get_shot_params(workgroupId, tid, 1 /* qubits per op */); + let shot = &shots[params.shot_idx]; + let scale = shot.renormalize; + let lowMask = (1u << q1) - 1; + let highMask = (1u << u32(QUBIT_COUNT)) - 1 - lowMask; + let qubit_is_0_mask = i32(shots[params.shot_idx].qubit_is_0_mask); + let qubit_is_1_mask = i32(shots[params.shot_idx].qubit_is_1_mask); + + var summed_probs: vec4f = vec4f(); + + /* This loop is where all the real work happens. Try to keep this tight and efficient. + + We want a 'structure of arrays' like access pattern here for efficiency, so we process the state vector + in blocks where each thread in the workgroup(s) handle an adjacent entry to be processed. + + Each thread should start at the state vector shot start + 'thread_idx_in_shot', which is sequential across the workgroup threads + Each next entry for the thread is WORKGROUPS_PER_SHOT * THREADS_PER_WORKGROUP away. + */ + var entry_index = params.thread_idx_in_shot; + + for (var i = 0; i < params.op_iterations; i++) { + let offset0: i32 = (entry_index & lowMask) | ((entry_index & highMask) << 1); + let offset1: i32 = offset0 | (1 << q1); + + // See if we can skip doing any work for this pair, because the state vector entries to processes + // are both definitely 0.0, as we know they are for states where other qubits are in definite opposite state. + let skip_processing = ((offset0 & qubit_is_0_mask) != 0) || ((~offset1 & qubit_is_1_mask) != 0); + + if (!skip_processing) { + if shot.op_type == OPID_RZ { + // For RZ, we can skip reading/writing the |0> amplitude, as it is unchanged. + // Just apply the phase to the |1> amplitude. Probabilities also don't change. + let amp1: vec2f = stateVector[params.shot_state_vector_start + offset1]; + let new1 = cplxMul(amp1, shot.unitary[5]); + stateVector[params.shot_state_vector_start + offset1] = new1; + } else { + let amp0: vec2f = stateVector[params.shot_state_vector_start + offset0]; + let amp1: vec2f = stateVector[params.shot_state_vector_start + offset1]; + + let new0 = scale * (cplxMul(amp0, shot.unitary[0]) + cplxMul(amp1, shot.unitary[1])); + let new1 = scale * (cplxMul(amp0, shot.unitary[4]) + cplxMul(amp1, shot.unitary[5])); + + stateVector[params.shot_state_vector_start + offset0] = new0; + stateVector[params.shot_state_vector_start + offset1] = new1; + + if shot.op_type == OPID_MRESETZ || shot.op_type == OPID_LOSS_NOISE || scale != 1.0 { + // For MResetZ, loss-commit, or renormalization, update the probabilities for all qubits + update_all_qubit_probs(u32(offset0), new0, tid); + update_all_qubit_probs(u32(offset1), new1, tid); + } else { + summed_probs[0] += cplxMag2(new0); + summed_probs[1] += cplxMag2(new1); + } + } + } + entry_index += params.total_threads_per_shot; + } + + if scale == 1.0 && shot.op_type != OPID_RZ && shot.op_type != OPID_MRESETZ && shot.op_type != OPID_LOSS_NOISE { + // Update this thread's totals for the two qubits in the workgroup storage + qubitProbabilities[tid].zero[q1] = summed_probs[0]; + qubitProbabilities[tid].one[q1] = summed_probs[1]; + } +} + +fn apply_2q_op(workgroupId: u32, tid: u32, q1: u32, q2: u32) { + let params = get_shot_params(workgroupId, tid, 2 /* qubits per op */); + let shot = &shots[params.shot_idx]; + let update_probs = shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; + + // Sometimes a 2-qubit op may be converted to a no-op (ID) due to qubit loss etc., so skip processing in that case + // Calculate masks to split the index into low, mid, and high bits around the two qubits + let lowQubit = select(q1, q2, q1 > q2); + let hiQubit = select(q1, q2, q1 < q2); + + // Number of bits in each section + let lowBitCount = lowQubit; + let midBitCount = hiQubit - lowQubit - 1; + let hiBitCount = u32(QUBIT_COUNT) - hiQubit - 1; + + // The masks below help extract the low, mid, and high bits from the counter to use around the two qubits locations + let lowMask = (1 << lowBitCount) - 1; + let midMask = (1 << (lowBitCount + midBitCount)) - 1 - lowMask; + let hiMask = (1 << u32(QUBIT_COUNT)) - 1 - midMask - lowMask; + + // Each iteration processes 4 amplitudes (the four affected by the 2-qubit gate), so quarter as many iterations as chunk size + var entry_index = params.thread_idx_in_shot; + var summed_probs: vec4f = vec4f(); + + for (var i = 0; i < params.op_iterations; i++) { + // q1 is the control, q2 is the target + let offset00: i32 = (entry_index & lowMask) | ((entry_index & midMask) << 1) | ((entry_index & hiMask) << 2); + let offset01: i32 = offset00 | (1 << q2); + let offset10: i32 = offset00 | (1 << q1); + let offset11: i32 = offset10 | (1 << q2); + + let can_skip_processing = + (((u32(offset00) & shot.qubit_is_0_mask) != 0) || + ((~(u32(offset11)) & shot.qubit_is_1_mask) != 0)); + if !can_skip_processing { + switch shot.op_type { + case OPID_CZ { + let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; + stateVector[params.shot_state_vector_start + offset11] = cplxNeg(amp11); + // CZ doesn't change any probabilities, so no need to update summed_probs + } + case OPID_RZZ { + // Firt and last entries are unchanged, only need to update the middle two + let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; + let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; + // Unitary matrix second entry in the second row is 5, third entry in the third row is 10 + stateVector[params.shot_state_vector_start + offset01] = cplxMul(amp01, shot.unitary[5]); + stateVector[params.shot_state_vector_start + offset10] = cplxMul(amp10, shot.unitary[10]); + } + case OPID_CX { + // Need to read all 4 to update the probabilities correctly, but only swap the |10> and |11> entries + let amp00: vec2f = stateVector[params.shot_state_vector_start + offset00]; + let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; + let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; + let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; + stateVector[params.shot_state_vector_start + offset10] = amp11; + stateVector[params.shot_state_vector_start + offset11] = amp10; + summed_probs[0] += (cplxMag2(amp00) + cplxMag2(amp01)); + summed_probs[1] += (cplxMag2(amp11) + cplxMag2(amp10)); + summed_probs[2] += (cplxMag2(amp00) + cplxMag2(amp11)); + summed_probs[3] += (cplxMag2(amp01) + cplxMag2(amp10)); + } + case OPID_CY { + // Like CX, but swap |10> and |11> with +/- i phases. + let amp00: vec2f = stateVector[params.shot_state_vector_start + offset00]; + let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; + let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; + let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; + stateVector[params.shot_state_vector_start + offset10] = vec2f(amp11.y, -amp11.x); // -i * |11> + stateVector[params.shot_state_vector_start + offset11] = vec2f(-amp10.y, amp10.x); // i * |10> + summed_probs[0] += (cplxMag2(amp00) + cplxMag2(amp01)); + summed_probs[1] += (cplxMag2(amp11) + cplxMag2(amp10)); + summed_probs[2] += (cplxMag2(amp00) + cplxMag2(amp11)); + summed_probs[3] += (cplxMag2(amp01) + cplxMag2(amp10)); + } + default { + // Assume OPID_SHOT_BUFF_2Q + // Get the state vector entries + let states = array( + stateVector[params.shot_state_vector_start + offset00], + stateVector[params.shot_state_vector_start + offset01], + stateVector[params.shot_state_vector_start + offset10], + stateVector[params.shot_state_vector_start + offset11] + ); + // Apply the unitary from the shot buffer + let result00 = innerProduct(getUnitaryRow(params.shot_idx, 0), states); + let result01 = innerProduct(getUnitaryRow(params.shot_idx, 1), states); + let result10 = innerProduct(getUnitaryRow(params.shot_idx, 2), states); + let result11 = innerProduct(getUnitaryRow(params.shot_idx, 3), states); + // Write back the results + stateVector[params.shot_state_vector_start + offset00] = result00; + stateVector[params.shot_state_vector_start + offset01] = result01; + stateVector[params.shot_state_vector_start + offset10] = result10; + stateVector[params.shot_state_vector_start + offset11] = result11; + // Update the probabilities for the acted on qubits + summed_probs[0] += (cplxMag2(result00) + cplxMag2(result01)); + summed_probs[1] += (cplxMag2(result10) + cplxMag2(result11)); + summed_probs[2] += (cplxMag2(result00) + cplxMag2(result10)); + summed_probs[3] += (cplxMag2(result01) + cplxMag2(result11)); + } + } + } + + entry_index += params.total_threads_per_shot; + } + + // Update this thread's totals for the two qubits in the workgroup storage + if (update_probs) { + // Update all for other 2-qubit gates + qubitProbabilities[tid].zero[q1] = summed_probs[0]; + qubitProbabilities[tid].one[q1] = summed_probs[1]; + qubitProbabilities[tid].zero[q2] = summed_probs[2]; + qubitProbabilities[tid].one[q2] = summed_probs[3]; + } +} + +fn apply_correlated_noise(workgroupId: u32, tid: u32) { + let params = get_shot_params(workgroupId, tid, 0 /* need to walk all entries */); + // Probabilities are already updated in the prepare_op stage + // Here we just need to apply the bit-flips and phase-flips to the state vector amplitudes + + let shot = &shots[params.shot_idx]; + + // Get the bit-flip and phase-flip masks from the shot buffer (stored by prep_correlated_noise) + let bit_flip_mask = bitcast(shot.unitary[0].x); + let phase_flip_mask = bitcast(shot.unitary[0].y); + + // If no flips to apply, early exit + if (bit_flip_mask == 0u && phase_flip_mask == 0u) { + return; + } + + var entry_index = params.thread_idx_in_shot; + + for (var i = 0; i < params.op_iterations; i++) { + // Get the target index to swap the state with by flipping the bits as indicated in the bit_flip_mask + let target_index = entry_index ^ i32(bit_flip_mask); + + // If there are an odd number of phase flips for the entry, we need to negate the amplitude + let negate_index: f32 = select(1.0, -1.0, (countOneBits(entry_index & i32(phase_flip_mask)) & 1) != 0); + + if (bit_flip_mask == 0u && negate_index == -1.0) { + // No bit flips to perform, but need to negate this entry (phase flip only) + stateVector[params.shot_state_vector_start + entry_index] = cplxNeg(stateVector[params.shot_state_vector_start + entry_index]); + } else if (entry_index < target_index) { + // Bit flips are happening (as the indices are different), but to avoid double swapping only handle the swap + // when entry_index < target_index (avoid reprocessing when later we encounter the target_index entry as the entry_index) + + let amp_entry: vec2f = stateVector[params.shot_state_vector_start + entry_index]; + let amp_target: vec2f = stateVector[params.shot_state_vector_start + target_index]; + + // If there are an odd number of phase flips for the target, we need to negate that amplitude too + let negate_target: f32 = select(1.0, -1.0, (countOneBits(target_index & i32(phase_flip_mask)) & 1) != 0); + + // Swap and apply any negations for phase flips. + // Note this only applies -1 & 1 to the phase, not -i and i as the 'canonical' Y gate does. + // However, this is sufficient for simulating noise, as the global phase doesn't matter. + stateVector[params.shot_state_vector_start + entry_index] = cplxMul(amp_target, vec2f(negate_index, 0.0)); + stateVector[params.shot_state_vector_start + target_index] = cplxMul(amp_entry, vec2f(negate_target, 0.0)); + } + + // Jump ahead to the next entry to process + entry_index += params.total_threads_per_shot; + } +} + +// For the state vector index and amplitude probability, update all the qubit probabilities for this thread +fn update_all_qubit_probs(stateVectorIndex: u32, amplitude: vec2f, tid: u32) { + var mask: u32 = 1u; + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + let is_one: bool = (stateVectorIndex & mask) != 0u; + let prob: f32 = cplxMag2(amplitude); + if (is_one) { + qubitProbabilities[tid].one[q] += prob; + } else { + qubitProbabilities[tid].zero[q] += prob; + } + mask = mask << 1u; + } +} + +fn sum_thread_totals_to_shot(q: u32, shot_idx: i32, wkg_collation_idx: i32) { + var total_zero: f32 = 0.0; + var total_one: f32 = 0.0; + for (var j = 0; j < THREADS_PER_WORKGROUP; j++) { + total_zero += qubitProbabilities[j].zero[q]; + total_one += qubitProbabilities[j].one[q]; + } + if (wkg_collation_idx >= 0) { + // Write to the workgroup collation buffer for later summation into the shot state + workgroup_collation.sums[wkg_collation_idx].qubits[q] = vec2f(total_zero, total_one); + } else { + // Single workgroup per shot case - write directly to the shot state + let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; + if !within_threshold { + // Populate the diagnostics buffer, if not already set + let old_value = atomicCompareExchangeWeak( + &diagnostics.error_code, + 0u, + ERR_INVALID_THREAD_TOTAL); + if old_value.exchanged { + // This is the first error - fill in the details + let shot = &shots[shot_idx]; + diagnostics.extra1 = q; + diagnostics.extra2 = total_zero; + diagnostics.extra3 = total_one; + // DX12 backend has issues copying structs. See https://github.com/gfx-rs/wgpu/issues/8552 + // DX12-start-strip + diagnostics.shot = *shot; + diagnostics.op = ops[shot.op_idx]; + // DX12-end-strip + } + let err_index = (shot_idx + 1) * i32(RESULT_COUNT) - 1; + atomicCompareExchangeWeak( + &results[err_index], + 0u, + ERR_INVALID_THREAD_TOTAL); + } else { + shots[shot_idx].qubit_state[q].zero_probability = total_zero; + shots[shot_idx].qubit_state[q].one_probability = total_one; + } + } +} + +// Samples the correlated noise table to determine whether noise should be applied, and if so, +// which Pauli string was selected. If no noise is applied, the shot is set to ID and the caller +// can return early. +fn sample_correlated_noise(shot_idx: u32, op_idx: u32, noise_table_idx: u32) -> CorrelatedNoiseSample { + let shot = &shots[shot_idx]; + let table = &batch_data.correlated_noise_tables[noise_table_idx]; + + // Generate a Q1.63 random number (two u32 values for lo and hi 32 bits) + // Mask off the high bit of rand_hi to ensure the value is in [0, 1) range + let rand_lo = next_rand_u32(shot_idx); + let rand_hi = next_rand_u32(shot_idx) & 0x7FFFFFFFu; + + // Get the total noise probability from the table metadata + let noise_prob_lo = table.noise_probability_lo; + let noise_prob_hi = table.noise_probability_hi; + + // Check if noise should be applied at all by comparing the random number against the total noise probability + // If rand >= noise_probability, then no noise is applied + if (rand_hi > noise_prob_hi || (rand_hi == noise_prob_hi && rand_lo >= noise_prob_lo)) { + // No noise to apply - set the op to ID + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + shot.qubits_updated_last_op_mask = 0u; + return CorrelatedNoiseSample(0u, 0u, 0u); + } + + // Noise should be applied - binary search to find which Pauli string to apply + let start = i32(table.start_offset); + let count = i32(table.entry_count); + let entry_idx = binary_search_noise_table(rand_lo, rand_hi, start, count); + let entry = &batch_data.correlated_noise_entries[start + entry_idx]; + + return CorrelatedNoiseSample(1u, entry.paulis_lo, entry.paulis_hi); +} + +// Extracts the 3-bit term value for qubit position `i` from a Pauli + loss string. +// Terms use the encoding I=0, X=1, Z=2, Y=3, L=4. The low two bits double as the +// bit-flip (0x1) and phase-flip (0x2) indicators, and 0x4 marks loss. +// The Rust parsing stores terms with the rightmost (last) character at the lowest +// bits, so for position i we read the 3 bits at (qubit_count - 1 - i) * 3. +fn get_pauli_bits(paulis_lo: u32, paulis_hi: u32, qubit_count: u32, i: u32) -> u32 { + let bit_position = (qubit_count - 1u - i) * 3u; + if (bit_position + 3u <= 32u) { + return (paulis_lo >> bit_position) & 0x7u; + } else if (bit_position >= 32u) { + return (paulis_hi >> (bit_position - 32u)) & 0x7u; + } else { + // The 3-bit term straddles the boundary between the lo and hi words. + let low_part = paulis_lo >> bit_position; + let high_part = paulis_hi << (32u - bit_position); + return (low_part | high_part) & 0x7u; + } +} + +// Commits correlated noise masks into the shot state: stores the masks, swaps probabilities and +// tracking bits for bit-flipped qubits, records any loss, and sets the shot up for the correlated +// noise execute stage. Qubits in `loss_mask` are scheduled for loss; following loss-commit ops +// perform the measure + reset. +fn commit_correlated_noise(shot_idx: u32, op_idx: u32, bit_flip_mask: u32, phase_flip_mask: u32, loss_mask: u32) { + let shot = &shots[shot_idx]; + + // Schedule loss for any qubit whose sampled term was loss. The actual + // measure + reset is performed by the loss-commit ops emitted after the + // correlated-noise op. + shot.pending_loss_mask |= loss_mask; + + // Store the masks in the shot buffer for the execute stage + // We use the unitary entries to store these masks (reinterpreted as floats) + shot.unitary[0] = vec2f(bitcast(bit_flip_mask), bitcast(phase_flip_mask)); + + // For bit-flipped qubits, we need to swap the 0 and 1 probabilities and masks + // This is done in prepare_op, not execute_op, since it's a simple swap + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + let qubit_mask = 1u << q; + if ((bit_flip_mask & qubit_mask) != 0u) { + // Swap the probabilities + let temp = shot.qubit_state[q].zero_probability; + shot.qubit_state[q].zero_probability = shot.qubit_state[q].one_probability; + shot.qubit_state[q].one_probability = temp; + + // Swap the bits in qubit_is_0_mask and qubit_is_1_mask + let was_0 = (shot.qubit_is_0_mask & qubit_mask) != 0u; + let was_1 = (shot.qubit_is_1_mask & qubit_mask) != 0u; + if (was_0) { + shot.qubit_is_0_mask &= ~qubit_mask; + shot.qubit_is_1_mask |= qubit_mask; + } else if (was_1) { + shot.qubit_is_1_mask &= ~qubit_mask; + shot.qubit_is_0_mask |= qubit_mask; + } + } + } + + // Set up the shot state for the correlated noise execution + shot.op_type = OPID_CORRELATED_NOISE; + shot.op_idx = op_idx; + // No probabilities need to be recomputed in execute_op since we've already swapped them here + shot.qubits_updated_last_op_mask = 0u; +} + +// Performas a binary search on a correlated noise probability table +// +// Preconditions: +// - table is sorted ascending, with every entry higher than the prior +// - table entries are cumulative probabilities totaling <= 1.0 +// - 'start' is the offset into the buffer array where this table's entries begin +// - 'count' is the number of entries in this table +// - 'rand_lo' and 'rand_hi' form a Q1.63 format random number in [0.0, 1.0) to use for the search +// - This will only called if a result should be found, i.e., +// - count > 0 +// - rand < table[start + count - 1].probability +// +// Returns the index of the found entry relative to 'start', which is the smallest index where "rand < table[start + index].probability" +fn binary_search_noise_table(rand_lo: u32, rand_hi: u32, start: i32, count: i32) -> i32 { + var low: i32 = 0; + var high: i32 = count; + + while (low < high) { + let mid: i32 = low + (high - low) / 2; + let p_lo = batch_data.correlated_noise_entries[start + mid].probability_lo; + let p_hi = batch_data.correlated_noise_entries[start + mid].probability_hi; + + if (rand_hi < p_hi || (rand_hi == p_hi && rand_lo < p_lo)) { + high = mid; + } else { + low = mid + 1; + } + } + return low; +} + +fn get_measure_qubit(shot_idx: u32, op_idx: u32) -> u32 { + return ops[op_idx].q1; +} + +fn get_measure_result(shot_idx: u32, op_idx: u32) -> u32 { + return ops[op_idx].q2; +} + +// Get the qubit id at the given index from the correlated noise op's qubit args +// Qubit args are stored in the unitary matrix elements as f32 values +fn get_correlated_noise_qubit(op_idx: u32, index: u32) -> u32 { + // Qubit ids are stored in the unitary as f32 values, starting at unitary[0].x, unitary[0].y, etc. + let vec_idx = index / 2u; + let component = index % 2u; + if (component == 0u) { + return u32(ops[op_idx].unitary[vec_idx].x); + } else { + return u32(ops[op_idx].unitary[vec_idx].y); + } +} + +// Prepare the shot state for executing a correlated noise operation. +// Resolves qubit IDs from the op's unitary matrix, samples the noise table, builds masks, and applies. +fn prep_correlated_noise(shot_idx: u32, op_idx: u32) { + let op = &ops[op_idx]; + let noise_table_idx = op.q1; + let qubit_count = op.q2; + + let sample = sample_correlated_noise(shot_idx, op_idx, noise_table_idx); + if (sample.should_apply == 0u) { return; } + + // Build bit-flip, phase-flip, and loss masks using qubit IDs from the op's unitary matrix + var bit_flip_mask: u32 = 0u; + var phase_flip_mask: u32 = 0u; + var loss_mask: u32 = 0u; + for (var i: u32 = 0u; i < qubit_count; i++) { + let pauli_bits = get_pauli_bits(sample.paulis_lo, sample.paulis_hi, qubit_count, i); + let qubit_mask = 1u << get_correlated_noise_qubit(op_idx, i); + if ((pauli_bits & 0x4u) != 0u) { + // Loss term (L = 4): the qubit is lost, no Pauli is applied to it. + loss_mask |= qubit_mask; + } else { + if ((pauli_bits & 0x1u) != 0u) { bit_flip_mask |= qubit_mask; } + if ((pauli_bits & 0x2u) != 0u) { phase_flip_mask |= qubit_mask; } + } + } + + commit_correlated_noise(shot_idx, op_idx, bit_flip_mask, phase_flip_mask, loss_mask); +} + +//#region Adaptive QIR utility functions + +// ----------------------------------------------------------------------------- +// Adaptive interpreter — register file access +// ----------------------------------------------------------------------------- + +fn read_reg(shot_idx: u32, reg: u32) -> u32 { + return shots[shot_idx].interp.registers[reg]; +} + +fn write_reg(shot_idx: u32, reg: u32, val: u32) { + shots[shot_idx].interp.registers[reg] = val; +} + +fn read_reg_i32(shot_idx: u32, reg: u32) -> i32 { + return bitcast(read_reg(shot_idx, reg)); +} + +fn write_reg_i32(shot_idx: u32, reg: u32, val: i32) { + write_reg(shot_idx, reg, bitcast(val)); +} + +fn read_reg_f32(shot_idx: u32, reg: u32) -> f32 { + return bitcast(read_reg(shot_idx, reg)); +} + +fn write_reg_f32(shot_idx: u32, reg: u32, val: f32) { + write_reg(shot_idx, reg, bitcast(val)); +} + +// ----------------------------------------------------------------------------- +// Adaptive interpreter — instruction fetch and opcode extraction +// ----------------------------------------------------------------------------- + +fn fetch_instr(pc: u32) -> Instruction { + return batch_data.program.instructions[pc]; +} + +fn get_opcode(packed: u32) -> u32 { return packed & 0xFFu; } +fn get_subcond(packed: u32) -> u32 { return (packed >> 8u) & 0xFFu; } +fn get_flags(packed: u32) -> u32 { return (packed >> 16u) & 0xFFu; } +fn is_src0_imm(flags: u32) -> bool { return (flags & 1u) != 0u; } +fn is_src1_imm(flags: u32) -> bool { return (flags & 2u) != 0u; } + +fn resolve_i32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> i32 { + if (flags & (1u << operand_idx)) != 0u { + return bitcast(operand); // immediate + } + return read_reg_i32(shot_idx, operand); // register +} + +fn resolve_u32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> u32 { + if (flags & (1u << operand_idx)) != 0u { + return operand; + } + return read_reg(shot_idx, operand); +} + +fn resolve_f32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> f32 { + if (flags & (1u << operand_idx)) != 0u { + return bitcast(operand); // immediate (IEEE 754 bit pattern) + } + return read_reg_f32(shot_idx, operand); +} + +// Resolves q1 for the current quantum instruction. +fn resolve_q1(shot_idx: u32) -> u32 { + let state = shots[shot_idx].interp; + let instr = fetch_instr(state.pc - 1); + if (instr.opcode & FLAG_AUX1_IMM) != 0 { + return instr.aux1; + } + return read_reg(shot_idx, instr.aux1); +} + +// Resolves q2 for the current quantum instruction. +fn resolve_q2(shot_idx: u32) -> u32 { + let state = shots[shot_idx].interp; + let instr = fetch_instr(state.pc - 1); + if (instr.opcode & FLAG_AUX2_IMM) != 0 { + return instr.aux2; + } + return read_reg(shot_idx, instr.aux2); +} + +// Resolves the rotation angle for the current quantum instruction. +// The angle is stored in the instruction's src0 field (register or immediate). +fn resolve_gate_angle(shot_idx: u32) -> f32 { + let state = shots[shot_idx].interp; + let instr = fetch_instr(state.pc - 1); + let flags = get_flags(instr.opcode); + return resolve_f32(shot_idx, instr.src0, flags, 0u); +} + +fn get_measure_qubit_adaptive(shot_idx: u32, op_idx: u32) -> u32 { + return resolve_q1(shot_idx); +} + +fn get_measure_result_adaptive(shot_idx: u32, op_idx: u32) -> u32 { + return resolve_q2(shot_idx); +} + +// Read a measurement result from the existing results buffer. +// Results are stored as atomic at shot_idx * RESULT_COUNT + result_id. +fn read_measurement_result(shot_idx: u32, result_id: u32) -> bool { + return atomicLoad(&results[shot_idx * RESULT_COUNT + result_id]) == 1u; +} + +// Return true if the id corresponds to a rotation gate. +fn is_rotation_gate(id: u32) -> bool { + return (12 <= id && id <= 14) || (17 <= id && id <= 19); +} + +// Return true if the angle for the current rotation gate is dynamic. +fn is_dynamic_angle(shot_idx: u32) -> bool { + let state = shots[shot_idx].interp; + let instr = fetch_instr(state.pc - 1); + return (instr.opcode | FLAG_SRC0_IMM) != 0; +} + +// Commit a sampled qubit loss on an explicitly given qubit (measure + reset to +// |0> and mark the qubit lost). The lost qubit is carried to the execute stage +// in `op_idx`, and `op_type` is set to OPID_LOSS_NOISE so execute applies the +// reset matrix to that explicit qubit. +fn prep_loss_commit(shot_idx: u32, qubit: u32) { + let shot = &shots[shot_idx]; + let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); + shot.qubit_state[qubit].heat = -1.0; + prep_measure_reset_instrument(shot_idx, qubit, result, true /* resets_to_zero */); + shot.op_idx = qubit; // execute reads the lost qubit from op_idx + shot.op_type = OPID_LOSS_NOISE; +} + +// Prepare correlated noise for the adaptive path. +// Qubit IDs are read from call_arg_table (register indices), following the same +// pattern as OP_CALL argument passing. +fn prep_correlated_noise_adaptive(shot_idx: u32, op_idx: u32, qubit_count: u32, arg_offset: u32) { + let noise_table_idx = ops[op_idx].q1; + + let sample = sample_correlated_noise(shot_idx, op_idx, noise_table_idx); + if (sample.should_apply == 0u) { return; } + + // Build bit-flip, phase-flip, and loss masks using qubit IDs from registers via call_arg_table + var bit_flip_mask: u32 = 0u; + var phase_flip_mask: u32 = 0u; + var loss_mask: u32 = 0u; + for (var i: u32 = 0u; i < qubit_count; i++) { + let pauli_bits = get_pauli_bits(sample.paulis_lo, sample.paulis_hi, qubit_count, i); + let arg_reg = batch_data.program.call_arg_table[arg_offset + i]; + let qubit_mask = 1u << read_reg(shot_idx, arg_reg); + if ((pauli_bits & 0x4u) != 0u) { + // Loss term (L = 4): the qubit is lost, no Pauli is applied to it. + loss_mask |= qubit_mask; + } else { + if ((pauli_bits & 0x1u) != 0u) { bit_flip_mask |= qubit_mask; } + if ((pauli_bits & 0x2u) != 0u) { phase_flip_mask |= qubit_mask; } + } + } + + commit_correlated_noise(shot_idx, op_idx, bit_flip_mask, phase_flip_mask, loss_mask); +} + +//#endregion + +//#region Kernels + +// ******************************* +// PREPARE OP +// This stage prepares the shot state for the next operation to execute (and any updates needed from the prior op) +// +// Each op is prepared by one thread. This is how we deal with some of the challenges with synchronization +// when multiple workgroups with multiple threads are used for a shot in the EXECUTE stage. The 'execute_op' +// does work that is 'embarrassingly parallel' across the state vector amplitudes, but the PREPARE_OP stage +// deal with preparing for that work, and collating results back into the shot state afterwards. +// +// This allows us to use the GPU 'dispatch' mechanism to ensure consistencty across shots without complex, +// synchronization code, as the GPU guarantees that all threads in a dispatch complete before the next dispatch +// starts, and all buffer writes are visible to the next dispatch. +// ******************************* + +// NOTE: Run with workgroup size of 1 for now, as threads may diverge too much in prepare_op stage causing performance issues. +// TODO: Try to increase later if lack of parallelism is a bottleneck. (Update the dispatch call accordingly). +@compute @workgroup_size(1) +fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { + // For the 'prepare_op' stage, each thread dispatched handles one shot, so the globalId.x is the shot index + let shot_idx = globalId.x; + let shot = &shots[shot_idx]; + + // WebGPU guarantees that buffers are zero-initialized, so next_op_idx will correctly be 0 on the first dispatch + let op_idx = shot.next_op_idx; + + // If we've gone past the end, set the op type to id and exit, so the execute stage is a no-op + if (op_idx >= u32(arrayLength(&ops))) { + // TODO: Set error/diagnostic info here + shot.op_type = OPID_ID; + shot.renormalize = 1.0; + shot.qubits_updated_last_op_mask = 0u; + return; + } + + let op = &ops[op_idx]; + + // Update the shot state based on the results of the last executed op (if needed) + if (shot.qubits_updated_last_op_mask != 0) { + update_qubit_state(shot_idx); + } + + shot_init_per_op(shot_idx); + shot.unitary = op.unitary; + + // Handle MResetZ, MZ, and ResetZ operations. These have unique handling and no associated noise ops, so prep and exit + if (op.id == OPID_MRESETZ) { + prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, true /* resets_to_zero */); + shot.next_op_idx = op_idx + 1u; // No associated noise ops, so just advance by 1 + return; + } + if (op.id == OPID_MZ) { + prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, false /* resets_to_zero */); + shot.next_op_idx = op_idx + 1u; + return; + } + if (op.id == OPID_RESETZ) { + prep_measure_reset(shot_idx, op_idx, false /* is_loss */, false /* stores_result */, true /* resets_to_zero */); + shot.next_op_idx = op_idx + 1u; + return; + } + + // Loss-commit op: lose this qubit if and only if the preceding noise sampler + // set its bit in pending_loss_mask; otherwise act as identity. + if (op.id == OPID_LOSS_NOISE) { + shot.next_op_idx = op_idx + 1u; + let loss_bit = 1u << op.q1; + if ((shot.pending_loss_mask & loss_bit) != 0u) { + shot.pending_loss_mask &= ~loss_bit; + prep_measure_reset(shot_idx, op_idx, true /* is_loss */, false /* stores_result */, true /* resets_to_zero */); + } else { + shot.op_type = OPID_ID; + shot.op_idx = op_idx; + shot.qubits_updated_last_op_mask = 0u; + } + return; + } + + /* Handle noise: + - For the 1-qubit op case, there could be pauli and loss noise after the op itself. We want to check for loss first and + only apply pauli noise if the qubit wasn't lost. (If lost, the pauli noise and even the gate itself don't matter). + - For the 2-qubit op case, there will only be optional pauli noise after the op itself. (Loss is applied via separate + Id ops on each qubit after the 2-qubit op). + */ + + let pauli_op_idx = get_pauli_noise_idx(op_idx); + // Advance past this gate and its (optional) inline Pauli/loss noise op. Any + // loss-commit ops that follow are separate ops handled on later iterations. + shot.next_op_idx = max(op_idx, pauli_op_idx) + 1u; + + // Handle correlated noise operations + if (op.id == OPID_CORRELATED_NOISE) { + prep_correlated_noise(shot_idx, op_idx); + return; + } + + // Before doing further work, if any qubit for the gate is lost, dispatch + // the gate's configured loss policy (stamped on op.policy). + let has_lost_operand = gate_has_lost_operand(shot_idx, op_idx, op.q1, op.q2); + if (has_lost_operand) { + handle_lost_operand_policy(shot_idx, op_idx, op.q1, op.q2); + } + + if pauli_op_idx != 0 { + if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { + // A 1-qubit gate has a single operand; if it is lost there is no + // surviving qubit to receive Pauli noise, so skip the noise. + if (!has_lost_operand) { + apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, op.q1); + } + return; + } else { + if (has_lost_operand) { + // The gate body was handled by the loss policy above. Still apply + // the attached Pauli noise to the surviving operand (if any). + apply_2q_pauli_noise_on_survivor(shot_idx, op_idx, pauli_op_idx, op.q1, op.q2); + } else { + apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, op.q1, op.q2); + } + return; + } + } + + // If the gate has any lost operands (and no attached noise), the gate logic + // was completely handled inside `handle_lost_operand_policy`. + if (has_lost_operand) { + return; + } + + // No noise to apply, just set up the shot to execute the op as-is + shot.op_idx = op_idx; + shot.op_type = op.id; + + // Turn any Rxx, Ryy, or Rzz gates into a gate from the shot buffer + // NOTE: Should probably just do this for all gates + if (op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP) { + shot.op_type = OPID_SHOT_BUFF_2Q; // Indicate to use the matrix in the shot buffer + } + + if (op.id >= OPID_X && op.id < OPID_CX) { + shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer + } + + if (is_1q_phase_gate(op.id)) { + // For phase gates, treat everything as RZ for execution purposes + shot.op_type = OPID_RZ; + } + + // Set this so the next prepare_op stage knows which qubits to update probabilities for + switch shot.op_type { + case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { + shot.qubits_updated_last_op_mask = 0u; + } + case OPID_SHOT_BUFF_1Q { + shot.qubits_updated_last_op_mask = 1u << op.q1; + } + case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { + shot.qubits_updated_last_op_mask = (1u << op.q1) | (1u << op.q2); + } + default { + // TODO: Set error/diagnostic info here + } + } +} + +@compute @workgroup_size(THREADS_PER_WORKGROUP) +fn initialize( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) tid: u32) { + // Get the params + let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); + + // We want every thread to zero out its portion of the state vector for the shot + // We also want threads executing in lockstep to update adjacent entries for better memory access patterns + for (var i = 0; i < params.op_iterations; i++) { + let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; + stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); + } + + // NOTE: No need to synchronize here, as each thread is writing to unique locations + if (params.thread_idx_in_shot == 0) { + // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot + stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); + reset_all(params.shot_idx); + } +} + +@compute @workgroup_size(THREADS_PER_WORKGROUP) +fn initialize_adaptive( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) tid: u32) { + // Get the params + let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); + + // We want every thread to zero out its portion of the state vector for the shot + // We also want threads executing in lockstep to update adjacent entries for better memory access patterns + for (var i = 0; i < params.op_iterations; i++) { + let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; + stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); + } + + // NOTE: No need to synchronize here, as each thread is writing to unique locations + if (params.thread_idx_in_shot == 0) { + // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot + stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); + reset_all(params.shot_idx); + + // Zero the results buffer for this shot so stale exit codes from + // prior runs do not leak via atomicCompareExchangeWeak in OP_RET. + let results_base = u32(params.shot_idx) * RESULT_COUNT; + for (var r = 0u; r < RESULT_COUNT; r++) { + atomicStore(&results[results_base + r], 0u); + } + + // Initialize memory from constant_data + for (var m = 0u; m < CONSTANT_DATA_SIZE; m++) { + shots[params.shot_idx].interp.memory[m] = batch_data.program.constant_data[m]; + } + // Zero the alloca region for CPU-GPU parity + for (var m = CONSTANT_DATA_SIZE; m < MAX_MEMORY; m++) { + shots[params.shot_idx].interp.memory[m] = 0u; + } + } +} + +// ----------------------------------------------------------------------------- +// Adaptive interpreter — interpret_classical entry point +// ----------------------------------------------------------------------------- +// +// This is the main classical bytecode interpreter for the GPU-based adaptive +// quantum simulator. It implements a register-based virtual machine that +// executes classical (non-quantum) instructions on the GPU, one thread per +// shot. Each shot has its own independent interpreter state (program counter, +// registers, call stack) allowing many shots to run in parallel with +// potentially divergent control flow paths (e.g., after mid-circuit +// measurements). +// +// ## Execution Model +// +// The interpreter runs cooperatively with the quantum simulation pipeline: +// +// 1. The host dispatches `interpret_classical` for all shots. +// 2. Each shot executes classical instructions in a loop until one of: +// (a) A quantum operation is encountered → status = QUANTUM_PENDING, +// which tells the host to run the quantum simulation kernels +// (prepare_op → execute) before re-entering this function. +// (b) A `ret` instruction terminates the shot → status = TERMINATED. +// (c) The step limit (MAX_CLASSICAL_STEPS) is hit → status = YIELD, +// which prevents any single dispatch from running forever; the host +// simply re-dispatches to continue. +// (d) An unknown opcode is hit → status = ERROR. +// +// ## Instruction Encoding +// +// Each instruction occupies 2 × vec4 (8 u32 words) in the `bytecode` +// buffer, fetched by `fetch_instr(pc)` into the `Instr` struct with fields: +// +// opcode : packed opcode word (bits [7:0] = primary op, [15:8] = sub- +// condition for comparisons, [23:16] = flags for immediates) +// dst : destination register index (or immediate for RET) +// src0 : first source operand (register index or immediate) +// src1 : second source operand (register index or immediate) +// aux0–3 : auxiliary fields whose meaning varies per opcode (e.g., block +// IDs, function IDs, qubit indices, phi-table offsets, etc.) +// +// The `resolve_u32` / `resolve_i32` helpers read an operand as either a +// register value or an inline immediate based on the FLAG_SRC0_IMM / +// FLAG_SRC1_IMM bits in the flags byte. This lets the compiler embed small +// constants directly in the instruction stream without extra CONST ops. + + +@compute @workgroup_size(1) +fn interpret_classical(@builtin(global_invocation_id) gid: vec3) { + // Each GPU thread handles exactly one shot. The global invocation ID + // maps directly to the shot index. + let shot_idx = gid.x; + let state = shots[shot_idx].interp; + + // -- Early-exit for shots that already finished or errored -- + let status = state.status; + if status == STATUS_TERMINATED || status == STATUS_ERROR { + return; + } + + // -- Drain pending qubit losses before resuming classical execution -- + // The most recent noise op (per-gate Pauli/loss or correlated) may have + // sampled one or more qubits as lost, recorded in pending_loss_mask. Commit + // each as its own measure+reset quantum op (one per round) before running + // any more bytecode, so loss is applied with the correct correlation. + if shots[shot_idx].pending_loss_mask != 0u { + let q = firstTrailingBit(shots[shot_idx].pending_loss_mask); + shots[shot_idx].pending_loss_mask &= ~(1u << q); + shots[shot_idx].interp.pending_op_idx = q; + shots[shot_idx].interp.pending_op_type = PENDING_OP_LOSS_COMMIT; + shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; + return; + } + + // If we were paused (QUANTUM_PENDING after a quantum op, or YIELD after + // hitting the step limit), transition back to RUNNING so the main loop + // resumes executing instructions from where it left off. + if status != STATUS_RUNNING { + shots[shot_idx].interp.status = STATUS_RUNNING; + } + + // -- Load interpreter registers from GPU memory into local variables -- + // Using local vars for the hot-path state avoids repeated global memory + // loads/stores on every instruction. They are written back at the end. + var pc: u32 = state.pc; // program counter + var block_id: u32 = state.current_block_id; + var prev_block: u32 = state.previous_block_id; // for PHI + var steps: u32 = 0u; // counts instructions executed this dispatch + var should_break: bool = false; // set to true to exit the main loop + + // -- Main interpreter loop -- + // Fetches and executes one instruction per iteration. Exits when the + // shot terminates, yields for quantum work, hits the step limit, or + // encounters an error. + loop { + // Guard against infinite loops in classical code: after executing + // MAX_CLASSICAL_STEPS instructions, yield back to the host which + // will re-dispatch this kernel to continue. + if steps >= MAX_CLASSICAL_STEPS { + // Only yield if the shot hasn't already errored (an error + // status must not be overwritten by a yield). + if state.status != STATUS_ERROR { + shots[shot_idx].interp.status = STATUS_YIELD; + } + break; + } + + // Fetch the instruction at the current PC. Each instruction is + // 2 × vec4 (8 words) in the bytecode buffer. + let instr = fetch_instr(pc); + + // Unpack the opcode word into its three components: + // op — primary opcode (bits 7:0), determines which case below runs + // subcond — sub-condition code (bits 15:8), used only by ICMP/FCMP to + // select the specific comparison predicate (eq, ne, slt, etc.) + // flags — immediate-mode flags (bits 23:16), tells resolve_* whether + // src0/src1 are register indices or inline immediates + let op = get_opcode(instr.opcode); + let subcond = get_subcond(instr.opcode); + let flags = get_flags(instr.opcode); + + // -- Opcode dispatch -- + // The switch below implements every bytecode instruction. Instructions + // are grouped by category. Most follow a common pattern: + // 1. Read operands via resolve_u32/i32 (register or immediate) + // 2. Compute the result + // 3. Write back to the destination register via write_reg* + // 4. Advance pc++ + // + // Control-flow ops (JUMP, BRANCH, SWITCH, CALL) modify pc and + // block_id directly instead of incrementing pc. + // + // Quantum ops (QUANTUM_GATE, MEASURE, RESET) write pending-op + // metadata to the interpreter state and set should_break=true to + // pause execution and hand control back to the host for quantum + // kernel dispatch. + switch op { + + // ------------------------------------------------------------- + // CONTROL FLOW + // ------------------------------------------------------------- + + // NOP: No operation. Simply advances the program counter. + case OP_NOP { + pc++; + } + + // RET: Terminates this shot's execution. + // The exit code (from dst, which may be an immediate) is stored + // both in the per-shot interpreter state and atomically into the + // results buffer. The atomic-compare-exchange ensures only the + // first non-zero exit code is recorded for this shot (useful for + // error reporting). The termination count in the diagnostics + // buffer is incremented so the host can detect when all shots + // have finished. + case OP_RET { + let exit_code = resolve_u32(shot_idx, instr.dst, flags, 2u); + shots[shot_idx].interp.exit_code = exit_code; + // Atomically store exit code into the last slot of this shot's + // result region, but only if it has not already been set. + let err_index = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_index], 0u, exit_code); + shots[shot_idx].interp.status = STATUS_TERMINATED; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + } + + // JUMP: Unconditional branch to a target block. + // Encoding: dst = target block ID. + // Updates prev_block (needed by subsequent PHI instructions in + // the target block) and sets pc to the first instruction of the + // target block via the block_table lookup. + case OP_JUMP { + prev_block = block_id; + block_id = instr.dst; + pc = batch_data.program.block_table[instr.dst].instr_offset; + } + + // BRANCH: Conditional branch (if/else). + // Encoding: src0 = condition (register or immediate), + // aux0 = true-branch block ID, + // aux1 = false-branch block ID. + // Evaluates the condition: if non-zero, jumps to aux0; otherwise + // jumps to aux1. Like JUMP, updates prev_block for PHI nodes. + case OP_BRANCH { + let cond = resolve_u32(shot_idx, instr.src0, flags, 0u) != 0u; + prev_block = block_id; + if cond { + block_id = instr.aux0; + pc = batch_data.program.block_table[instr.aux0].instr_offset; + } else { + block_id = instr.aux1; + pc = batch_data.program.block_table[instr.aux1].instr_offset; + } + } + + // SWITCH: Multi-way branch (like a C switch statement). + // Encoding: src0 = value to match, + // aux0 = default block ID, + // aux1 = offset into switch_table, + // aux2 = number of case entries. + // Each switch_table entry is a vec2(match_value, target_block). + // Linearly scans the case table; if a match is found, jumps to + // that block. If no match, falls through to the default block. + case OP_SWITCH { + let val = resolve_u32(shot_idx, instr.src0, flags, 0u); + let default_block = instr.aux0; + let case_offset = instr.aux1; + let case_count = instr.aux2; + var target_block = default_block; + for (var i = 0u; i < case_count; i++) { + let entry = batch_data.program.switch_table[case_offset + i]; + if entry.case_val == val { + target_block = entry.target_block; + break; + } + } + prev_block = block_id; + block_id = target_block; + pc = batch_data.program.block_table[target_block].instr_offset; + } + + // CALL: Invokes a function. + // Encoding: dst = register to receive the return value, + // aux0 = function ID (index into function_table), + // aux1 = argument count, + // aux2 = offset into call_arg_table. + // + // The function_table entry is vec4(entry_block, param_count, + // param_base_reg, reserved). + // + // Steps: + // 1. Push a return frame onto the per-shot call stack. Each + // frame stores: (return_block, return_pc, return_reg, + // reserved) — 4 u32 words. The stack supports up to 8 frames. + // 2. Copy each argument from caller registers (looked up via + // call_arg_table) into callee parameter registers starting + // at param_base_reg. + // 3. Jump to the function's entry block. + case OP_CALL { + let func_id = instr.aux0; + let arg_count = instr.aux1; + let arg_offset = instr.aux2; + let func = batch_data.program.function_table[func_id]; + // Push return info onto the call stack + let sp = shots[shot_idx].interp.call_sp; + // Guard: prevent call stack overflow (max 8 frames) + if sp >= 8u { + shots[shot_idx].interp.exit_code = ERR_CALL_STACK_OVERFLOW; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_CALL_STACK_OVERFLOW); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + break; + } + shots[shot_idx].interp.call_stack_frames[sp].block_id = block_id; // return_block — resume here on return + shots[shot_idx].interp.call_stack_frames[sp].return_pc = pc + 1u; // return_pc — instruction after the CALL + shots[shot_idx].interp.call_stack_frames[sp].return_reg = instr.dst; // return_reg — where to write result + shots[shot_idx].interp.call_sp = sp + 1u; + // Copy caller arguments into the callee's parameter registers + let param_base = func.param_base_reg; + for (var i = 0u; i < arg_count; i++) { + let arg_reg = batch_data.program.call_arg_table[arg_offset + i]; + write_reg(shot_idx, param_base + i, read_reg(shot_idx, arg_reg)); + } + // Transfer control to the function entry block + block_id = func.entry_block_id; + pc = batch_data.program.block_table[block_id].instr_offset; + } + + // CALL_RETURN: Returns from a function call. + // Encoding: src0 = register holding the return value. + // + // Pops the top frame from the call stack to restore block_id and + // pc to the instruction after the CALL. If the caller specified a + // return register (not 0xFFFFFFFF), copies the return value into + // that register. + case OP_CALL_RETURN { + if shots[shot_idx].interp.call_sp == 0u { + shots[shot_idx].interp.exit_code = ERR_CALL_STACK_UNDERFLOW; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_CALL_STACK_UNDERFLOW); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + break; + } + + let sp = shots[shot_idx].interp.call_sp - 1; + shots[shot_idx].interp.call_sp = sp; + block_id = shots[shot_idx].interp.call_stack_frames[sp].block_id; + pc = shots[shot_idx].interp.call_stack_frames[sp].return_pc; + let return_reg = shots[shot_idx].interp.call_stack_frames[sp].return_reg; + if return_reg != VOID_RETURN { + write_reg(shot_idx, return_reg, read_reg(shot_idx, instr.src0)); + } + } + + // ------------------------------------------------------------- + // QUANTUM OPERATIONS — pause the interpreter, yield to the host + // ------------------------------------------------------------- + // When the interpreter hits a quantum instruction, it cannot + // execute it directly (quantum simulation runs in separate GPU + // kernels with parallel state-vector processing). Instead, it + // writes the pending operation details into the interpreter + // state for the host to read, sets status = QUANTUM_PENDING, + // advances pc past the instruction, and breaks out of the loop. + // + // The host then dispatches prepare_op (which reads the + // pending op metadata and configures the shot for the quantum + // kernel) followed by the execute kernel (which applies the + // gate/measurement/reset to the state vector). After that, the + // host re-dispatches interpret_classical to continue. + // + // Qubit IDs may be static (embedded in aux1/aux2 by the + // compiler) or dynamic (computed at runtime and stored in + // registers). + + // QUANTUM_GATE: Request a 1- or 2-qubit gate. + // Encoding: aux0 = quantum op table index, + // aux1 = qubit 1 (or register if not sentinel), + // aux2 = qubit 2 (or register if not sentinel). + case OP_QUANTUM_GATE { + shots[shot_idx].interp.pending_op_idx = instr.aux0; + shots[shot_idx].interp.pending_op_type = 0u; // type 0 = gate + // Qubit IDs are resolved in prepare_op via resolve_q1/resolve_q2, + // which use the FLAG_AUX1_IMM / FLAG_AUX2_IMM bits to decide + // between immediate values and register lookups. + shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; + pc++; + should_break = true; + } + + // MEASURE: Request a qubit measurement. + // Encoding: aux0 = quantum op table index, + // aux1 = qubit to measure (or register). + // Only q1 is used; q2 is set to sentinel (unused). + case OP_MEASURE { + shots[shot_idx].interp.pending_op_idx = instr.aux0; + shots[shot_idx].interp.pending_op_type = 1u; // type 1 = gate + // Qubit and result IDs are resolved in prepare_op via + // resolve_q1 (aux1) and resolve_q2 (aux2). + shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; + pc++; + should_break = true; + } + + // RESET: Request a qubit reset (measure + conditional X). + // Encoding: aux0 = quantum op table index, + // aux1 = qubit to reset (or register). + case OP_RESET { + shots[shot_idx].interp.pending_op_idx = instr.aux0; + shots[shot_idx].interp.pending_op_type = 2u; // type 2 = reset + // Qubit ID is resolved in prepare_op via resolve_q1 (aux1). + shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; + pc++; + should_break = true; + } + + // ------------------------------------------------------------- + // QUANTUM RESULT ACCESS + // ------------------------------------------------------------- + + // READ_RESULT: Load a prior measurement outcome into a register. + // Encoding: src0 = result ID (index into the results buffer), + // dst = destination register. + // The measurement result (0 or 1) was written by an earlier + // MEASURE quantum op. This reads it atomically from the shared + // results buffer and stores 0u or 1u into the destination + // register, allowing classical code to branch on measurement + // outcomes. + case OP_READ_RESULT { + let result_id = instr.src0; + let result_val = read_measurement_result(shot_idx, result_id); + write_reg(shot_idx, instr.dst, select(0u, 1u, result_val)); + pc++; + } + + // RECORD_OUTPUT: Marker for output recording. + // On the GPU this is a no-op — the host reads the results buffer + // directly after all shots terminate. The instruction exists to + // maintain compatibility with the QIR adaptive profile bytecode. + case OP_RECORD_OUTPUT { + pc++; + } + + // READ_LOSS: Reports whether the measurement that produced a + // result observed a lost qubit. The per-shot ``results`` buffer + // encodes loss as the value 2u (0u = Zero, 1u = One, 2u = Loss), + // so we compare against 2u and write 1u when the result was a loss, + // else 0u. + case OP_READ_LOSS { + let result_id = instr.src0; + let val = atomicLoad(&results[shot_idx * RESULT_COUNT + result_id]); + write_reg(shot_idx, instr.dst, select(0u, 1u, val == 2u)); + pc++; + } + + // ------------------------------------------------------------- + // INTEGER ARITHMETIC + // ------------------------------------------------------------- + // All integer arithmetic ops follow the pattern: + // dst = src0 src1 + // Operands are resolved via resolve_i32/u32, which checks the + // FLAG_SRC0_IMM / FLAG_SRC1_IMM bits to determine if the field + // is a register index or an inline immediate constant. + + // ADD: Signed integer addition. dst = src0 + src1. + case OP_ADD { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a + b); + pc++; + } + + // SUB: Signed integer subtraction. dst = src0 - src1. + case OP_SUB { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a - b); + pc++; + } + + // MUL: Signed integer multiplication. dst = src0 * src1. + case OP_MUL { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a * b); + pc++; + } + + // UDIV: Unsigned integer division. dst = src0 / src1. + case OP_UDIV { + let a = resolve_u32(shot_idx, instr.src0, flags, 0u); + let b = resolve_u32(shot_idx, instr.src1, flags, 1u); + write_reg(shot_idx, instr.dst, a / b); + pc++; + } + + // SDIV: Signed integer division (truncates toward zero). dst = src0 / src1. + case OP_SDIV { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a / b); + pc++; + } + + // UREM: Unsigned integer remainder. dst = src0 % src1. + case OP_UREM { + let a = resolve_u32(shot_idx, instr.src0, flags, 0u); + let b = resolve_u32(shot_idx, instr.src1, flags, 1u); + write_reg(shot_idx, instr.dst, a % b); + pc++; + } + + // SREM: Signed integer remainder. + // Computes a - b * trunc(a/b) manually rather than using the % + // operator, because WGSL i32 division truncates toward zero but + // the built-in % may not preserve the sign of the dividend on + // all GPU backends. This matches LLVM's srem semantics. + case OP_SREM { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a - b * (a / b)); + pc++; + } + + // ------------------------------------------------------------- + // BITWISE / SHIFT OPERATIONS + // ------------------------------------------------------------- + // Operate on the raw u32 bit pattern of the register values. + + // AND: Bitwise AND. dst = src0 & src1. + case OP_AND { + write_reg(shot_idx, instr.dst, + resolve_u32(shot_idx, instr.src0, flags, 0u) & resolve_u32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // OR: Bitwise OR. dst = src0 | src1. + case OP_OR { + write_reg(shot_idx, instr.dst, + resolve_u32(shot_idx, instr.src0, flags, 0u) | resolve_u32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // XOR: Bitwise exclusive OR. dst = src0 ^ src1. + case OP_XOR { + write_reg(shot_idx, instr.dst, + resolve_u32(shot_idx, instr.src0, flags, 0u) ^ resolve_u32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // SHL: Logical shift left. dst = src0 << src1. + case OP_SHL { + write_reg(shot_idx, instr.dst, + resolve_u32(shot_idx, instr.src0, flags, 0u) << resolve_u32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // LSHR: Logical shift right (zero-fill). dst = src0 >> src1. + case OP_LSHR { + write_reg(shot_idx, instr.dst, + resolve_u32(shot_idx, instr.src0, flags, 0u) >> resolve_u32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // ASHR: Arithmetic shift right (sign-extending). dst = src0 >> src1. + // Uses i32 to preserve the sign bit during the shift. + case OP_ASHR { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_u32(shot_idx, instr.src1, flags, 1u); + write_reg_i32(shot_idx, instr.dst, a >> b); + pc++; + } + + // ------------------------------------------------------------- + // INTEGER COMPARISON (ICMP) + // ------------------------------------------------------------- + // Compares two integer operands using the sub-condition code + // encoded in bits [15:8] of the opcode word. The result is + // written as 0u (false) or 1u (true) to the destination register. + // Signed comparisons (SLT, SLE, SGT, SGE) use i32 directly; + // unsigned comparisons (ULT, ULE, UGT, UGE) bitcast to u32. + // These mirror LLVM icmp predicates. + case OP_ICMP { + let a = resolve_i32(shot_idx, instr.src0, flags, 0u); + let b = resolve_i32(shot_idx, instr.src1, flags, 1u); + var result: bool = false; + switch subcond { + case ICMP_EQ { result = (a == b); } + case ICMP_NE { result = (a != b); } + case ICMP_SLT { result = (a < b); } + case ICMP_SLE { result = (a <= b); } + case ICMP_SGT { result = (a > b); } + case ICMP_SGE { result = (a >= b); } + case ICMP_ULT { result = (bitcast(a) < bitcast(b)); } + case ICMP_ULE { result = (bitcast(a) <= bitcast(b)); } + case ICMP_UGT { result = (bitcast(a) > bitcast(b)); } + case ICMP_UGE { result = (bitcast(a) >= bitcast(b)); } + default { + shots[shot_idx].interp.status = ERR_INVALID_INSTRUCTION; + shots[shot_idx].interp.exit_code = ERR_INVALID_INSTRUCTION; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_INVALID_INSTRUCTION); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + } + } + write_reg(shot_idx, instr.dst, select(0u, 1u, result)); + pc++; + } + + // ------------------------------------------------------------- + // FLOAT COMPARISON (FCMP) + // ------------------------------------------------------------- + // Compares two f32 operands using the sub-condition code. + // "O" prefix = ordered (both operands are not NaN). The result + // is written as 0u/1u. Mirrors LLVM fcmp ordered predicates. + case OP_FCMP { + let a = resolve_f32(shot_idx, instr.src0, flags, 0u); + let b = resolve_f32(shot_idx, instr.src1, flags, 1u); + var result: bool = false; + switch subcond { + case FCMP_OEQ { result = (a == b); } + case FCMP_ONE { result = (a != b); } + case FCMP_OLT { result = (a < b); } + case FCMP_OLE { result = (a <= b); } + case FCMP_OGT { result = (a > b); } + case FCMP_OGE { result = (a >= b); } + default { + shots[shot_idx].interp.exit_code = ERR_INVALID_INSTRUCTION; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_INVALID_INSTRUCTION); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + } + } + write_reg(shot_idx, instr.dst, select(0u, 1u, result)); + pc++; + } + + // ------------------------------------------------------------- + // FLOAT ARITHMETIC + // ------------------------------------------------------------- + // These operate on f32 values stored in registers via bitcast. + // Operands are always register-based (no immediate flags for + // float ops). + + // FADD: Float addition. dst = src0 + src1. + case OP_FADD { + write_reg_f32(shot_idx, instr.dst, + resolve_f32(shot_idx, instr.src0, flags, 0u) + resolve_f32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // FSUB: Float subtraction. dst = src0 - src1. + case OP_FSUB { + write_reg_f32(shot_idx, instr.dst, + resolve_f32(shot_idx, instr.src0, flags, 0u) - resolve_f32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // FMUL: Float multiplication. dst = src0 * src1. + case OP_FMUL { + write_reg_f32(shot_idx, instr.dst, + resolve_f32(shot_idx, instr.src0, flags, 0u) * resolve_f32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // FDIV: Float division. dst = src0 / src1. + case OP_FDIV { + write_reg_f32(shot_idx, instr.dst, + resolve_f32(shot_idx, instr.src0, flags, 0u) / resolve_f32(shot_idx, instr.src1, flags, 1u)); + pc++; + } + + // FREM: Float remainder. LLVM docs say this instruction has + // the same semantics as C's fmod, which is implemented as: + // dst = src0 - trunc(src0/src1) * src1 + case OP_FREM { + let a = resolve_f32(shot_idx, instr.src0, flags, 0u); + let b = resolve_f32(shot_idx, instr.src1, flags, 1u); + write_reg_f32(shot_idx, instr.dst, a - trunc(a / b) * b); + pc++; + } + + // ------------------------------------------------------------- + // TYPE CONVERSIONS + // ------------------------------------------------------------- + // Maps LLVM-style type conversion instructions. Many are + // identity ops on the GPU since all integer registers are 32-bit + // and all floats are f32. They exist to keep the bytecode in + // 1:1 correspondence with the compiled QIR instructions. + + // ZEXT: Zero-extend — identity on 32-bit GPU (values already u32). + case OP_ZEXT { + write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // SEXT: Sign-extend from a narrower bit width to i32. + // aux0 encodes the source bit width (e.g., 1 for i1→i32). + // The shift-left then arithmetic-shift-right trick propagates + // the sign bit from position (src_bits-1) into all higher bits. + case OP_SEXT { + let val = resolve_i32(shot_idx, instr.src0, flags, 0u); + let src_bits = instr.aux0; // source type bit width + if src_bits > 0u && src_bits < 32u { + let shift = 32u - src_bits; + write_reg_i32(shot_idx, instr.dst, (val << shift) >> shift); + } else { + write_reg_i32(shot_idx, instr.dst, val); + } + pc++; + } + + // TRUNC: Truncate — identity on 32-bit GPU (already the target width). + case OP_TRUNC { + write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // FPEXT: Float widen (e.g., f32→f64) — identity since GPU only uses f32. + case OP_FPEXT { + write_reg_f32(shot_idx, instr.dst, resolve_f32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // FPTRUNC: Float narrow (e.g., f64→f32) — identity since GPU only uses f32. + case OP_FPTRUNC { + write_reg_f32(shot_idx, instr.dst, resolve_f32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // INTTOPTR: Integer to pointer cast — identity, pointers are u32 on GPU. + case OP_INTTOPTR { + write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // FPTOSI: Float to signed integer conversion. dst = i32(src0). + case OP_FPTOSI { + write_reg_i32(shot_idx, instr.dst, i32(resolve_f32(shot_idx, instr.src0, flags, 0u))); + pc++; + } + + // SITOFP: Signed integer to float conversion. dst = f32(src0). + case OP_SITOFP { + write_reg_f32(shot_idx, instr.dst, f32(resolve_i32(shot_idx, instr.src0, flags, 0u))); + pc++; + } + + // FPTOUI: Float to unsigned integer conversion. dst = u32(src0). + case OP_FPTOUI { + write_reg(shot_idx, instr.dst, u32(resolve_f32(shot_idx, instr.src0, flags, 0u))); + pc++; + } + + // UITOFP: Unsigned integer to float conversion. dst = f32(src0). + case OP_UITOFP { + write_reg_f32(shot_idx, instr.dst, f32(resolve_u32(shot_idx, instr.src0, flags, 0u))); + pc++; + } + + // ------------------------------------------------------------- + // PHI NODE (SSA resolution at runtime) + // ------------------------------------------------------------- + // In SSA form, PHI nodes select a value based on which + // predecessor block the control flow came from. The compiler + // emits a phi_table with (predecessor_block_id, value_register) + // pairs for each PHI instruction. + // + // Encoding: dst = destination register, + // aux0 = offset into phi_table, + // aux1 = number of predecessor entries. + // + // At runtime, we scan the entries to find the one whose block + // ID matches prev_block, then copy that register's value into + // the destination. This is how the interpreter handles SSA + // control-flow merges without explicit move instructions on + // every edge. + case OP_PHI { + let offset = instr.aux0; + let count = instr.aux1; + for (var i = 0u; i < count; i++) { + let entry = batch_data.program.phi_table[offset + i]; + if entry.block_id == prev_block { + write_reg(shot_idx, instr.dst, read_reg(shot_idx, entry.val_reg)); + break; + } + } + pc++; + } + + // ------------------------------------------------------------- + // DATA MOVEMENT + // ------------------------------------------------------------- + + // SELECT: Conditional move (ternary operator). + // Encoding: src0 = condition, aux0 = true-value, + // aux1 = false-value, dst = destination. + // dst = cond ? aux0 : aux1 + case OP_SELECT { + let cond = resolve_u32(shot_idx, instr.src0, flags, 0u) != 0u; + let true_val = resolve_u32(shot_idx, instr.aux0, flags, 3u); + let false_val = resolve_u32(shot_idx, instr.aux1, flags, 4u); + write_reg(shot_idx, instr.dst, select(false_val, true_val, cond)); + pc++; + } + + // MOV: Register-to-register move (or immediate-to-register if flagged). + // dst = src0 (resolved through flags for possible immediate). + case OP_MOV { + write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); + pc++; + } + + // CONST: Load an immediate constant into a register. + // dst = src0 (always treated as a literal value, not a register). + case OP_CONST { + write_reg(shot_idx, instr.dst, instr.src0); + pc++; + } + + // ------------------------------------------------------------- + // MEMORY OPERATIONS + // ------------------------------------------------------------- + + // ALLOCA: Reserve memory and write the address to dst. + // Encoding: src0 = number of words, src1 = compile-time assigned address. + case OP_ALLOCA { + let num_words = resolve_u32(shot_idx, instr.src0, flags, 0u); + let addr = resolve_u32(shot_idx, instr.src1, flags, 1u); + if addr + num_words > MAX_MEMORY { + shots[shot_idx].interp.exit_code = ERR_ALLOCA_OUT_OF_BOUNDS; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_ALLOCA_OUT_OF_BOUNDS); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + break; + } + write_reg(shot_idx, instr.dst, addr); + pc++; + } + + // LOAD: Read a value from memory at the given address. + // Encoding: src0 = memory address, dst = destination register. + case OP_LOAD { + let addr = resolve_u32(shot_idx, instr.src0, flags, 0u); + if addr >= MAX_MEMORY { + shots[shot_idx].interp.exit_code = ERR_MEMORY_OUT_OF_BOUNDS; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_MEMORY_OUT_OF_BOUNDS); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + break; + } + let val = shots[shot_idx].interp.memory[addr]; + write_reg(shot_idx, instr.dst, val); + pc++; + } + + // STORE: Write a value to memory at the given address. + // Encoding: src0 = value to store, src1 = memory address. + case OP_STORE { + let val = resolve_u32(shot_idx, instr.src0, flags, 0u); + let addr = resolve_u32(shot_idx, instr.src1, flags, 1u); + if addr >= MAX_MEMORY { + shots[shot_idx].interp.exit_code = ERR_MEMORY_OUT_OF_BOUNDS; + let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; + atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_MEMORY_OUT_OF_BOUNDS); + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + break; + } + shots[shot_idx].interp.memory[addr] = val; + pc++; + } + + // GEP: Get element pointer — compute address from base + index * elem_size. + // Encoding: src0 = base address, src1 = index, aux0 = element size. + case OP_GEP { + let base = resolve_u32(shot_idx, instr.src0, flags, 0u); + let index = resolve_u32(shot_idx, instr.src1, flags, 1u); + let elem_size = resolve_u32(shot_idx, instr.aux0, flags, 3u); + let addr = base + index * elem_size; + write_reg(shot_idx, instr.dst, addr); + pc++; + } + + // Unknown opcode — flag the shot as errored. + default { + shots[shot_idx].interp.status = STATUS_ERROR; + atomicAdd(&diagnostics.termination_count, 1u); + should_break = true; + } + } + steps++; + if should_break { break; } + } + + // -- Persist interpreter state back to GPU memory -- + // Write the local variables back so the next dispatch (after quantum ops + // or a yield) can resume exactly where this invocation left off. + shots[shot_idx].interp.pc = pc; + shots[shot_idx].interp.current_block_id = block_id; + shots[shot_idx].interp.previous_block_id = prev_block; +} + +// ----------------------------------------------------------------------------- +// Adaptive interpreter — prepare_op entry point +// ----------------------------------------------------------------------------- +// Prepares a quantum operation for shots that have STATUS_QUANTUM_PENDING. +// Shots not in that state are set to OPID_ID so execute is a no-op. + +@compute @workgroup_size(1) +fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { + let shot_idx = globalId.x; + let shot = &shots[shot_idx]; + let state = shots[shot_idx].interp; + let status = state.status; + + // Only process shots that are quantum-pending + if status != STATUS_QUANTUM_PENDING { + // Set op_type to ID so execute is a no-op for this shot + shot.op_type = OPID_ID; + shot.renormalize = 1.0; + shot.qubits_updated_last_op_mask = 0u; + return; + } + + // Update shot state from prior op execution + if shot.qubits_updated_last_op_mask != 0 { + update_qubit_state(shot_idx); + } + shot_init_per_op(shot_idx); + + let op_idx = state.pending_op_idx; + let op_type = state.pending_op_type; + + // Loss commit: pending_op_idx holds the lost qubit (not an ops-pool index). + // Measure + reset that qubit; the execute stage applies it via op_idx. + if op_type == PENDING_OP_LOSS_COMMIT { + prep_loss_commit(shot_idx, op_idx); + return; + } + + let op = &ops[op_idx]; + + // Correlated noise: qubit IDs are stored as register indices in + // call_arg_table; read aux1 (qubit count) and aux2 (arg offset) + // from the instruction that triggered this quantum op. + if op_type == 0u && op.id == OPID_CORRELATED_NOISE { + let pc = state.pc; + let noise_instr = fetch_instr(pc - 1u); + let qubit_count = noise_instr.aux1; + let arg_offset = noise_instr.aux2; + shot.op_idx = op_idx; + shot.op_type = op.id; + prep_correlated_noise(shot_idx, op_idx, qubit_count, arg_offset); + shots[shot_idx].interp.status = STATUS_RUNNING; + return; + } + + let q1 = resolve_q1(shot_idx); + let q2 = resolve_q2(shot_idx); + + shot.unitary = op.unitary; + + switch op_type { + case 0u { // Gate + // For rotation gates, recompute the unitary from the dynamic angle stored + // in the instruction's src0 field if needed. The op pool unitary was built + // at upload time and may not reflect a runtime-computed angle. + if is_rotation_gate(op.id) && is_dynamic_angle(shot_idx) { + if op.id == OPID_RX || op.id == OPID_RY || op.id == OPID_RZ { + let angle = resolve_gate_angle(shot_idx); + let half = angle * 0.5; + let c = cos(half); + let s = sin(half); + if op.id == OPID_RX { + // [[cos(θ/2), -i·sin(θ/2)], [-i·sin(θ/2), cos(θ/2)]] + shot.unitary[0] = vec2f(c, 0.0); + shot.unitary[1] = vec2f(0.0, -s); + shot.unitary[4] = vec2f(0.0, -s); + shot.unitary[5] = vec2f(c, 0.0); + } else if op.id == OPID_RY { + // [[cos(θ/2), -sin(θ/2)], [sin(θ/2), cos(θ/2)]] + shot.unitary[0] = vec2f(c, 0.0); + shot.unitary[1] = vec2f(-s, 0.0); + shot.unitary[4] = vec2f(s, 0.0); + shot.unitary[5] = vec2f(c, 0.0); + } else { + // RZ: [[1, 0], [0, e^(iθ)]] + shot.unitary[0] = vec2f(1.0, 0.0); + shot.unitary[1] = vec2f(0.0, 0.0); + shot.unitary[4] = vec2f(0.0, 0.0); + shot.unitary[5] = vec2f(cos(angle), sin(angle)); + } + } else if op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_RZZ { + let angle = resolve_gate_angle(shot_idx); + let half = angle * 0.5; + let c = cos(half); + let s = sin(half); + if op.id == OPID_RXX { + // exp(-i·θ/2·X⊗X) + shot.unitary[0] = vec2f(c, 0.0); + shot.unitary[3] = vec2f(0.0, -s); + shot.unitary[5] = vec2f(c, 0.0); + shot.unitary[6] = vec2f(0.0, -s); + shot.unitary[9] = vec2f(0.0, -s); + shot.unitary[10] = vec2f(c, 0.0); + shot.unitary[12] = vec2f(0.0, -s); + shot.unitary[15] = vec2f(c, 0.0); + } else if op.id == OPID_RYY { + // exp(-i·θ/2·Y⊗Y) + shot.unitary[0] = vec2f(c, 0.0); + shot.unitary[3] = vec2f(0.0, s); + shot.unitary[5] = vec2f(c, 0.0); + shot.unitary[6] = vec2f(0.0, -s); + shot.unitary[9] = vec2f(0.0, -s); + shot.unitary[10] = vec2f(c, 0.0); + shot.unitary[12] = vec2f(0.0, s); + shot.unitary[15] = vec2f(c, 0.0); + } else { + // RZZ: diag(1, e^(iθ), e^(iθ), 1) + shot.unitary[0] = vec2f(1.0, 0.0); + shot.unitary[5] = vec2f(cos(angle), sin(angle)); + shot.unitary[10] = vec2f(cos(angle), sin(angle)); + shot.unitary[15] = vec2f(1.0, 0.0); + } + } + } + + shot.op_idx = op_idx; + shot.op_type = op.id; + + // If any operand is lost, dispatch the gate's configured loss + // policy (stamped on op.policy). + let has_lost_operand = gate_has_lost_operand(shot_idx, op_idx, q1, q2); + if (has_lost_operand) { + handle_lost_operand_policy(shot_idx, op_idx, q1, q2); + } + + // Check for noise ops after this gate in the ops pool + let pauli_op_idx = get_pauli_noise_idx(op_idx); + + // Handle Pauli noise (loss, if sampled, is recorded in pending_loss_mask) + if pauli_op_idx != 0u { + if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { + // A 1-qubit gate has a single operand; if it is lost there + // is no surviving qubit to receive Pauli noise. + if (!has_lost_operand) { + apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1); + } + } else { + if (has_lost_operand) { + // The gate body was handled by the loss policy above; + // apply the noise to the surviving operand (if any). + apply_2q_pauli_noise_on_survivor(shot_idx, op_idx, pauli_op_idx, q1, q2); + } else { + apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1, q2); + } + } + shots[shot_idx].interp.status = STATUS_RUNNING; + return; + } + + // If the gate has any lost operands (and no attached noise), the gate + // logic was completely handled inside `handle_lost_operand_policy`. + if (has_lost_operand) { + shots[shot_idx].interp.status = STATUS_RUNNING; + return; + } + + // No noise — set up the op for execution + + // Turn multi-qubit matrix ops into shot buffer ops + if op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP { + shot.op_type = OPID_SHOT_BUFF_2Q; + } + + // Turn 1Q matrix ops into shot buffer ops + if op.id >= OPID_X && op.id < OPID_CX { + shot.op_type = OPID_SHOT_BUFF_1Q; + } + + // Phase gates all execute as RZ + if is_1q_phase_gate(op.id) { + shot.op_type = OPID_RZ; + } + + // Set qubits_updated mask so next round knows which probabilities to update + switch shot.op_type { + case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { + shot.qubits_updated_last_op_mask = 0u; + } + case OPID_SHOT_BUFF_1Q { + shot.qubits_updated_last_op_mask = 1u << q1; + } + case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { + shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); + } + default {} + } + } + case 1u { // Measure + // Check for noise ops before the measure op + // (noise is applied as Id+noise, then original measure, matching non-adaptive pattern) + let pauli_op_idx = get_pauli_noise_idx(op_idx); + + if pauli_op_idx != 0u { + // Apply noise to the Id gate before measure, then the measure itself + // The non-adaptive path inserts Id+noise before measure; here the Id + // is at op_idx and the original measure op follows after noise ops + if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { + apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1); + } else { + apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1, q2); + } + shots[shot_idx].interp.status = STATUS_RUNNING; + return; + } + + // No noise — standard measure + let resets = op.id == OPID_MRESETZ; + prep_measure_reset(shot_idx, op_idx, false, true, resets); + } + case 2u { // Reset + prep_measure_reset(shot_idx, op_idx, false, false, true); + } + default { + shot.op_type = OPID_ID; + } + } + + // Mark shot as running so interpret_classical resumes next round + shots[shot_idx].interp.status = STATUS_RUNNING; +} + +@compute @workgroup_size(THREADS_PER_WORKGROUP) +fn execute( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) tid: u32) { + let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; + let shot = &shots[shot_idx]; + + // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating + // Correlated noise also updates probabilities in prepare_op, so can skip doing that here + let update_probs = shot.op_type != OPID_ID && shot.op_type != OPID_CORRELATED_NOISE && + shot.op_type != OPID_RZ && shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; + + if (shot.op_type == OPID_ID) { + // IGNORE + } else if (shot.op_type == OPID_CORRELATED_NOISE) { + apply_correlated_noise(workgroupId.x, tid); + } else if (is_1q_op(shot.op_type)) { + let q1: u32 = ops[shot.op_idx].q1; + apply_1q_op(workgroupId.x, tid, q1); + } else /* 2 qubit op */ { + let q1: u32 = ops[shot.op_idx].q1; + let q2: u32 = ops[shot.op_idx].q2; + apply_2q_op(workgroupId.x, tid, q1, q2); + } + + // workgroupBarrier can't be conditional in DX12 backend, so we have to do an unconditional one here + // outside of the skip_work conditional above. + workgroupBarrier(); + + // If the workgroup is done updating, have the first thread reduce the per-thread probabilities into the + // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. + // Skip for correlated noise since probabilities were already updated in prepare_op. + if (tid == 0 && update_probs) { + let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; + let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { + sum_thread_totals_to_shot(q, shot_idx, workgroup_collation_idx); + } + } + } +} + +// TODO: Seems to only differ in OPID_LOSS_NOISE case. See if can unify +@compute @workgroup_size(THREADS_PER_WORKGROUP) +fn execute_adaptive( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) tid: u32) { + let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; + let shot_idx_u32: u32 = u32(shot_idx); + let shot = &shots[shot_idx]; + + // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating + // Correlated noise also updates probabilities in prepare_op, so can skip doing that here + let update_probs = shot.op_type != OPID_ID && shot.op_type != OPID_CORRELATED_NOISE && + shot.op_type != OPID_RZ && shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; + + if (shot.op_type == OPID_ID) { + // IGNORE + } else if (shot.op_type == OPID_CORRELATED_NOISE) { + apply_correlated_noise(workgroupId.x, tid); + } else if (shot.op_type == OPID_LOSS_NOISE) { + // Loss commit: the lost qubit is carried in op_idx (set by prep_loss_commit). + apply_1q_op(workgroupId.x, tid, shot.op_idx); + } else if (is_1q_op(shot.op_type)) { + let q1: u32 = resolve_q1(shot_idx_u32); + apply_1q_op(workgroupId.x, tid, q1); + } else /* 2 qubit op */ { + let q1: u32 = resolve_q1(shot_idx_u32); + let q2: u32 = resolve_q2(shot_idx_u32); + apply_2q_op(workgroupId.x, tid, q1, q2); + } + + // workgroupBarrier can't be conditional in DX12 backend, so we have to do an unconditional one here + // outside of the skip_work conditional above. + workgroupBarrier(); + + // If the workgroup is done updating, have the first thread reduce the per-thread probabilities into the + // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. + // Skip for correlated noise since probabilities were already updated in prepare_op. + if (tid == 0 && update_probs) { + let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; + let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { + sum_thread_totals_to_shot(q, shot_idx, workgroup_collation_idx); + } + } + } +} + +//#endregion From f8032712219c3dd18ed61e595e127f29286f4d02 Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Thu, 16 Jul 2026 22:51:24 -0700 Subject: [PATCH 2/7] Unification fixes --- .../src/gpu_full_state_simulator/unified.wgsl | 39 ++++++------------- 1 file changed, 11 insertions(+), 28 deletions(-) diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/unified.wgsl index 9af9afcdb8c..b927f14b56f 100644 --- a/source/simulators/src/gpu_full_state_simulator/unified.wgsl +++ b/source/simulators/src/gpu_full_state_simulator/unified.wgsl @@ -860,10 +860,11 @@ fn prep_measure_reset_instrument(shot_idx: u32, qubit: u32, result: u32, resets_ & ~(shot.qubit_is_0_mask | shot.qubit_is_1_mask); } -fn prep_measure_reset(shot_idx: u32, op_idx: u32, is_loss: bool, stores_result: bool, resets_to_zero: bool) { +// `qubit` and `result_id` are resolved by the caller: the base pipeline reads +// them from the ops pool (ops[op_idx].q1/.q2), while the adaptive interpreter +// resolves them from registers/immediates (resolve_q1/resolve_q2). +fn prep_measure_reset(shot_idx: u32, op_idx: u32, qubit: u32, result_id: u32, is_loss: bool, stores_result: bool, resets_to_zero: bool) { let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - let qubit = get_measure_qubit(shot_idx, op_idx); // Choose measurement result based on qubit probabilities and random number let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); @@ -872,8 +873,6 @@ fn prep_measure_reset(shot_idx: u32, op_idx: u32, is_loss: bool, stores_result: // Instead, mark the qubit as lost by setting the heat to -1.0 if !is_loss { if stores_result { - let result_id = get_measure_result(shot_idx, op_idx); // Result id to store the measurement result in is stored in q2 - // If the qubit is already marked as lost, just report that and exit. It's already in the zero // state so nothing to update or renormalize. The execute op should be a no-op (ID) if shot.qubit_state[qubit].heat == -1.0 { @@ -1967,14 +1966,6 @@ fn binary_search_noise_table(rand_lo: u32, rand_hi: u32, start: i32, count: i32) return low; } -fn get_measure_qubit(shot_idx: u32, op_idx: u32) -> u32 { - return ops[op_idx].q1; -} - -fn get_measure_result(shot_idx: u32, op_idx: u32) -> u32 { - return ops[op_idx].q2; -} - // Get the qubit id at the given index from the correlated noise op's qubit args // Qubit args are stored in the unitary matrix elements as f32 values fn get_correlated_noise_qubit(op_idx: u32, index: u32) -> u32 { @@ -2111,14 +2102,6 @@ fn resolve_gate_angle(shot_idx: u32) -> f32 { return resolve_f32(shot_idx, instr.src0, flags, 0u); } -fn get_measure_qubit_adaptive(shot_idx: u32, op_idx: u32) -> u32 { - return resolve_q1(shot_idx); -} - -fn get_measure_result_adaptive(shot_idx: u32, op_idx: u32) -> u32 { - return resolve_q2(shot_idx); -} - // Read a measurement result from the existing results buffer. // Results are stored as atomic at shot_idx * RESULT_COUNT + result_id. fn read_measurement_result(shot_idx: u32, result_id: u32) -> bool { @@ -2229,17 +2212,17 @@ fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { // Handle MResetZ, MZ, and ResetZ operations. These have unique handling and no associated noise ops, so prep and exit if (op.id == OPID_MRESETZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, true /* resets_to_zero */); + prep_measure_reset(shot_idx, op_idx, op.q1, op.q2, false /* is_loss */, true /* stores_result */, true /* resets_to_zero */); shot.next_op_idx = op_idx + 1u; // No associated noise ops, so just advance by 1 return; } if (op.id == OPID_MZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, false /* resets_to_zero */); + prep_measure_reset(shot_idx, op_idx, op.q1, op.q2, false /* is_loss */, true /* stores_result */, false /* resets_to_zero */); shot.next_op_idx = op_idx + 1u; return; } if (op.id == OPID_RESETZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, false /* stores_result */, true /* resets_to_zero */); + prep_measure_reset(shot_idx, op_idx, op.q1, op.q2, false /* is_loss */, false /* stores_result */, true /* resets_to_zero */); shot.next_op_idx = op_idx + 1u; return; } @@ -2251,7 +2234,7 @@ fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { let loss_bit = 1u << op.q1; if ((shot.pending_loss_mask & loss_bit) != 0u) { shot.pending_loss_mask &= ~loss_bit; - prep_measure_reset(shot_idx, op_idx, true /* is_loss */, false /* stores_result */, true /* resets_to_zero */); + prep_measure_reset(shot_idx, op_idx, op.q1, op.q2, true /* is_loss */, false /* stores_result */, true /* resets_to_zero */); } else { shot.op_type = OPID_ID; shot.op_idx = op_idx; @@ -3312,7 +3295,7 @@ fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { let arg_offset = noise_instr.aux2; shot.op_idx = op_idx; shot.op_type = op.id; - prep_correlated_noise(shot_idx, op_idx, qubit_count, arg_offset); + prep_correlated_noise_adaptive(shot_idx, op_idx, qubit_count, arg_offset); shots[shot_idx].interp.status = STATUS_RUNNING; return; } @@ -3479,10 +3462,10 @@ fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { // No noise — standard measure let resets = op.id == OPID_MRESETZ; - prep_measure_reset(shot_idx, op_idx, false, true, resets); + prep_measure_reset(shot_idx, op_idx, q1, q2, false, true, resets); } case 2u { // Reset - prep_measure_reset(shot_idx, op_idx, false, false, true); + prep_measure_reset(shot_idx, op_idx, q1, q2, false, false, true); } default { shot.op_type = OPID_ID; From 0b0707faf3afa96e460a2599b592331a8fe1e69b Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Fri, 17 Jul 2026 00:22:15 -0700 Subject: [PATCH 3/7] Unified and tested --- source/simulators/src/bytecode/runtime.rs | 2 +- .../gpu_full_state_simulator/gpu_context.rs | 57 ++-- .../gpu_full_state_simulator/gpu_resources.rs | 227 +++++++-------- .../src/gpu_full_state_simulator/unified.wgsl | 258 +++++++----------- 4 files changed, 235 insertions(+), 309 deletions(-) diff --git a/source/simulators/src/bytecode/runtime.rs b/source/simulators/src/bytecode/runtime.rs index 125a23a193a..bb7486279cd 100644 --- a/source/simulators/src/bytecode/runtime.rs +++ b/source/simulators/src/bytecode/runtime.rs @@ -21,7 +21,7 @@ use crate::{ // --------------------------------------------------------------------------- // Opcode constants — must stay in sync with the Python `_adaptive_bytecode.py` -// and the WGSL `simulator_adaptive.wgsl` shader. +// and the WGSL `unified.wgsl` shader. // --------------------------------------------------------------------------- // Flags (pre-shifted to bit 16+) diff --git a/source/simulators/src/gpu_full_state_simulator/gpu_context.rs b/source/simulators/src/gpu_full_state_simulator/gpu_context.rs index c87500504f5..5d45e390750 100644 --- a/source/simulators/src/gpu_full_state_simulator/gpu_context.rs +++ b/source/simulators/src/gpu_full_state_simulator/gpu_context.rs @@ -6,7 +6,7 @@ use std::mem::size_of; use bytemuck::{Zeroable, cast_slice}; -use crate::bytecode::AdaptiveProgram; +use crate::bytecode::{AdaptiveProgram, Block, Function, Instruction, PhiNodeEntry, SwitchCase}; use crate::correlated_noise::NoiseTables; use crate::gpu_resources::GpuResources; use crate::noise_config::NoiseConfig; @@ -407,12 +407,18 @@ impl GpuContext { self.pipeline_is_dirty = true; } - // Upload combined noise batch_data buffer + // Upload combined noise batch_data buffer. + // The unified shader's BatchData always ends with the adaptive + // `program` struct, so even in base mode the buffer must reserve + // space for it (the bytes stay zeroed and unused). Its size is + // determined by the same clamped table sizes stamped into the shader. let noise_metadata_bytes: &[u8] = cast_slice(&self.noise_tables.metadata); let noise_entries_bytes: &[u8] = cast_slice(&self.noise_tables.entries); let noise_table_padded_size = self.run_params.noise_table_count * 16; let noise_entry_padded_size = self.run_params.noise_entry_count * 16; - let total_size = noise_table_padded_size + noise_entry_padded_size; + let total_size = noise_table_padded_size + + noise_entry_padded_size + + program_region_size(&self.run_params); let mut batch_buf = vec![0u8; total_size]; batch_buf[..noise_metadata_bytes.len()].copy_from_slice(noise_metadata_bytes); batch_buf[noise_table_padded_size..noise_table_padded_size + noise_entries_bytes.len()] @@ -425,20 +431,8 @@ impl GpuContext { if self.pipeline_is_dirty { // The pipeline is marked as dirty if the qubit or result count changed (shot count doesn't impact it) - self.resources.create_shaders( - params.qubit_count, - params.result_count, - // The next two params are derived from qubit count and result count, so will only change if those do - params.workgroups_per_shot, - params.entries_per_thread, - // The below are constants so will not change from run to run - THREADS_PER_WORKGROUP, - MAX_QUBIT_COUNT, - MAX_QUBITS_PER_WORKGROUP, - // These only change if the size of the noise table changes - params.noise_table_count, - params.noise_entry_count, - )?; + self.resources + .create_shaders(&self.run_params, self.is_adaptive)?; } self.resources.ensure_run_buffers( @@ -460,16 +454,16 @@ impl GpuContext { let state_vector_size_per_shot = entries_per_shot * 8; // Each entry is a complex containing 2 * f32 // Figure out the per-shot GPU struct size first, since it's needed for batching limits. - // In adaptive mode, each shot's GPU struct includes the interpreter state + registers + memory. + // The unified shader's ShotData always embeds the interpreter state + registers + memory, + // so both base and adaptive modes reserve that space (base clamps the counts to a minimum + // of 1, matching the shader constants; the region is unused in base mode). // WGSL requires struct size to be a multiple of the struct's alignment (8 bytes due to vec2f). - let gpu_shot_size = if self.is_adaptive { + let gpu_shot_size = { let raw = SIZEOF_SHOTDATA + size_of::() - + params.num_registers * 4 - + params.max_memory * 4; + + params.num_registers.max(1) * 4 + + params.max_memory.max(1) * 4; (raw + 7) & !7 // round up to 8-byte alignment - } else { - SIZEOF_SHOTDATA }; // Figure out some limits based on buffer size limits, structure sizes, and number of qubits @@ -681,8 +675,7 @@ impl GpuContext { } if self.pipeline_is_dirty { - let params = &self.run_params; - self.resources.create_shaders_adaptive(params)?; + self.resources.create_shaders(&self.run_params, true)?; } if self.program_is_dirty || self.noise_config_is_dirty || self.batch_data_is_dirty { @@ -1024,3 +1017,17 @@ fn i32_to_usize(value: i32) -> usize { fn u32_to_i32(value: u32) -> i32 { i32::try_from(value).unwrap_or_else(|_| panic!("{value} should fit in a i32")) } + +/// Byte size of the adaptive `Program` region at the end of the unified shader's +/// `BatchData` struct, using the same clamped (minimum 1) table sizes that are +/// stamped into the shader constants. In base mode the program data is unused, +/// but the binding must still reserve this space for the layout to match. +fn program_region_size(params: &RunParams) -> usize { + params.num_instructions.max(1) * size_of::>() + + params.num_blocks.max(1) * size_of::>() + + params.num_functions.max(1) * size_of::>() + + params.num_phi_entries.max(1) * size_of::>() + + params.num_switch_cases.max(1) * size_of::>() + + params.num_call_args.max(1) * size_of::() + + params.num_constant_data.max(1) * size_of::() +} diff --git a/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs b/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs index a8de06403da..f8d83fe75cf 100644 --- a/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs +++ b/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs @@ -391,18 +391,15 @@ impl GpuResources { Ok(()) } - #[allow(clippy::too_many_arguments)] - pub fn create_shaders( + /// Creates the compute pipelines from the unified shader source. + /// + /// The same `unified.wgsl` source powers both the base (linear op-list) and + /// adaptive (QIR bytecode interpreter) simulators; `is_adaptive` selects + /// which code paths are compiled in via the `IS_ADAPTIVE` constant. + pub(crate) fn create_shaders( &mut self, - qubit_count: i32, - result_count: i32, - workgroups_per_shot: i32, - entries_per_thread: i32, - threads_per_workgroup: i32, - max_qubit_count: i32, - max_qubits_per_workgroup: i32, - noise_table_count: usize, - noise_entry_count: usize, + params: &RunParams, + is_adaptive: bool, ) -> Result<(), String> { let adapter = self.adapter.as_ref().ok_or("GPU adapter not initialized")?; let device = self.device.as_ref().ok_or("GPU device not initialized")?; @@ -411,127 +408,51 @@ impl GpuResources { .as_ref() .ok_or("Bind group layout not initialized")?; // This is created with the device, so should exist here - // Create the shader module and bind group layout - let raw_shader_src = concat!( - include_str!("common.wgsl"), - include_str!("simulator_base.wgsl"), - ); - let mut shader_src = raw_shader_src - .replace("{{QUBIT_COUNT}}", &qubit_count.to_string()) - .replace("{{RESULT_COUNT}}", &(result_count + 1).to_string()) // +1 for result code per shot - .replace("{{WORKGROUPS_PER_SHOT}}", &workgroups_per_shot.to_string()) - .replace("{{ENTRIES_PER_THREAD}}", &entries_per_thread.to_string()) - .replace( - "{{THREADS_PER_WORKGROUP}}", - &threads_per_workgroup.to_string(), - ) - .replace("{{MAX_QUBIT_COUNT}}", &max_qubit_count.to_string()) - .replace( - "{{MAX_QUBITS_PER_WORKGROUP}}", - &max_qubits_per_workgroup.to_string(), - ) - .replace("{{NOISE_TABLE_COUNT}}", &noise_table_count.to_string()) - .replace("{{NOISE_ENTRY_COUNT}}", &noise_entry_count.to_string()); - - // Strip out DX12-incompatible code sections if needed - if adapter.get_info().backend == wgpu::Backend::Dx12 { - shader_src = strip_dx12_sections(&shader_src); - } - - let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("GPU Simulator Shader Module"), - source: wgpu::ShaderSource::Wgsl(shader_src.into()), - }); - - let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("GPU simulator pipeline layout"), - bind_group_layouts: &[Some(bind_group_layout)], - immediate_size: 0, - }); - - let get_kernel = |name: &str| -> ComputePipeline { - device.create_compute_pipeline(&ComputePipelineDescriptor { - label: Some(&format!("GPU kernel - {name}")), - layout: Some(&pipeline_layout), - module: &shader_module, - entry_point: Some(name), - compilation_options: Default::default(), - cache: None, - }) - }; - - let dummy_kernel = get_kernel("initialize"); - - self.device_resources.kernels = Some(GpuKernels { - init_op: get_kernel("initialize"), - prepare_op: get_kernel("prepare_op"), - execute_op: get_kernel("execute"), - interpret_classical: dummy_kernel, - }); - - Ok(()) - } - - #[allow(clippy::too_many_arguments)] - pub(crate) fn create_shaders_adaptive(&mut self, params: &RunParams) -> Result<(), String> { - let adapter = self.adapter.as_ref().ok_or("GPU adapter not initialized")?; - let device = self.device.as_ref().ok_or("GPU device not initialized")?; - let bind_group_layout = self - .bind_group_layout - .as_ref() - .ok_or("Bind group layout not initialized")?; // This is created with the device, so should exist here + // WGSL forbids zero-length fixed-size arrays, so the adaptive program + // table sizes are clamped to a minimum of 1. In base mode these tables + // are unused (the adaptive code paths are compiled out), so the minimal + // size just keeps the module valid. + let replacements = [ + ("QUBIT_COUNT", params.qubit_count.to_string()), + ("RESULT_COUNT", (params.result_count + 1).to_string()), // +1 for result code per shot + ( + "WORKGROUPS_PER_SHOT", + params.workgroups_per_shot.to_string(), + ), + ("ENTRIES_PER_THREAD", params.entries_per_thread.to_string()), + ("THREADS_PER_WORKGROUP", THREADS_PER_WORKGROUP.to_string()), + ("MAX_QUBIT_COUNT", MAX_QUBIT_COUNT.to_string()), + ( + "MAX_QUBITS_PER_WORKGROUP", + MAX_QUBITS_PER_WORKGROUP.to_string(), + ), + ("NOISE_TABLE_COUNT", params.noise_table_count.to_string()), + ("NOISE_ENTRY_COUNT", params.noise_entry_count.to_string()), + ("MAX_REGISTERS", params.num_registers.max(1).to_string()), + ("MAX_MEMORY", params.max_memory.max(1).to_string()), + ( + "INSTRUCTIONS_SIZE", + params.num_instructions.max(1).to_string(), + ), + ("BLOCK_TABLE_SIZE", params.num_blocks.max(1).to_string()), + ( + "FUNCTION_TABLE_SIZE", + params.num_functions.max(1).to_string(), + ), + ("PHI_TABLE_SIZE", params.num_phi_entries.max(1).to_string()), + ( + "SWITCH_CASES_SIZE", + params.num_switch_cases.max(1).to_string(), + ), + ("CALL_ARGS_SIZE", params.num_call_args.max(1).to_string()), + ( + "CONSTANT_DATA_SIZE", + params.num_constant_data.max(1).to_string(), + ), + ("IS_ADAPTIVE", is_adaptive.to_string()), + ]; - // Create the shader module and bind group layout - let raw_shader_src = concat!( - include_str!("common.wgsl"), - include_str!("simulator_adaptive.wgsl"), - ); - let mut shader_src = raw_shader_src - .replace("{{QUBIT_COUNT}}", ¶ms.qubit_count.to_string()) - .replace("{{RESULT_COUNT}}", &(params.result_count + 1).to_string()) // +1 for result code per shot - .replace( - "{{WORKGROUPS_PER_SHOT}}", - ¶ms.workgroups_per_shot.to_string(), - ) - .replace( - "{{ENTRIES_PER_THREAD}}", - ¶ms.entries_per_thread.to_string(), - ) - .replace( - "{{THREADS_PER_WORKGROUP}}", - &THREADS_PER_WORKGROUP.to_string(), - ) - .replace("{{MAX_QUBIT_COUNT}}", &MAX_QUBIT_COUNT.to_string()) - .replace( - "{{MAX_QUBITS_PER_WORKGROUP}}", - &MAX_QUBITS_PER_WORKGROUP.to_string(), - ) - .replace("{{MAX_REGISTERS}}", ¶ms.num_registers.to_string()) - .replace("{{MAX_MEMORY}}", ¶ms.max_memory.to_string()) - .replace( - "{{INSTRUCTIONS_SIZE}}", - ¶ms.num_instructions.to_string(), - ) - .replace("{{BLOCK_TABLE_SIZE}}", ¶ms.num_blocks.to_string()) - .replace("{{FUNCTION_TABLE_SIZE}}", ¶ms.num_functions.to_string()) - .replace("{{PHI_TABLE_SIZE}}", ¶ms.num_phi_entries.to_string()) - .replace( - "{{SWITCH_CASES_SIZE}}", - ¶ms.num_switch_cases.to_string(), - ) - .replace("{{CALL_ARGS_SIZE}}", ¶ms.num_call_args.to_string()) - .replace( - "{{CONSTANT_DATA_SIZE}}", - ¶ms.num_constant_data.to_string(), - ) - .replace( - "{{NOISE_TABLE_COUNT}}", - ¶ms.noise_table_count.to_string(), - ) - .replace( - "{{NOISE_ENTRY_COUNT}}", - ¶ms.noise_entry_count.to_string(), - ); + let mut shader_src = stamp_shader_consts(include_str!("unified.wgsl"), &replacements); // Strip out DX12-incompatible code sections if needed if adapter.get_info().backend == wgpu::Backend::Dx12 { @@ -927,6 +848,50 @@ impl GpuResources { } } +/// Stamps compile-time constant values into the shader source. +/// +/// `unified.wgsl` declares its host-substituted constants with a literal +/// default followed by a `// REPLACE` marker, e.g. +/// `const QUBIT_COUNT: i32 = 8; // REPLACE`. Keeping a valid literal default +/// (rather than a `{{PLACEHOLDER}}` token) lets editor tooling parse the file. +/// This walks the source line by line and, for every `// REPLACE` line, swaps +/// the literal default for the value provided in `replacements`, keyed by the +/// constant's name. Every `// REPLACE` constant must have a matching entry. +fn stamp_shader_consts(source: &str, replacements: &[(&str, String)]) -> String { + let mut result = String::with_capacity(source.len()); + + for line in source.lines() { + let trimmed = line.trim_start(); + if trimmed.starts_with("const ") && line.trim_end().ends_with("// REPLACE") { + // Extract the constant name: between "const " and ':'. + let after_const = &trimmed["const ".len()..]; + let name = after_const + .split(':') + .next() + .map(str::trim) + .expect("REPLACE const line must declare a name"); + let value = &replacements + .iter() + .find(|(n, _)| *n == name) + .unwrap_or_else(|| panic!("no replacement value provided for const `{name}`")) + .1; + + // Rebuild the line as: " String { let mut result = String::new(); diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/unified.wgsl index b927f14b56f..133cbf66923 100644 --- a/source/simulators/src/gpu_full_state_simulator/unified.wgsl +++ b/source/simulators/src/gpu_full_state_simulator/unified.wgsl @@ -27,6 +27,12 @@ const SWITCH_CASES_SIZE: u32 = 0; // REPLACE const CALL_ARGS_SIZE: u32 = 0; // REPLACE const CONSTANT_DATA_SIZE: u32 = 0; // REPLACE +// Selects the adaptive (QIR bytecode interpreter) code paths when true, or the +// base (linear op-list) paths when false. String-replaced by the host per run. +// Because it is a `const`, the compiler folds the branches and eliminates the +// unused path, so there is no runtime cost to the shared kernels below. +const IS_ADAPTIVE: bool = false; // REPLACE + //#endregion //#region Error codes @@ -1549,8 +1555,8 @@ fn apply_1q_op(workgroupId: u32, tid: u32, q1: u32) { let params = get_shot_params(workgroupId, tid, 1 /* qubits per op */); let shot = &shots[params.shot_idx]; let scale = shot.renormalize; - let lowMask = (1u << q1) - 1; - let highMask = (1u << u32(QUBIT_COUNT)) - 1 - lowMask; + let lowMask = (1 << q1) - 1; + let highMask = (1 << u32(QUBIT_COUNT)) - 1 - lowMask; let qubit_is_0_mask = i32(shots[params.shot_idx].qubit_is_0_mask); let qubit_is_1_mask = i32(shots[params.shot_idx].qubit_is_1_mask); @@ -2166,6 +2172,70 @@ fn prep_correlated_noise_adaptive(shot_idx: u32, op_idx: u32, qubit_count: u32, //#region Kernels +// Shared kernel helpers used by both the base and adaptive code paths. + +// Zero this shot's slice of the state vector and set the |0...0> amplitude to 1, +// then reset the shot's tracking state. Shared by the initialize kernel. +fn init_state_vector(params: ShotParams) { + // We want every thread to zero out its portion of the state vector for the shot + // We also want threads executing in lockstep to update adjacent entries for better memory access patterns + for (var i = 0; i < params.op_iterations; i++) { + let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; + stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); + } + + // NOTE: No need to synchronize here, as each thread is writing to unique locations + if (params.thread_idx_in_shot == 0) { + // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot + stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); + reset_all(params.shot_idx); + } +} + +// Finalize the setup of a plain (no-noise, no-loss) gate op for execution: +// translate the op id into the execute-stage op_type (shot-buffer conversions, +// phase gates as RZ) and record which qubit probabilities to update next round. +// `q1`/`q2` are the resolved operands (ops-pool values for base, register/ +// immediate resolved for adaptive). Shared by both prepare_op paths. +fn finalize_gate_op(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + + shot.op_idx = op_idx; + shot.op_type = op.id; + + // Turn any Rxx, Ryy, or Rzz gates into a gate from the shot buffer + // NOTE: Should probably just do this for all gates + if (op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP) { + shot.op_type = OPID_SHOT_BUFF_2Q; // Indicate to use the matrix in the shot buffer + } + + if (op.id >= OPID_X && op.id < OPID_CX) { + shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer + } + + if (is_1q_phase_gate(op.id)) { + // For phase gates, treat everything as RZ for execution purposes + shot.op_type = OPID_RZ; + } + + // Set this so the next prepare_op stage knows which qubits to update probabilities for + switch shot.op_type { + case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { + shot.qubits_updated_last_op_mask = 0u; + } + case OPID_SHOT_BUFF_1Q { + shot.qubits_updated_last_op_mask = 1u << q1; + } + case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { + shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); + } + default { + // TODO: Set error/diagnostic info here + } + } +} + // ******************************* // PREPARE OP // This stage prepares the shot state for the next operation to execute (and any updates needed from the prior op) @@ -2182,10 +2252,8 @@ fn prep_correlated_noise_adaptive(shot_idx: u32, op_idx: u32, qubit_count: u32, // NOTE: Run with workgroup size of 1 for now, as threads may diverge too much in prepare_op stage causing performance issues. // TODO: Try to increase later if lack of parallelism is a bottleneck. (Update the dispatch call accordingly). -@compute @workgroup_size(1) -fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { +fn prepare_op_base_impl(shot_idx: u32) { // For the 'prepare_op' stage, each thread dispatched handles one shot, so the globalId.x is the shot index - let shot_idx = globalId.x; let shot = &shots[shot_idx]; // WebGPU guarantees that buffers are zero-initialized, so next_op_idx will correctly be 0 on the first dispatch @@ -2295,39 +2363,7 @@ fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { } // No noise to apply, just set up the shot to execute the op as-is - shot.op_idx = op_idx; - shot.op_type = op.id; - - // Turn any Rxx, Ryy, or Rzz gates into a gate from the shot buffer - // NOTE: Should probably just do this for all gates - if (op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP) { - shot.op_type = OPID_SHOT_BUFF_2Q; // Indicate to use the matrix in the shot buffer - } - - if (op.id >= OPID_X && op.id < OPID_CX) { - shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer - } - - if (is_1q_phase_gate(op.id)) { - // For phase gates, treat everything as RZ for execution purposes - shot.op_type = OPID_RZ; - } - - // Set this so the next prepare_op stage knows which qubits to update probabilities for - switch shot.op_type { - case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { - shot.qubits_updated_last_op_mask = 0u; - } - case OPID_SHOT_BUFF_1Q { - shot.qubits_updated_last_op_mask = 1u << op.q1; - } - case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { - shot.qubits_updated_last_op_mask = (1u << op.q1) | (1u << op.q2); - } - default { - // TODO: Set error/diagnostic info here - } - } + finalize_gate_op(shot_idx, op_idx, op.q1, op.q2); } @compute @workgroup_size(THREADS_PER_WORKGROUP) @@ -2337,41 +2373,11 @@ fn initialize( // Get the params let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); - // We want every thread to zero out its portion of the state vector for the shot - // We also want threads executing in lockstep to update adjacent entries for better memory access patterns - for (var i = 0; i < params.op_iterations; i++) { - let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; - stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); - } - - // NOTE: No need to synchronize here, as each thread is writing to unique locations - if (params.thread_idx_in_shot == 0) { - // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot - stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); - reset_all(params.shot_idx); - } -} - -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn initialize_adaptive( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - // Get the params - let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); - - // We want every thread to zero out its portion of the state vector for the shot - // We also want threads executing in lockstep to update adjacent entries for better memory access patterns - for (var i = 0; i < params.op_iterations; i++) { - let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; - stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); - } - - // NOTE: No need to synchronize here, as each thread is writing to unique locations - if (params.thread_idx_in_shot == 0) { - // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot - stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); - reset_all(params.shot_idx); + // Zero the state vector and set the |0...0> amplitude to 1.0 for this shot. + init_state_vector(params); + // The adaptive interpreter needs additional per-shot state initialized. + if (IS_ADAPTIVE && params.thread_idx_in_shot == 0) { // Zero the results buffer for this shot so stale exit codes from // prior runs do not leak via atomicCompareExchangeWeak in OP_RET. let results_base = u32(params.shot_idx) * RESULT_COUNT; @@ -3251,9 +3257,7 @@ fn interpret_classical(@builtin(global_invocation_id) gid: vec3) { // Prepares a quantum operation for shots that have STATUS_QUANTUM_PENDING. // Shots not in that state are set to OPID_ID so execute is a no-op. -@compute @workgroup_size(1) -fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { - let shot_idx = globalId.x; +fn prepare_op_adaptive_impl(shot_idx: u32) { let shot = &shots[shot_idx]; let state = shots[shot_idx].interp; let status = state.status; @@ -3412,35 +3416,7 @@ fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { } // No noise — set up the op for execution - - // Turn multi-qubit matrix ops into shot buffer ops - if op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP { - shot.op_type = OPID_SHOT_BUFF_2Q; - } - - // Turn 1Q matrix ops into shot buffer ops - if op.id >= OPID_X && op.id < OPID_CX { - shot.op_type = OPID_SHOT_BUFF_1Q; - } - - // Phase gates all execute as RZ - if is_1q_phase_gate(op.id) { - shot.op_type = OPID_RZ; - } - - // Set qubits_updated mask so next round knows which probabilities to update - switch shot.op_type { - case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { - shot.qubits_updated_last_op_mask = 0u; - } - case OPID_SHOT_BUFF_1Q { - shot.qubits_updated_last_op_mask = 1u << q1; - } - case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { - shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); - } - default {} - } + finalize_gate_op(shot_idx, op_idx, q1, q2); } case 1u { // Measure // Check for noise ops before the measure op @@ -3476,56 +3452,23 @@ fn prepare_op_adaptive(@builtin(global_invocation_id) globalId: vec3) { shots[shot_idx].interp.status = STATUS_RUNNING; } -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn execute( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let shot = &shots[shot_idx]; - - // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating - // Correlated noise also updates probabilities in prepare_op, so can skip doing that here - let update_probs = shot.op_type != OPID_ID && shot.op_type != OPID_CORRELATED_NOISE && - shot.op_type != OPID_RZ && shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; - - if (shot.op_type == OPID_ID) { - // IGNORE - } else if (shot.op_type == OPID_CORRELATED_NOISE) { - apply_correlated_noise(workgroupId.x, tid); - } else if (is_1q_op(shot.op_type)) { - let q1: u32 = ops[shot.op_idx].q1; - apply_1q_op(workgroupId.x, tid, q1); - } else /* 2 qubit op */ { - let q1: u32 = ops[shot.op_idx].q1; - let q2: u32 = ops[shot.op_idx].q2; - apply_2q_op(workgroupId.x, tid, q1, q2); - } - - // workgroupBarrier can't be conditional in DX12 backend, so we have to do an unconditional one here - // outside of the skip_work conditional above. - workgroupBarrier(); - - // If the workgroup is done updating, have the first thread reduce the per-thread probabilities into the - // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. - // Skip for correlated noise since probabilities were already updated in prepare_op. - if (tid == 0 && update_probs) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { - sum_thread_totals_to_shot(q, shot_idx, workgroup_collation_idx); - } - } +// Single prepare_op entry point. Dispatches to the base or adaptive +// implementation based on the compile-time IS_ADAPTIVE flag; the unused +// implementation is eliminated by the compiler. +@compute @workgroup_size(1) +fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { + if (IS_ADAPTIVE) { + prepare_op_adaptive_impl(globalId.x); + } else { + prepare_op_base_impl(globalId.x); } } -// TODO: Seems to only differ in OPID_LOSS_NOISE case. See if can unify @compute @workgroup_size(THREADS_PER_WORKGROUP) -fn execute_adaptive( +fn execute( @builtin(workgroup_id) workgroupId: vec3, @builtin(local_invocation_index) tid: u32) { let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let shot_idx_u32: u32 = u32(shot_idx); let shot = &shots[shot_idx]; // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating @@ -3537,15 +3480,27 @@ fn execute_adaptive( // IGNORE } else if (shot.op_type == OPID_CORRELATED_NOISE) { apply_correlated_noise(workgroupId.x, tid); - } else if (shot.op_type == OPID_LOSS_NOISE) { + } else if (IS_ADAPTIVE && shot.op_type == OPID_LOSS_NOISE) { // Loss commit: the lost qubit is carried in op_idx (set by prep_loss_commit). apply_1q_op(workgroupId.x, tid, shot.op_idx); } else if (is_1q_op(shot.op_type)) { - let q1: u32 = resolve_q1(shot_idx_u32); + var q1: u32; + if (IS_ADAPTIVE) { + q1 = resolve_q1(u32(shot_idx)); + } else { + q1 = ops[shot.op_idx].q1; + } apply_1q_op(workgroupId.x, tid, q1); } else /* 2 qubit op */ { - let q1: u32 = resolve_q1(shot_idx_u32); - let q2: u32 = resolve_q2(shot_idx_u32); + var q1: u32; + var q2: u32; + if (IS_ADAPTIVE) { + q1 = resolve_q1(u32(shot_idx)); + q2 = resolve_q2(u32(shot_idx)); + } else { + q1 = ops[shot.op_idx].q1; + q2 = ops[shot.op_idx].q2; + } apply_2q_op(workgroupId.x, tid, q1, q2); } @@ -3557,7 +3512,6 @@ fn execute_adaptive( // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. // Skip for correlated noise since probabilities were already updated in prepare_op. if (tid == 0 && update_probs) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { From 69f07778130256e6abea63da795c435290998615 Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Fri, 17 Jul 2026 01:01:00 -0700 Subject: [PATCH 4/7] Region grouping --- .../src/gpu_full_state_simulator/unified.wgsl | 264 +++++++++++------- 1 file changed, 166 insertions(+), 98 deletions(-) diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/unified.wgsl index 133cbf66923..8d1820fd650 100644 --- a/source/simulators/src/gpu_full_state_simulator/unified.wgsl +++ b/source/simulators/src/gpu_full_state_simulator/unified.wgsl @@ -660,6 +660,8 @@ fn next_rand_f32(shot_idx: u32) -> f32 { //#endregion +//#region Operation classification helpers + fn is_1q_phase_gate(op_id: u32) -> bool { return (op_id == OPID_S || op_id == OPID_SAdj || op_id == OPID_T || op_id == OPID_TAdj || op_id == OPID_RZ); } @@ -670,6 +672,10 @@ fn is_1q_op(op_id: u32) -> bool { op_id == OPID_MAT1Q || op_id == OPID_SHOT_BUFF_1Q); } +//#endregion + +//#region Per-shot setup and reset + fn shot_init_per_op(shot_idx: u32) { let shot = &shots[shot_idx]; @@ -736,6 +742,10 @@ fn reset_all(shot_idx: i32) { // unitary will be set in prepare_op } +//#endregion + +//#region Qubit probability tracking + fn update_qubit_state(shot_idx: u32) { let shot = &shots[shot_idx]; @@ -819,6 +829,68 @@ fn update_qubit_state(shot_idx: u32) { } } +// For the state vector index and amplitude probability, update all the qubit probabilities for this thread +fn update_all_qubit_probs(stateVectorIndex: u32, amplitude: vec2f, tid: u32) { + var mask: u32 = 1u; + for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { + let is_one: bool = (stateVectorIndex & mask) != 0u; + let prob: f32 = cplxMag2(amplitude); + if (is_one) { + qubitProbabilities[tid].one[q] += prob; + } else { + qubitProbabilities[tid].zero[q] += prob; + } + mask = mask << 1u; + } +} + +fn sum_thread_totals_to_shot(q: u32, shot_idx: i32, wkg_collation_idx: i32) { + var total_zero: f32 = 0.0; + var total_one: f32 = 0.0; + for (var j = 0; j < THREADS_PER_WORKGROUP; j++) { + total_zero += qubitProbabilities[j].zero[q]; + total_one += qubitProbabilities[j].one[q]; + } + if (wkg_collation_idx >= 0) { + // Write to the workgroup collation buffer for later summation into the shot state + workgroup_collation.sums[wkg_collation_idx].qubits[q] = vec2f(total_zero, total_one); + } else { + // Single workgroup per shot case - write directly to the shot state + let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; + if !within_threshold { + // Populate the diagnostics buffer, if not already set + let old_value = atomicCompareExchangeWeak( + &diagnostics.error_code, + 0u, + ERR_INVALID_THREAD_TOTAL); + if old_value.exchanged { + // This is the first error - fill in the details + let shot = &shots[shot_idx]; + diagnostics.extra1 = q; + diagnostics.extra2 = total_zero; + diagnostics.extra3 = total_one; + // DX12 backend has issues copying structs. See https://github.com/gfx-rs/wgpu/issues/8552 + // DX12-start-strip + diagnostics.shot = *shot; + diagnostics.op = ops[shot.op_idx]; + // DX12-end-strip + } + let err_index = (shot_idx + 1) * i32(RESULT_COUNT) - 1; + atomicCompareExchangeWeak( + &results[err_index], + 0u, + ERR_INVALID_THREAD_TOTAL); + } else { + shots[shot_idx].qubit_state[q].zero_probability = total_zero; + shots[shot_idx].qubit_state[q].one_probability = total_one; + } + } +} + +//#endregion + +//#region Measurement and reset ops + // Build a measure-and-reset (or measure-only) instrument for `qubit` given a // measured `result`, store it in the shot buffer, set up renormalization, and // mark the qubit as no longer in a definite basis state so the execute stage @@ -912,39 +984,9 @@ fn prep_measure_reset(shot_idx: u32, op_idx: u32, qubit: u32, result_id: u32, is shot.op_type = OPID_MRESETZ; } -// Starting from the given index, return the next index if pauli noise, else 0 -fn get_pauli_noise_idx(op_idx: u32) -> u32 { - if (arrayLength(&ops) > (op_idx + 1)) { - let op = &ops[op_idx + 1]; - if (op.id == OPID_PAULI_NOISE_1Q || op.id == OPID_PAULI_NOISE_2Q) { - return op_idx + 1u; - } - } - return 0u; -} - -// From the starting index given, return the next index if loss noise, else 0 -fn get_loss_idx(op_idx: u32) -> u32 { - if (arrayLength(&ops) > (op_idx + 1)) { - let op = &ops[op_idx + 1]; - if (op.id == OPID_LOSS_NOISE) { - return op_idx + 1u; - } - } - return 0u; -} +//#endregion -// Returns true if the gate at `op_idx` touches at least one lost qubit. -// `q1`/`q2` are the (resolved) operands of the gate. -fn gate_has_lost_operand(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) -> bool { - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - if (shot.qubit_state[q1].heat == -1.0) { - return true; - } - let is_2q = !is_1q_op(op.id); - return is_2q && (shot.qubit_state[q2].heat == -1.0); -} +//#region Unitary construction helpers // Builds a 4x4 (in shot.unitary) that applies the 1-qubit matrix `m` (given as // m00,m01,m10,m11) to `target_is_q2 ? q2 : q1` and identity to the other qubit @@ -995,6 +1037,33 @@ fn finish_2q_shot_buffer(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); } +//#endregion + +//#region Qubit loss handling + +// From the starting index given, return the next index if loss noise, else 0 +fn get_loss_idx(op_idx: u32) -> u32 { + if (arrayLength(&ops) > (op_idx + 1)) { + let op = &ops[op_idx + 1]; + if (op.id == OPID_LOSS_NOISE) { + return op_idx + 1u; + } + } + return 0u; +} + +// Returns true if the gate at `op_idx` touches at least one lost qubit. +// `q1`/`q2` are the (resolved) operands of the gate. +fn gate_has_lost_operand(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) -> bool { + let shot = &shots[shot_idx]; + let op = &ops[op_idx]; + if (shot.qubit_state[q1].heat == -1.0) { + return true; + } + let is_2q = !is_1q_op(op.id); + return is_2q && (shot.qubit_state[q2].heat == -1.0); +} + // Loses a single surviving `qubit` for the PROPAGATE policy: samples a // measurement outcome, collapses the qubit to that outcome and resets it to // |0>, and marks it lost (heat = -1.0). The collapse is expressed as a 2-qubit @@ -1032,16 +1101,6 @@ fn propagate_loss_to_qubit(shot_idx: u32, op_idx: u32, q1: u32, q2: u32, qubit: finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); } -// Records an error `code` for `shot_idx` in both the diagnostics buffer and the -// shot's result-code slot, mirroring the reporting done elsewhere in this file. -// Used for conditions the host guarantees never occur (e.g. a loss policy that -// is not valid for a given gate). -fn report_shot_error(shot_idx: u32, code: u32) { - atomicCompareExchangeWeak(&diagnostics.error_code, 0u, code); - let err_index = (shot_idx + 1u) * RESULT_COUNT - 1u; - atomicCompareExchangeWeak(&results[err_index], 0u, code); -} - // Handles a gate whose operand(s) include at least one lost qubit, according to // the loss policy stamped on the op's `policy` field. `q1`/`q2` are the // (resolved) operands. The gate body is fully handled here (degraded unitary, @@ -1197,6 +1256,35 @@ fn handle_lost_operand_policy(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { shot.op_idx = op_idx; } +//#endregion + +//#region Error reporting + +// Records an error `code` for `shot_idx` in both the diagnostics buffer and the +// shot's result-code slot, mirroring the reporting done elsewhere in this file. +// Used for conditions the host guarantees never occur (e.g. a loss policy that +// is not valid for a given gate). +fn report_shot_error(shot_idx: u32, code: u32) { + atomicCompareExchangeWeak(&diagnostics.error_code, 0u, code); + let err_index = (shot_idx + 1u) * RESULT_COUNT - 1u; + atomicCompareExchangeWeak(&results[err_index], 0u, code); +} + +//#endregion + +//#region Independent Pauli noise + +// Starting from the given index, return the next index if pauli noise, else 0 +fn get_pauli_noise_idx(op_idx: u32) -> u32 { + if (arrayLength(&ops) > (op_idx + 1)) { + let op = &ops[op_idx + 1]; + if (op.id == OPID_PAULI_NOISE_1Q || op.id == OPID_PAULI_NOISE_2Q) { + return op_idx + 1u; + } + } + return 0u; +} + fn apply_1q_pauli_noise(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32) { // NOTE: Assumes that whatever prepared the program ensured that noise_op.q1 matches op.q1 and // that op is a 1-qubit gate. `q1` is the resolved target qubit (may be @@ -1520,6 +1608,10 @@ fn apply_2q_pauli_noise_on_survivor(shot_idx: u32, op_idx: u32, noise_idx: u32, shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << survivor); } +//#endregion + +//#region Shot and kernel params + fn get_shot_params( workgroupId: u32, tid: u32, @@ -1551,6 +1643,10 @@ fn get_shot_params( ); } +//#endregion + +//#region Gate application (execute stage) + fn apply_1q_op(workgroupId: u32, tid: u32, q1: u32) { let params = get_shot_params(workgroupId, tid, 1 /* qubits per op */); let shot = &shots[params.shot_idx]; @@ -1783,63 +1879,9 @@ fn apply_correlated_noise(workgroupId: u32, tid: u32) { } } -// For the state vector index and amplitude probability, update all the qubit probabilities for this thread -fn update_all_qubit_probs(stateVectorIndex: u32, amplitude: vec2f, tid: u32) { - var mask: u32 = 1u; - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - let is_one: bool = (stateVectorIndex & mask) != 0u; - let prob: f32 = cplxMag2(amplitude); - if (is_one) { - qubitProbabilities[tid].one[q] += prob; - } else { - qubitProbabilities[tid].zero[q] += prob; - } - mask = mask << 1u; - } -} +//#endregion -fn sum_thread_totals_to_shot(q: u32, shot_idx: i32, wkg_collation_idx: i32) { - var total_zero: f32 = 0.0; - var total_one: f32 = 0.0; - for (var j = 0; j < THREADS_PER_WORKGROUP; j++) { - total_zero += qubitProbabilities[j].zero[q]; - total_one += qubitProbabilities[j].one[q]; - } - if (wkg_collation_idx >= 0) { - // Write to the workgroup collation buffer for later summation into the shot state - workgroup_collation.sums[wkg_collation_idx].qubits[q] = vec2f(total_zero, total_one); - } else { - // Single workgroup per shot case - write directly to the shot state - let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; - if !within_threshold { - // Populate the diagnostics buffer, if not already set - let old_value = atomicCompareExchangeWeak( - &diagnostics.error_code, - 0u, - ERR_INVALID_THREAD_TOTAL); - if old_value.exchanged { - // This is the first error - fill in the details - let shot = &shots[shot_idx]; - diagnostics.extra1 = q; - diagnostics.extra2 = total_zero; - diagnostics.extra3 = total_one; - // DX12 backend has issues copying structs. See https://github.com/gfx-rs/wgpu/issues/8552 - // DX12-start-strip - diagnostics.shot = *shot; - diagnostics.op = ops[shot.op_idx]; - // DX12-end-strip - } - let err_index = (shot_idx + 1) * i32(RESULT_COUNT) - 1; - atomicCompareExchangeWeak( - &results[err_index], - 0u, - ERR_INVALID_THREAD_TOTAL); - } else { - shots[shot_idx].qubit_state[q].zero_probability = total_zero; - shots[shot_idx].qubit_state[q].one_probability = total_one; - } - } -} +//#region Correlated noise // Samples the correlated noise table to determine whether noise should be applied, and if so, // which Pauli string was selected. If no noise is applied, the shot is set to ID and the caller @@ -2014,6 +2056,8 @@ fn prep_correlated_noise(shot_idx: u32, op_idx: u32) { commit_correlated_noise(shot_idx, op_idx, bit_flip_mask, phase_flip_mask, loss_mask); } +//#endregion + //#region Adaptive QIR utility functions // ----------------------------------------------------------------------------- @@ -2172,6 +2216,8 @@ fn prep_correlated_noise_adaptive(shot_idx: u32, op_idx: u32, qubit_count: u32, //#region Kernels +//#region Shared kernel helpers + // Shared kernel helpers used by both the base and adaptive code paths. // Zero this shot's slice of the state vector and set the |0...0> amplitude to 1, @@ -2236,6 +2282,10 @@ fn finalize_gate_op(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { } } +//#endregion + +//#region Base prepare_op implementation + // ******************************* // PREPARE OP // This stage prepares the shot state for the next operation to execute (and any updates needed from the prior op) @@ -2366,6 +2416,10 @@ fn prepare_op_base_impl(shot_idx: u32) { finalize_gate_op(shot_idx, op_idx, op.q1, op.q2); } +//#endregion + +//#region initialize kernel + @compute @workgroup_size(THREADS_PER_WORKGROUP) fn initialize( @builtin(workgroup_id) workgroupId: vec3, @@ -2396,6 +2450,10 @@ fn initialize( } } +//#endregion + +//#region interpret_classical kernel + // ----------------------------------------------------------------------------- // Adaptive interpreter — interpret_classical entry point // ----------------------------------------------------------------------------- @@ -3251,6 +3309,10 @@ fn interpret_classical(@builtin(global_invocation_id) gid: vec3) { shots[shot_idx].interp.previous_block_id = prev_block; } +//#endregion + +//#region Adaptive prepare_op implementation + // ----------------------------------------------------------------------------- // Adaptive interpreter — prepare_op entry point // ----------------------------------------------------------------------------- @@ -3452,6 +3514,10 @@ fn prepare_op_adaptive_impl(shot_idx: u32) { shots[shot_idx].interp.status = STATUS_RUNNING; } +//#endregion + +//#region prepare_op and execute kernels + // Single prepare_op entry point. Dispatches to the base or adaptive // implementation based on the compile-time IS_ADAPTIVE flag; the unused // implementation is eliminated by the compiler. @@ -3522,3 +3588,5 @@ fn execute( } //#endregion + +//#endregion From 9f2ff4051d60733f6ac63e7ebc31ad2206da926e Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Fri, 17 Jul 2026 01:04:40 -0700 Subject: [PATCH 5/7] Remove dead code --- .../src/gpu_full_state_simulator/unified.wgsl | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/unified.wgsl index 8d1820fd650..32f94bb4fb5 100644 --- a/source/simulators/src/gpu_full_state_simulator/unified.wgsl +++ b/source/simulators/src/gpu_full_state_simulator/unified.wgsl @@ -1041,17 +1041,6 @@ fn finish_2q_shot_buffer(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { //#region Qubit loss handling -// From the starting index given, return the next index if loss noise, else 0 -fn get_loss_idx(op_idx: u32) -> u32 { - if (arrayLength(&ops) > (op_idx + 1)) { - let op = &ops[op_idx + 1]; - if (op.id == OPID_LOSS_NOISE) { - return op_idx + 1u; - } - } - return 0u; -} - // Returns true if the gate at `op_idx` touches at least one lost qubit. // `q1`/`q2` are the (resolved) operands of the gate. fn gate_has_lost_operand(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) -> bool { @@ -2099,8 +2088,6 @@ fn fetch_instr(pc: u32) -> Instruction { fn get_opcode(packed: u32) -> u32 { return packed & 0xFFu; } fn get_subcond(packed: u32) -> u32 { return (packed >> 8u) & 0xFFu; } fn get_flags(packed: u32) -> u32 { return (packed >> 16u) & 0xFFu; } -fn is_src0_imm(flags: u32) -> bool { return (flags & 1u) != 0u; } -fn is_src1_imm(flags: u32) -> bool { return (flags & 2u) != 0u; } fn resolve_i32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> i32 { if (flags & (1u << operand_idx)) != 0u { From b62c3c76818621cbeae120b2d456eb0af0ac5b7f Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Fri, 17 Jul 2026 13:44:40 -0700 Subject: [PATCH 6/7] Remove old wgsl files --- .../src/gpu_full_state_simulator/common.wgsl | 1520 -------------- .../simulator_adaptive.wgsl | 1772 ----------------- .../simulator_base.wgsl | 449 ----- 3 files changed, 3741 deletions(-) delete mode 100644 source/simulators/src/gpu_full_state_simulator/common.wgsl delete mode 100644 source/simulators/src/gpu_full_state_simulator/simulator_adaptive.wgsl delete mode 100644 source/simulators/src/gpu_full_state_simulator/simulator_base.wgsl diff --git a/source/simulators/src/gpu_full_state_simulator/common.wgsl b/source/simulators/src/gpu_full_state_simulator/common.wgsl deleted file mode 100644 index a68911e15cc..00000000000 --- a/source/simulators/src/gpu_full_state_simulator/common.wgsl +++ /dev/null @@ -1,1520 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -// See https://webgpufundamentals.org/webgpu/lessons/webgpu-wgsl.html for an overview -// See https://www.w3.org/TR/WGSL/ for the details -// See https://webgpu.github.io/webgpu-samples/ for examples - -// WGSL has pipeline overridables, but they're a pain and limited, so just string replace constants here -const QUBIT_COUNT: i32 = {{QUBIT_COUNT}}; -const RESULT_COUNT: u32 = {{RESULT_COUNT}}; -const WORKGROUPS_PER_SHOT: i32 = {{WORKGROUPS_PER_SHOT}}; -const ENTRIES_PER_THREAD: i32 = {{ENTRIES_PER_THREAD}}; -const THREADS_PER_WORKGROUP: i32 = {{THREADS_PER_WORKGROUP}}; -const MAX_QUBIT_COUNT: i32 = {{MAX_QUBIT_COUNT}}; -const MAX_QUBITS_PER_WORKGROUP: i32 = {{MAX_QUBITS_PER_WORKGROUP}}; - -const ERR_INVALID_PROBS = 1u; -const ERR_INVALID_THREAD_TOTAL = 2u; -// A loss policy was stamped onto a gate that does not support it. The host -// validates loss policies per gate before submission, so this indicates a bug. -// Errors 3-31 are reserved for base.wgsl and adaptive.wgsl -const ERR_UNSUPPORTED_LOSS_POLICY = 32u; - -// Tolerance for probabilities to sum to 1.0 -const PROB_THRESHOLD: f32 = 0.0001; - -// Always use 32 threads per workgroup for max concurrency on most current GPU hardware -const MAX_WORKGROUP_SUM_PARTITIONS: i32 = 1 << u32(MAX_QUBIT_COUNT - MAX_QUBITS_PER_WORKGROUP); - -// Operation IDs -const OPID_ID = 0u; -const OPID_RESETZ = 1u; -const OPID_X = 2u; -const OPID_Y = 3u; -const OPID_Z = 4u; -const OPID_H = 5u; -const OPID_S = 6u; -const OPID_SAdj = 7u; -const OPID_T = 8u; -const OPID_TAdj = 9u; -const OPID_RX = 12u; -const OPID_RY = 13u; -const OPID_RZ = 14u; -const OPID_CX = 15u; -const OPID_CZ = 16u; -const OPID_RXX = 17u; -const OPID_RYY = 18u; -const OPID_RZZ = 19u; -const OPID_MZ = 21u; -const OPID_MRESETZ = 22u; -const OPID_SWAP = 24u; -const OPID_MAT1Q = 25u; -const OPID_MAT2Q = 26u; -const OPID_CY = 29u; - -const OPID_PAULI_NOISE_1Q = 128u; -const OPID_PAULI_NOISE_2Q = 129u; -const OPID_LOSS_NOISE = 130u; -const OPID_CORRELATED_NOISE = 131u; - -// If the application of noise results in a custom matrix, it will have been stored in the shot buffer -// These OPIDs indicate to use that matrix and for how many qubits. (The qubit ids are in the original Op) -const OPID_SHOT_BUFF_1Q = 256u; -const OPID_SHOT_BUFF_2Q = 257u; - -struct WorkgroupSums { - qubits: array, // Each vec2f holds (zero_probability, one_probability) -}; - -struct WorkgroupCollationBuffer { - sums: array, -}; - -fn is_1q_phase_gate(op_id: u32) -> bool { - return (op_id == OPID_S || op_id == OPID_SAdj || op_id == OPID_T || op_id == OPID_TAdj || op_id == OPID_RZ); -} - -fn is_1q_op(op_id: u32) -> bool { - return ((op_id >= OPID_ID && op_id <= OPID_RZ) || - op_id == OPID_MZ || op_id == OPID_MRESETZ || - op_id == OPID_MAT1Q || op_id == OPID_SHOT_BUFF_1Q); -} - -fn shot_init_per_op(shot_idx: u32) { - let shot = &shots[shot_idx]; - - // Default to 1.0 renormalization (i.e., no renormalization needed). MResetZ or noise affecting the - // overall probability distribution (e.g. loss or amplitude damping) will update this if needed. - shot.renormalize = 1.0; - shot.qubits_updated_last_op_mask = 0u; - - // Generate the next set of random numbers to use for noise and measurement - shot.rand_pauli = next_rand_f32(shot_idx); - shot.rand_damping = next_rand_f32(shot_idx); - shot.rand_dephase = next_rand_f32(shot_idx); - shot.rand_measure = next_rand_f32(shot_idx); - // Reserved draw: qubit loss is now sampled from the combined `rand_pauli` - // distribution rather than its own value, but we still advance the RNG by - // one draw here to keep the per-op random stream (and thus seeded results) - // identical to the previous loss model. - next_rand_f32(shot_idx); -} - -// Resets the entire shot state, including RNG, probabilities, and per-qubit tracking. -fn reset_all(shot_idx: i32) { - let shot = &shots[shot_idx]; - - // One of the main goals of the shot_id is to seed the RNG state uniquely per shot - let rng_seed = uniforms.rng_seed; - let shot_id = u32(uniforms.batch_start_shot_id + shot_idx); - - // Due to DX12 backend issues, we can't just assign a zeroed struct, so manually reset all fields - // DX12-start-strip - *shot = ShotData(); - // DX12-end-strip - shot.shot_id = shot_id; - - // After init, start execution from the first op - shot.next_op_idx = 0u; - - shot.rng_state.x[0] = rng_seed ^ hash_pcg(shot_id); - shot.rng_state.x[1] = rng_seed ^ hash_pcg(shot_id + 1); - shot.rng_state.x[2] = rng_seed ^ hash_pcg(shot_id + 2); - shot.rng_state.x[3] = rng_seed ^ hash_pcg(shot_id + 3); - shot.rng_state.x[4] = rng_seed ^ hash_pcg(shot_id + 4); - - shot.op_type = 0; - shot.op_idx = 0; - - // rand_* will be initialized in shot_init_per_op when preparing the first op - shot.duration = 0.0; - shot.renormalize = 1.0; - - shot.qubit_is_0_mask = (1u << u32(QUBIT_COUNT)) - 1u; // All qubits are |0> - shot.qubit_is_1_mask = 0u; - shot.qubits_updated_last_op_mask = 0; - shot.pending_loss_mask = 0u; - - // Initialize all qubit probabilities to 100% |0> - for (var i: i32 = 0; i < QUBIT_COUNT; i++) { - shot.qubit_state[i].zero_probability = 1.0; - shot.qubit_state[i].one_probability = 0.0; - shot.qubit_state[i].heat = 0.0; - shot.qubit_state[i].idle_since = 0.0; - } - - // unitary will be set in prepare_op -} - -fn update_qubit_state(shot_idx: u32) { - let shot = &shots[shot_idx]; - - // If any qubits were updated in the last op, we may need to sum workgroup probabilities into the shot state - // This is only needed if multiple workgroups were used for the shot execution. If not, then the - // single workgroup for the shot would have written directly to the shot state already. - - // For each qubit that was updated in the last op - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - let qubit_mask: u32 = 1u << q; - if ((shot.qubits_updated_last_op_mask & qubit_mask) != 0u) { - // Sum the workgroup collation entries for this qubit into the shot state - // Note: We ignore the fact a qubit may be 'lost' here. It should already be - // in the |0> state if lost, so summing the probabilities is still valid. - var total_zero: f32 = 0.0; - var total_one: f32 = 0.0; - - if (WORKGROUPS_PER_SHOT > 1) { - // Offset into workgroup collation buffer based on shot index - let offset = shot_idx * u32(WORKGROUPS_PER_SHOT); - for (var wkg_idx: u32 = 0u; wkg_idx < u32(WORKGROUPS_PER_SHOT); wkg_idx++) { - let sums = workgroup_collation.sums[wkg_idx + offset]; - total_zero = total_zero + sums.qubits[q].x; - total_one = total_one + sums.qubits[q].y; - } - } else { - // Single workgroup per shot case - just read directly from the shot - total_zero = shot.qubit_state[q].zero_probability; - total_one = shot.qubit_state[q].one_probability; - } - - // Update the shot state with the summed probabilities - // Round to 0 or 1 if extremely close to mitigate minor floating point errors - // TODO: Use PROB_THRESHOLD constant here? - if (total_zero < 0.000001) { total_zero = 0.0; } - if (total_one < 0.000001) { total_one = 0.0; } - if (total_zero > 0.999999) { total_zero = 1.0; } - if (total_one > 0.999999) { total_one = 1.0; } - - shot.qubit_state[q].zero_probability = total_zero; - shot.qubit_state[q].one_probability = total_one; - - // NOTE: Any kind of operation with a NaN float value results in a NaN, or false for logical comparisons - // So beware of conditions that may not behave as expected if NaN values are possible. - let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; - if !within_threshold { - // Populate the diagnostics buffer, if not already set - let old_value = atomicCompareExchangeWeak( - &diagnostics.error_code, - 0u, - ERR_INVALID_PROBS); - if old_value.exchanged { - // This is the first error - fill in the details - diagnostics.extra1 = q; - diagnostics.extra2 = total_zero; - diagnostics.extra3 = total_one; - // DX12 backend has issues assigning structs. See https://github.com/gfx-rs/wgpu/issues/8552 - // DX12-start-strip - diagnostics.shot = *shot; - diagnostics.op = ops[shot.op_idx]; - // DX12-end-strip - } - // Store the error value (if none set already) - let err_index = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak( - &results[err_index], - 0u, - ERR_INVALID_PROBS); - } - - // Update the masks for definite states - shot.qubit_is_0_mask = select( - shot.qubit_is_0_mask & ~qubit_mask, - shot.qubit_is_0_mask | qubit_mask, - total_zero == 1.0); - shot.qubit_is_1_mask = select( - shot.qubit_is_1_mask & ~qubit_mask, - shot.qubit_is_1_mask | qubit_mask, - total_one == 1.0); - } - } -} - -// Build a measure-and-reset (or measure-only) instrument for `qubit` given a -// measured `result`, store it in the shot buffer, set up renormalization, and -// mark the qubit as no longer in a definite basis state so the execute stage -// recomputes its probabilities. Shared by `prep_measure_reset` and -// `prep_loss_commit`; the caller sets `shot.op_idx` and `shot.op_type`. -fn prep_measure_reset_instrument(shot_idx: u32, qubit: u32, result: u32, resets_to_zero: bool) { - let shot = &shots[shot_idx]; - - // Construct the measurement/reset instrument based on the measured result - // Put the instrument into the shot buffer for the execute_op stage to apply - if resets_to_zero { - // Reset variants (MResetZ, ResetZ): - // Result=0: [[1,0],[0,0]] - project onto |0⟩ (already there) - // Result=1: [[0,1],[0,0]] - swap |1⟩ into |0⟩ slot (reset) - shot.unitary[0] = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); - shot.unitary[1] = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); - shot.unitary[4] = vec2f(); - shot.unitary[5] = vec2f(); - } else { - // Measure-only (MZ): - // Result=0: [[1,0],[0,0]] - project onto |0⟩ - // Result=1: [[0,0],[0,1]] - project onto |1⟩ (keep in place) - shot.unitary[0] = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); - shot.unitary[1] = vec2f(); - shot.unitary[4] = vec2f(); - shot.unitary[5] = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); - } - - shot.renormalize = select( - 1.0 / sqrt(shot.qubit_state[qubit].zero_probability), - 1.0 / sqrt(shot.qubit_state[qubit].one_probability), - result == 1u); - - // We don't want the measurement pass to skip over this qubit, so ensure it's marked as not in a definite state - shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << qubit); - shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << qubit); - - // Set the qubits_updated_last_op_mask to all except those that were already in a definite - // state (so we don't waste time updating probabilities that are already known). Note that - // next 'prepare_op' should set the just measured qubit into a definite 0 or 1 state. - shot.qubits_updated_last_op_mask = - // A mask with all qubits set - ((1u << u32(QUBIT_COUNT)) - 1u) - // Exclude qubits already in definite states - & ~(shot.qubit_is_0_mask | shot.qubit_is_1_mask); -} - -fn prep_measure_reset(shot_idx: u32, op_idx: u32, is_loss: bool, stores_result: bool, resets_to_zero: bool) { - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - let qubit = get_measure_qubit(shot_idx, op_idx); - - // Choose measurement result based on qubit probabilities and random number - let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); - - // If this is being called due to loss noise, we don't write the result back to the results buffer - // Instead, mark the qubit as lost by setting the heat to -1.0 - if !is_loss { - if stores_result { - let result_id = get_measure_result(shot_idx, op_idx); // Result id to store the measurement result in is stored in q2 - - // If the qubit is already marked as lost, just report that and exit. It's already in the zero - // state so nothing to update or renormalize. The execute op should be a no-op (ID) - if shot.qubit_state[qubit].heat == -1.0 { - atomicStore(&results[(shot_idx * RESULT_COUNT) + result_id], 2u); - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - // Qubit get reloaded after a Measurement, so set the heat back to 0.0 - shot.qubit_state[qubit].heat = 0.0; - return; - } else { - atomicStore(&results[(shot_idx * RESULT_COUNT) + result_id], result); - } - } else { - // No result to store (e.g. ResetZ). If the qubit is lost, it's already in the zero - // state so nothing to update. Just set to ID and return. - if shot.qubit_state[qubit].heat == -1.0 { - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - return; - } - } - } else { - shot.qubit_state[qubit].heat = -1.0; - } - - prep_measure_reset_instrument(shot_idx, qubit, result, resets_to_zero); - - shot.op_idx = op_idx; - // Use OPID_MRESETZ as the op_type for all three variants in execute stage - // (they all use the same matrix-apply + update_all_qubit_probs path) - shot.op_type = OPID_MRESETZ; -} - -// Starting from the given index, return the next index if pauli noise, else 0 -fn get_pauli_noise_idx(op_idx: u32) -> u32 { - if (arrayLength(&ops) > (op_idx + 1)) { - let op = &ops[op_idx + 1]; - if (op.id == OPID_PAULI_NOISE_1Q || op.id == OPID_PAULI_NOISE_2Q) { - return op_idx + 1u; - } - } - return 0u; -} - -// From the starting index given, return the next index if loss noise, else 0 -fn get_loss_idx(op_idx: u32) -> u32 { - if (arrayLength(&ops) > (op_idx + 1)) { - let op = &ops[op_idx + 1]; - if (op.id == OPID_LOSS_NOISE) { - return op_idx + 1u; - } - } - return 0u; -} - -// Loss policy values. These are stamped onto a gate op's `q3` field by the host -// (see `LossPolicy::as_u32` on the Rust side) and tell the shader how to handle -// the gate when one of its operands is lost. `0` means "no policy stamped", -// which the shader treats the same as SKIP. -const LOSS_POLICY_SKIP = 0u; -const LOSS_POLICY_PROPAGATE = 1u; -const LOSS_POLICY_DEGRADE = 2u; -const LOSS_POLICY_RESIDUAL_S_DAGGER = 3u; -const LOSS_POLICY_APPLY_ANYWAY = 4u; - -// Returns true if the gate at `op_idx` touches at least one lost qubit. -// `q1`/`q2` are the (resolved) operands of the gate. -fn gate_has_lost_operand(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) -> bool { - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - if (shot.qubit_state[q1].heat == -1.0) { - return true; - } - let is_2q = !is_1q_op(op.id); - return is_2q && (shot.qubit_state[q2].heat == -1.0); -} - -// Builds a 4x4 (in shot.unitary) that applies the 1-qubit matrix `m` (given as -// m00,m01,m10,m11) to `target_is_q2 ? q2 : q1` and identity to the other qubit -// of the pair. The lost qubit is in the |0> state, so the identity factor keeps -// it there. The 2-qubit basis is |q1 q2>, so the row/col index is -// (2 * q1_bit + q2_bit). -fn set_1q_on_pair_unitary(shot_idx: u32, target_is_q2: bool, - m00: vec2f, m01: vec2f, m10: vec2f, m11: vec2f) { - let shot = &shots[shot_idx]; - // Zero the whole 4x4 first. - for (var i = 0u; i < 16u; i++) { - shot.unitary[i] = vec2f(0.0, 0.0); - } - if target_is_q2 { - // Acts on q2 (low bit): block-diagonal diag(M, M). - // Top-left block (q1 = 0): - shot.unitary[0] = m00; shot.unitary[1] = m01; - shot.unitary[4] = m10; shot.unitary[5] = m11; - // Bottom-right block (q1 = 1): - shot.unitary[10] = m00; shot.unitary[11] = m01; - shot.unitary[14] = m10; shot.unitary[15] = m11; - } else { - // Acts on q1 (high bit): M (x) I. - shot.unitary[0] = m00; shot.unitary[2] = m01; - shot.unitary[8] = m10; shot.unitary[10] = m11; - shot.unitary[5] = m00; shot.unitary[7] = m01; - shot.unitary[13] = m10; shot.unitary[15] = m11; - } -} - -// Multiplies one row of the 4x4 pair unitary (in shot.unitary) by -i, in place. -// Folding a diag(1, -i) = S-dagger factor on one qubit into a 2-qubit matrix -// scales the rows whose target-qubit bit is 1 by -i. For a complex entry -// (x + y i), (x + y i) * -i = y - x i. -fn scale_pair_unitary_row_by_neg_i(shot_idx: u32, row: u32) { - let shot = &shots[shot_idx]; - for (var c = 0u; c < 4u; c++) { - let e = shot.unitary[row * 4u + c]; - shot.unitary[row * 4u + c] = vec2f(e.y, -e.x); - } -} - -// Sets up the shot to execute a 2-qubit shot-buffer op on the gate's operands. -fn finish_2q_shot_buffer(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { - let shot = &shots[shot_idx]; - shot.op_idx = op_idx; - shot.op_type = OPID_SHOT_BUFF_2Q; - shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); -} - -// Loses a single surviving `qubit` for the PROPAGATE policy: samples a -// measurement outcome, collapses the qubit to that outcome and resets it to -// |0>, and marks it lost (heat = -1.0). The collapse is expressed as a 2-qubit -// tensor on the gate's operands (reset on `qubit`, identity on the lost -// partner, which is already in |0>), reusing the standard shot-buffer execute -// path. `qubit` must be one of the gate's two operands `q1`/`q2`. -fn propagate_loss_to_qubit(shot_idx: u32, op_idx: u32, q1: u32, q2: u32, qubit: u32) { - let shot = &shots[shot_idx]; - - let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); - - // Reset instrument (project + move |1> into |0> slot), same as MResetZ: - // result==0: [[1,0],[0,0]] - // result==1: [[0,1],[0,0]] - let m00 = select(vec2f(1.0, 0.0), vec2f(0.0, 0.0), result == 1u); - let m01 = select(vec2f(0.0, 0.0), vec2f(1.0, 0.0), result == 1u); - let m10 = vec2f(0.0, 0.0); - let m11 = vec2f(0.0, 0.0); - - let target_is_q2 = (qubit == q2); - set_1q_on_pair_unitary(shot_idx, target_is_q2, m00, m01, m10, m11); - - // Renormalize by the measured branch probability. - shot.renormalize = select( - 1.0 / sqrt(shot.qubit_state[qubit].zero_probability), - 1.0 / sqrt(shot.qubit_state[qubit].one_probability), - result == 1u); - - // Mark the qubit lost and clear its definite-state bits so the probability - // pass recomputes it. - shot.qubit_state[qubit].heat = -1.0; - shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << qubit); - shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << qubit); - - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); -} - -// Records an error `code` for `shot_idx` in both the diagnostics buffer and the -// shot's result-code slot, mirroring the reporting done elsewhere in this file. -// Used for conditions the host guarantees never occur (e.g. a loss policy that -// is not valid for a given gate). -fn report_shot_error(shot_idx: u32, code: u32) { - atomicCompareExchangeWeak(&diagnostics.error_code, 0u, code); - let err_index = (shot_idx + 1u) * RESULT_COUNT - 1u; - atomicCompareExchangeWeak(&results[err_index], 0u, code); -} - -// Handles a gate whose operand(s) include at least one lost qubit, according to -// the loss policy stamped on the op's `policy` field. `q1`/`q2` are the -// (resolved) operands. The gate body is fully handled here (degraded unitary, -// loss propagation, or turned into Id); the caller must not run the original -// gate afterwards. Any attached Pauli noise is applied separately to the -// surviving operand via `apply_2q_pauli_noise_on_survivor`. -fn handle_lost_operand_policy(shot_idx: u32, op_idx: u32, q1: u32, q2: u32) { - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - let is_1q = is_1q_op(op.id); - let is_2q = !is_1q; - let policy = op.policy; - - // Loss policies only make sense for multi-qubit gates. - // If this is a single-qubit gate, skip it entirely. - if (is_1q) { - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - return; - } - - let q1_lost = shot.qubit_state[q1].heat == -1.0; - let q2_lost = is_2q && (shot.qubit_state[q2].heat == -1.0); - let has_survivor = is_2q && !(q1_lost && q2_lost); - // The surviving operand (only meaningful when has_survivor is true). - let survivor = select(q1, q2, q1_lost); - let survivor_is_q2 = q1_lost; - - // SWAP is special: it physically relocates the two qubits, so their loss - // state is always exchanged regardless of the policy (the policy only - // governs whether the unitary runs). Handle it explicitly here. - if (op.id == OPID_SWAP) { - switch policy { - case LOSS_POLICY_PROPAGATE { - propagate_loss_to_qubit(shot_idx, op_idx, q1, q2, survivor); - return; - } - case LOSS_POLICY_RESIDUAL_S_DAGGER { - // Match the CPU/stabilizer SWAP + residual S-dagger semantics: - // 1. Apply the full SWAP (shot.unitary already holds it). - // 2. Apply S-dagger = diag(1, -i) to the (originally) lost - // operand's position, which after the SWAP holds the - // survivor's amplitudes. - // 3. Exchange the per-qubit loss flag (heat) of the operands. - - // Fold the S-dagger into the SWAP matrix by scaling, by -i, the - // two rows of the |q1 q2> pair matrix whose lost-qubit bit is 1. - // q1 is the high bit (rows 2, 3); q2 is the low bit (rows 1, 3). - let lost_row = select(1u, 2u, q1_lost); - scale_pair_unitary_row_by_neg_i(shot_idx, lost_row); - scale_pair_unitary_row_by_neg_i(shot_idx, 3u); - // Exchange the per-qubit loss flag (heat) of the two operands. - let heat1 = shot.qubit_state[q1].heat; - shot.qubit_state[q1].heat = shot.qubit_state[q2].heat; - shot.qubit_state[q2].heat = heat1; - // The 2-qubit execute path skips amplitudes for qubits known to be - // in a definite state, which would skip the amplitudes SWAP needs to move. - // Clear those bits for both operands so the swap is actually applied. - shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~((1u << q1) | (1u << q2)); - shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~((1u << q1) | (1u << q2)); - // shot.unitary now holds (S-dagger on lost) * SWAP. - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); - return; - } - case LOSS_POLICY_APPLY_ANYWAY { - // Exchange the per-qubit loss flag (heat) of the two operands. - let heat1 = shot.qubit_state[q1].heat; - shot.qubit_state[q1].heat = shot.qubit_state[q2].heat; - shot.qubit_state[q2].heat = heat1; - // The 2-qubit execute path skips amplitudes for qubits known to be - // in a definite state, which would skip the amplitudes SWAP needs to move. - // Clear those bits for both operands so the swap is actually applied. - shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~((1u << q1) | (1u << q2)); - shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~((1u << q1) | (1u << q2)); - // shot.unitary already holds the SWAP matrix (set by the caller). - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); - return; - } - case LOSS_POLICY_SKIP { - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - return; - } - default { - // SWAP only supports SKIP, PROPAGATE, RESIDUAL_S_DAGGER, and - // APPLY_ANYWAY. Any other policy (e.g. DEGRADE) is rejected by - // the host, so reaching here indicates a bug. - report_shot_error(shot_idx, ERR_UNSUPPORTED_LOSS_POLICY); - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - return; - } - } - } - - // APPLY_ANYWAY is only valid for SWAP, which is handled above. Reaching here - // with it on any other gate is rejected by the host, so it indicates a bug. - if (policy == LOSS_POLICY_APPLY_ANYWAY) { - report_shot_error(shot_idx, ERR_UNSUPPORTED_LOSS_POLICY); - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - return; - } - - if (policy == LOSS_POLICY_PROPAGATE && has_survivor) { - propagate_loss_to_qubit(shot_idx, op_idx, q1, q2, survivor); - return; - } - - if (policy == LOSS_POLICY_RESIDUAL_S_DAGGER && has_survivor) { - // Apply S-dagger = diag(1, -i) to the surviving operand. - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(1.0, 0.0), vec2f(0.0, 0.0), - vec2f(0.0, 0.0), vec2f(0.0, -1.0)); - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); - return; - } - - // DEGRADE is only valid for the two-qubit rotations (Rxx/Ryy/Rzz), so the - // op is guaranteed to be one of them when a survivor exists. - if (policy == LOSS_POLICY_DEGRADE && has_survivor) { - // Degrade the two-qubit rotation to its single-qubit version on the - // survivor. The op's unitary[0] holds cos(θ/2) for Rxx/Ryy; we recover - // the angle to build the 1-qubit rotation matrix. - let cos_half = op.unitary[0].x; - if (op.id == OPID_RXX) { - // Rx(θ) = [[c, -i s], [-i s, c]], where s = sin(θ/2). - let s = op.unitary[3].y * -1.0; // unitary[3] = (0, -sin(θ/2)) - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(cos_half, 0.0), vec2f(0.0, -s), - vec2f(0.0, -s), vec2f(cos_half, 0.0)); - } else if (op.id == OPID_RYY) { - // Ry(θ) = [[c, -s], [s, c]], where s = sin(θ/2). - let s = op.unitary[3].y; // unitary[3] = (0, sin(θ/2)) for Ryy - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(cos_half, 0.0), vec2f(-s, 0.0), - vec2f(s, 0.0), vec2f(cos_half, 0.0)); - } else { - // Rzz -> Rz(θ). The GPU Rz convention is [[1, 0], [0, e^{iθ}]], - // and unitary[5] = e^{iθ} holds the full-angle phase. - let phase = op.unitary[5]; - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(1.0, 0.0), vec2f(0.0, 0.0), - vec2f(0.0, 0.0), phase); - } - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); - return; - } - - // SKIP, or any policy when both operands are lost (no survivor to act on): - // skip the gate entirely. - shot.op_type = OPID_ID; - shot.op_idx = op_idx; -} - -fn apply_1q_pauli_noise(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32) { - // NOTE: Assumes that whatever prepared the program ensured that noise_op.q1 matches op.q1 and - // that op is a 1-qubit gate. `q1` is the resolved target qubit (may be - // dynamic for the adaptive interpreter, where op.q1 is only a placeholder). - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - let noise_op = &ops[noise_idx]; - - // Categorical outcome probabilities by 3-bit term (X=1, Z=2, Y=3, L=4), - // stored at flat slot k = term in `unitary[k / 2][k % 2]`. The identity - // outcome (slot 0) is implicit. - let p_x = noise_op.unitary[0].y; - let p_z = noise_op.unitary[1].x; - let p_y = noise_op.unitary[1].y; - let p_loss = noise_op.unitary[2].x; - - shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer - - let rand = shot.rand_pauli; - if (rand < p_x) { - // Apply the X permutation (basically swap the rows) - shot.unitary[0] = op.unitary[4]; - shot.unitary[1] = op.unitary[5]; - shot.unitary[4] = op.unitary[0]; - shot.unitary[5] = op.unitary[1]; - } else if (rand < (p_x + p_y)) { - // Apply the Y permutation (swap rows with negated |0> state) - shot.unitary[0] = cplxNeg(op.unitary[4]); - shot.unitary[1] = cplxNeg(op.unitary[5]); - shot.unitary[4] = op.unitary[0]; - shot.unitary[5] = op.unitary[1]; - } else if (rand < (p_x + p_y + p_z)) { - // Apply Z error (negate |1> state) - shot.unitary[0] = op.unitary[0]; - shot.unitary[1] = op.unitary[1]; - shot.unitary[4] = cplxNeg(op.unitary[4]); - shot.unitary[5] = cplxNeg(op.unitary[5]); - } else { - // Either loss or no noise: the gate executes unmodified. If loss was - // sampled, schedule a loss commit for this qubit; a following - // loss-commit op performs the measure + reset. - if (rand < (p_x + p_z + p_y + p_loss)) { - shot.pending_loss_mask |= (1u << q1); - } - // No noise. Set the op_type back to the op.id value if it's Id, MResetZ, MZ, or ResetZ, as they get handled specially in execute_op - if (op.id == OPID_ID || op.id == OPID_MRESETZ || op.id == OPID_MZ || op.id == OPID_RESETZ) { - shot.op_type = op.id; - } - if (is_1q_phase_gate(op.id)) { - // For phase gates, treat everything as RZ for execution purposes - shot.op_type = OPID_RZ; - } - } - - shot.op_idx = op_idx; - if (shot.op_type == OPID_ID || shot.op_type == OPID_RZ) { - shot.qubits_updated_last_op_mask = 0u; - } else { - shot.qubits_updated_last_op_mask = 1u << q1; - }; -} - -fn apply_2q_pauli_noise(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32, q2: u32) { - let shot = &shots[shot_idx]; - let op = &ops[op_idx]; - let noise_op = &ops[noise_idx]; - - // The categorical distribution over the 25 (q1_term, q2_term) outcomes is - // stored at flat slot k = q1_term * 5 + q2_term in `unitary[k / 2][k % 2]`. - // Terms use the 3-bit encoding: I=0, X=1, Z=2, Y=3, L=4. The II slot (0) is - // implicit and carries the remaining probability. - var rand = shot.rand_pauli; - var q1_term = 0; - var q2_term = 0; - - // Find the terms to apply based on the random number and the probabilities - for (var a = 0; a < 5; a = a + 1) { - for (var b = 0; b < 5; b = b + 1) { - let k = a * 5 + b; - if (k == 0) { continue; } // II carries no stored probability - let slot = noise_op.unitary[k / 2]; - let p_ab = select(slot.x, slot.y, (k & 1) == 1); - if (rand < p_ab) { - q1_term = a; - q2_term = b; - // Break out of both loops - a = 5; - b = 5; - } else { - rand = rand - p_ab; - } - } - } - - // Schedule loss commits for any qubit whose sampled term is loss (L = 4). - // A following loss-commit op performs the measure + reset. - if (q1_term == 4) { shot.pending_loss_mask |= (1u << q1); } - if (q2_term == 4) { shot.pending_loss_mask |= (1u << q2); } - - // A Pauli fault (X, Z, Y = 1, 2, 3) is fused into the gate by permuting its - // rows. Loss (4) and identity (0) leave the gate unmodified for that qubit. - let q1_pauli = q1_term >= 1 && q1_term <= 3; - let q2_pauli = q2_term >= 1 && q2_term <= 3; - - if (q1_pauli || q2_pauli) { - // Get the rows of the 2 qubit unitary - var op_row_0 = getOpRow(op_idx, 0); - var op_row_1 = getOpRow(op_idx, 1); - var op_row_2 = getOpRow(op_idx, 2); - var op_row_3 = getOpRow(op_idx, 3); - - // Apply the Paulis to the matrices. Note this is just permuting the rows, and appliction - // commutes, so we can apply them in any order. High order bit is q1. Low order bit is q2. - // X on q1 is rows 2<>0 and 3<>1, X on q2 is rows 1<>0 and 3<>2, etc. - // Y on q1 is rows -2<>0 and -3<>1, Y on q2 is rows -1<>0 and -3<>2 - // Z on q1 is -2 and -3, Z on q2 is -1 and -3 - - // Apply the q1 permutations as needed - if (q1_term == 1) { - // Apply the X permutation - let old_row_0 = op_row_0; - let old_row_1 = op_row_1; - op_row_0 = op_row_2; - op_row_1 = op_row_3; - op_row_2 = old_row_0; - op_row_3 = old_row_1; - } else if (q1_term == 3) { - // Apply the Y permutation - let old_row_0 = op_row_0; - let old_row_1 = op_row_1; - op_row_0 = rowNeg(op_row_2); - op_row_1 = rowNeg(op_row_3); - op_row_2 = old_row_0; - op_row_3 = old_row_1; - } else if (q1_term == 2) { - // Apply Z permutation - op_row_2 = rowNeg(op_row_2); - op_row_3 = rowNeg(op_row_3); - } - // Apply the q2 permutations as needed - if (q2_term == 1) { - // Apply the X permutation - let old_row_0 = op_row_0; - let old_row_2 = op_row_2; - op_row_0 = op_row_1; - op_row_2 = op_row_3; - op_row_1 = old_row_0; - op_row_3 = old_row_2; - } else if (q2_term == 3) { - // Apply the Y permutation - let old_row_0 = op_row_0; - let old_row_2 = op_row_2; - op_row_0 = rowNeg(op_row_1); - op_row_2 = rowNeg(op_row_3); - op_row_1 = old_row_0; - op_row_3 = old_row_2; - } else if (q2_term == 2) { - // Apply Z permutation - op_row_1 = rowNeg(op_row_1); - op_row_3 = rowNeg(op_row_3); - } - // Write the rows back to the shot buffer unitary - setUnitaryRow(shot_idx, 0u, op_row_0); - setUnitaryRow(shot_idx, 1u, op_row_1); - setUnitaryRow(shot_idx, 2u, op_row_2); - setUnitaryRow(shot_idx, 3u, op_row_3); - shot.op_type = OPID_SHOT_BUFF_2Q; - } else { - // No Pauli fault to fuse (identity or loss only). Leave if CX, CY, CZ, or RZZ as they get handled specially in execute_op - if (op.id == OPID_CX || op.id == OPID_CY || op.id == OPID_CZ || op.id == OPID_RZZ) { - shot.op_type = op.id; - } else { - shot.op_type = OPID_SHOT_BUFF_2Q; - } - } - shot.op_idx = op_idx; - if (shot.op_type == OPID_CZ || shot.op_type == OPID_RZZ) { - shot.qubits_updated_last_op_mask = 0u; - } else { - shot.qubits_updated_last_op_mask = (1u << q1 ) | (1u << q2); - } -} - -// Left-multiplies the 4x4 pair unitary already in `shot.unitary` by a single -// Pauli (term: X=1, Z=2, Y=3) acting on `target_is_q2 ? q2 : q1`. This is the -// same row permutation/negation that `apply_2q_pauli_noise` fuses, just applied -// to the policy-degraded gate rather than the original op. Note the Y branch -// uses real signs (i.e. -i*Y), matching `apply_2q_pauli_noise`; the resulting -// global phase is unobservable for a Pauli noise channel. -fn fuse_1q_pauli_on_pair_unitary(shot_idx: u32, target_is_q2: bool, term: u32) { - let si = i32(shot_idx); - var row_0 = getUnitaryRow(si, 0u); - var row_1 = getUnitaryRow(si, 1u); - var row_2 = getUnitaryRow(si, 2u); - var row_3 = getUnitaryRow(si, 3u); - - if (!target_is_q2) { - // Acting on q1 (high bit): rows {0,1} <-> {2,3}. - if (term == 1u) { // X - let o0 = row_0; let o1 = row_1; - row_0 = row_2; row_1 = row_3; - row_2 = o0; row_3 = o1; - } else if (term == 3u) { // Y - let o0 = row_0; let o1 = row_1; - row_0 = rowNeg(row_2); row_1 = rowNeg(row_3); - row_2 = o0; row_3 = o1; - } else { // Z - row_2 = rowNeg(row_2); row_3 = rowNeg(row_3); - } - } else { - // Acting on q2 (low bit): rows {0,2} <-> {1,3}. - if (term == 1u) { // X - let o0 = row_0; let o2 = row_2; - row_0 = row_1; row_2 = row_3; - row_1 = o0; row_3 = o2; - } else if (term == 3u) { // Y - let o0 = row_0; let o2 = row_2; - row_0 = rowNeg(row_1); row_2 = rowNeg(row_3); - row_1 = o0; row_3 = o2; - } else { // Z - row_1 = rowNeg(row_1); row_3 = rowNeg(row_3); - } - } - - setUnitaryRow(shot_idx, 0u, row_0); - setUnitaryRow(shot_idx, 1u, row_1); - setUnitaryRow(shot_idx, 2u, row_2); - setUnitaryRow(shot_idx, 3u, row_3); -} - -// Applies the Pauli noise attached to a 2-qubit gate that had a lost operand. -// The gate body itself was already handled by `handle_lost_operand_policy` -// (which may have left a degraded 4x4 in `shot.unitary`, or turned the gate -// into Id for SKIP). This mirrors the CPU `apply_fault`: the joint (q1, q2) -// term is sampled, but only the operand still alive *after* the policy ran -// receives its term; a lost operand gets nothing. -// -// Because this is only reached when the gate has at least one lost operand, -// there is at most one surviving operand, so at most one single-qubit Pauli is -// fused. -fn apply_2q_pauli_noise_on_survivor(shot_idx: u32, op_idx: u32, noise_idx: u32, q1: u32, q2: u32) { - let shot = &shots[shot_idx]; - let noise_op = &ops[noise_idx]; - - // Surviving operand(s) after the policy ran (alive => heat != -1.0). - let q1_alive = shot.qubit_state[q1].heat != -1.0; - let q2_alive = shot.qubit_state[q2].heat != -1.0; - // Both lost (e.g. PROPAGATE collapsed the survivor): nothing to apply. - if (!q1_alive && !q2_alive) { - return; - } - - // Sample the joint (q1_term, q2_term) outcome (same encoding/layout as - // apply_2q_pauli_noise: I=0, X=1, Z=2, Y=3, L=4). - var rand = shot.rand_pauli; - var q1_term = 0; - var q2_term = 0; - for (var a = 0; a < 5; a = a + 1) { - for (var b = 0; b < 5; b = b + 1) { - let k = a * 5 + b; - if (k == 0) { continue; } - let slot = noise_op.unitary[k / 2]; - let p_ab = select(slot.x, slot.y, (k & 1) == 1); - if (rand < p_ab) { - q1_term = a; - q2_term = b; - a = 5; - b = 5; - } else { - rand = rand - p_ab; - } - } - } - - // The survivor's own term. (At most one operand is alive here.) - let survivor_is_q2 = !q1_alive; - let survivor = select(q1, q2, survivor_is_q2); - let term = select(q1_term, q2_term, survivor_is_q2); - - // Loss (4): schedule a loss commit for the survivor; a later loss-commit op - // performs the measure + reset. The gate set up by the policy still runs. - if (term == 4) { - shot.pending_loss_mask |= (1u << survivor); - return; - } - - // Identity (0): nothing to fuse; leave the policy's setup untouched. - if (term == 0) { - return; - } - - // Pauli (X=1, Z=2, Y=3): fuse onto the survivor. - if (shot.op_type == OPID_SHOT_BUFF_2Q) { - // The policy left a degraded 4x4 in shot.unitary; left-multiply it by - // the survivor Pauli. - fuse_1q_pauli_on_pair_unitary(shot_idx, survivor_is_q2, u32(term)); - } else { - // The policy turned the gate into Id (SKIP). Build a pair unitary that - // applies just the Pauli to the survivor and identity to the lost - // partner (which is in |0>). Real-sign Y matches the fuse path above. - if (term == 1) { // X - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(0.0, 0.0), vec2f(1.0, 0.0), - vec2f(1.0, 0.0), vec2f(0.0, 0.0)); - } else if (term == 3) { // Y (real-sign, i.e. -i*Y) - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(0.0, 0.0), vec2f(-1.0, 0.0), - vec2f(1.0, 0.0), vec2f(0.0, 0.0)); - } else { // Z - set_1q_on_pair_unitary(shot_idx, survivor_is_q2, - vec2f(1.0, 0.0), vec2f(0.0, 0.0), - vec2f(0.0, 0.0), vec2f(-1.0, 0.0)); - } - finish_2q_shot_buffer(shot_idx, op_idx, q1, q2); - } - - // The survivor's amplitudes may have been in a definite computational-basis - // state; clear its definite-state bits so the execute pass recomputes them - // after the Pauli (mirrors the SWAP handling in handle_lost_operand_policy). - shot.qubit_is_0_mask = shot.qubit_is_0_mask & ~(1u << survivor); - shot.qubit_is_1_mask = shot.qubit_is_1_mask & ~(1u << survivor); -} - -struct ShotParams { - shot_idx: i32, - shot_state_vector_start: i32, - workgroup_collation_idx: i32, - workgroup_idx_in_shot: i32, - thread_idx_in_shot: i32, - total_threads_per_shot: i32, - zero_entry_count: i32, - op_iterations: i32, -} - -fn get_shot_params( - workgroupId: u32, - tid: u32, - op_qubit_count: i32) -> ShotParams { - // Workgroups are per shot if 22 or less qubits, else 2 workgroups for 23 qubits, 4 for 24, etc.. - let shot_idx: i32 = i32(workgroupId) / WORKGROUPS_PER_SHOT; - let shot_state_vector_start: i32 = shot_idx * (1 << u32(QUBIT_COUNT)); - let workgroup_idx_in_shot: i32 = i32(workgroupId) % WORKGROUPS_PER_SHOT; - let thread_idx_in_shot: i32 = workgroup_idx_in_shot * THREADS_PER_WORKGROUP + i32(tid); - let total_threads_per_shot: i32 = WORKGROUPS_PER_SHOT * THREADS_PER_WORKGROUP; - - // If using multiple workgroups per shot, each workgroup will write its partial sums to the collation - // buffer for later summing by the prepare_op stage. If single workgroup per shot, no collation needed. - // Use -1 as a marker for single workgroup per shot case (in which case we should write directly to the shot). - let workgroup_collation_idx: i32 = select(-1, i32(workgroupId), WORKGROUPS_PER_SHOT > 1); - - let zero_entry_count: i32 = (1 << u32(QUBIT_COUNT)) >> u32(op_qubit_count); - let op_iterations: i32 = zero_entry_count / total_threads_per_shot; - - return ShotParams( - shot_idx, - shot_state_vector_start, - workgroup_collation_idx, - workgroup_idx_in_shot, - thread_idx_in_shot, - total_threads_per_shot, - zero_entry_count, - op_iterations - ); -} - -fn apply_1q_op(workgroupId: u32, tid: u32, q1: u32) { - let params = get_shot_params(workgroupId, tid, 1 /* qubits per op */); - let shot = &shots[params.shot_idx]; - let scale = shot.renormalize; - let lowMask = (1 << q1) - 1; - let highMask = (1 << u32(QUBIT_COUNT)) - 1 - lowMask; - let qubit_is_0_mask = i32(shots[params.shot_idx].qubit_is_0_mask); - let qubit_is_1_mask = i32(shots[params.shot_idx].qubit_is_1_mask); - - var summed_probs: vec4f = vec4f(); - - /* This loop is where all the real work happens. Try to keep this tight and efficient. - - We want a 'structure of arrays' like access pattern here for efficiency, so we process the state vector - in blocks where each thread in the workgroup(s) handle an adjacent entry to be processed. - - Each thread should start at the state vector shot start + 'thread_idx_in_shot', which is sequential across the workgroup threads - Each next entry for the thread is WORKGROUPS_PER_SHOT * THREADS_PER_WORKGROUP away. - */ - var entry_index = params.thread_idx_in_shot; - - for (var i = 0; i < params.op_iterations; i++) { - let offset0: i32 = (entry_index & lowMask) | ((entry_index & highMask) << 1); - let offset1: i32 = offset0 | (1 << q1); - - // See if we can skip doing any work for this pair, because the state vector entries to processes - // are both definitely 0.0, as we know they are for states where other qubits are in definite opposite state. - let skip_processing = ((offset0 & qubit_is_0_mask) != 0) || ((~offset1 & qubit_is_1_mask) != 0); - - if (!skip_processing) { - if shot.op_type == OPID_RZ { - // For RZ, we can skip reading/writing the |0> amplitude, as it is unchanged. - // Just apply the phase to the |1> amplitude. Probabilities also don't change. - let amp1: vec2f = stateVector[params.shot_state_vector_start + offset1]; - let new1 = cplxMul(amp1, shot.unitary[5]); - stateVector[params.shot_state_vector_start + offset1] = new1; - } else { - let amp0: vec2f = stateVector[params.shot_state_vector_start + offset0]; - let amp1: vec2f = stateVector[params.shot_state_vector_start + offset1]; - - let new0 = scale * (cplxMul(amp0, shot.unitary[0]) + cplxMul(amp1, shot.unitary[1])); - let new1 = scale * (cplxMul(amp0, shot.unitary[4]) + cplxMul(amp1, shot.unitary[5])); - - stateVector[params.shot_state_vector_start + offset0] = new0; - stateVector[params.shot_state_vector_start + offset1] = new1; - - if shot.op_type == OPID_MRESETZ || shot.op_type == OPID_LOSS_NOISE || scale != 1.0 { - // For MResetZ, loss-commit, or renormalization, update the probabilities for all qubits - update_all_qubit_probs(u32(offset0), new0, tid); - update_all_qubit_probs(u32(offset1), new1, tid); - } else { - summed_probs[0] += cplxMag2(new0); - summed_probs[1] += cplxMag2(new1); - } - } - } - entry_index += params.total_threads_per_shot; - } - - if scale == 1.0 && shot.op_type != OPID_RZ && shot.op_type != OPID_MRESETZ && shot.op_type != OPID_LOSS_NOISE { - // Update this thread's totals for the two qubits in the workgroup storage - qubitProbabilities[tid].zero[q1] = summed_probs[0]; - qubitProbabilities[tid].one[q1] = summed_probs[1]; - } -} - -fn apply_2q_op(workgroupId: u32, tid: u32, q1: u32, q2: u32) { - let params = get_shot_params(workgroupId, tid, 2 /* qubits per op */); - let shot = &shots[params.shot_idx]; - let update_probs = shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; - - // Sometimes a 2-qubit op may be converted to a no-op (ID) due to qubit loss etc., so skip processing in that case - // Calculate masks to split the index into low, mid, and high bits around the two qubits - let lowQubit = select(q1, q2, q1 > q2); - let hiQubit = select(q1, q2, q1 < q2); - - // Number of bits in each section - let lowBitCount = lowQubit; - let midBitCount = hiQubit - lowQubit - 1; - let hiBitCount = u32(QUBIT_COUNT) - hiQubit - 1; - - // The masks below help extract the low, mid, and high bits from the counter to use around the two qubits locations - let lowMask = (1 << lowBitCount) - 1; - let midMask = (1 << (lowBitCount + midBitCount)) - 1 - lowMask; - let hiMask = (1 << u32(QUBIT_COUNT)) - 1 - midMask - lowMask; - - // Each iteration processes 4 amplitudes (the four affected by the 2-qubit gate), so quarter as many iterations as chunk size - var entry_index = params.thread_idx_in_shot; - var summed_probs: vec4f = vec4f(); - - for (var i = 0; i < params.op_iterations; i++) { - // q1 is the control, q2 is the target - let offset00: i32 = (entry_index & lowMask) | ((entry_index & midMask) << 1) | ((entry_index & hiMask) << 2); - let offset01: i32 = offset00 | (1 << q2); - let offset10: i32 = offset00 | (1 << q1); - let offset11: i32 = offset10 | (1 << q2); - - let can_skip_processing = - (((u32(offset00) & shot.qubit_is_0_mask) != 0) || - ((~(u32(offset11)) & shot.qubit_is_1_mask) != 0)); - if !can_skip_processing { - switch shot.op_type { - case OPID_CZ { - let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; - stateVector[params.shot_state_vector_start + offset11] = cplxNeg(amp11); - // CZ doesn't change any probabilities, so no need to update summed_probs - } - case OPID_RZZ { - // Firt and last entries are unchanged, only need to update the middle two - let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; - let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; - // Unitary matrix second entry in the second row is 5, third entry in the third row is 10 - stateVector[params.shot_state_vector_start + offset01] = cplxMul(amp01, shot.unitary[5]); - stateVector[params.shot_state_vector_start + offset10] = cplxMul(amp10, shot.unitary[10]); - } - case OPID_CX { - // Need to read all 4 to update the probabilities correctly, but only swap the |10> and |11> entries - let amp00: vec2f = stateVector[params.shot_state_vector_start + offset00]; - let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; - let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; - let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; - stateVector[params.shot_state_vector_start + offset10] = amp11; - stateVector[params.shot_state_vector_start + offset11] = amp10; - summed_probs[0] += (cplxMag2(amp00) + cplxMag2(amp01)); - summed_probs[1] += (cplxMag2(amp11) + cplxMag2(amp10)); - summed_probs[2] += (cplxMag2(amp00) + cplxMag2(amp11)); - summed_probs[3] += (cplxMag2(amp01) + cplxMag2(amp10)); - } - case OPID_CY { - // Like CX, but swap |10> and |11> with +/- i phases. - let amp00: vec2f = stateVector[params.shot_state_vector_start + offset00]; - let amp01: vec2f = stateVector[params.shot_state_vector_start + offset01]; - let amp10: vec2f = stateVector[params.shot_state_vector_start + offset10]; - let amp11: vec2f = stateVector[params.shot_state_vector_start + offset11]; - stateVector[params.shot_state_vector_start + offset10] = vec2f(amp11.y, -amp11.x); // -i * |11> - stateVector[params.shot_state_vector_start + offset11] = vec2f(-amp10.y, amp10.x); // i * |10> - summed_probs[0] += (cplxMag2(amp00) + cplxMag2(amp01)); - summed_probs[1] += (cplxMag2(amp11) + cplxMag2(amp10)); - summed_probs[2] += (cplxMag2(amp00) + cplxMag2(amp11)); - summed_probs[3] += (cplxMag2(amp01) + cplxMag2(amp10)); - } - default { - // Assume OPID_SHOT_BUFF_2Q - // Get the state vector entries - let states = array( - stateVector[params.shot_state_vector_start + offset00], - stateVector[params.shot_state_vector_start + offset01], - stateVector[params.shot_state_vector_start + offset10], - stateVector[params.shot_state_vector_start + offset11] - ); - // Apply the unitary from the shot buffer - let result00 = innerProduct(getUnitaryRow(params.shot_idx, 0), states); - let result01 = innerProduct(getUnitaryRow(params.shot_idx, 1), states); - let result10 = innerProduct(getUnitaryRow(params.shot_idx, 2), states); - let result11 = innerProduct(getUnitaryRow(params.shot_idx, 3), states); - // Write back the results - stateVector[params.shot_state_vector_start + offset00] = result00; - stateVector[params.shot_state_vector_start + offset01] = result01; - stateVector[params.shot_state_vector_start + offset10] = result10; - stateVector[params.shot_state_vector_start + offset11] = result11; - // Update the probabilities for the acted on qubits - summed_probs[0] += (cplxMag2(result00) + cplxMag2(result01)); - summed_probs[1] += (cplxMag2(result10) + cplxMag2(result11)); - summed_probs[2] += (cplxMag2(result00) + cplxMag2(result10)); - summed_probs[3] += (cplxMag2(result01) + cplxMag2(result11)); - } - } - } - - entry_index += params.total_threads_per_shot; - } - - // Update this thread's totals for the two qubits in the workgroup storage - if (update_probs) { - // Update all for other 2-qubit gates - qubitProbabilities[tid].zero[q1] = summed_probs[0]; - qubitProbabilities[tid].one[q1] = summed_probs[1]; - qubitProbabilities[tid].zero[q2] = summed_probs[2]; - qubitProbabilities[tid].one[q2] = summed_probs[3]; - } -} - -fn apply_correlated_noise(workgroupId: u32, tid: u32) { - let params = get_shot_params(workgroupId, tid, 0 /* need to walk all entries */); - // Probabilities are already updated in the prepare_op stage - // Here we just need to apply the bit-flips and phase-flips to the state vector amplitudes - - let shot = &shots[params.shot_idx]; - - // Get the bit-flip and phase-flip masks from the shot buffer (stored by prep_correlated_noise) - let bit_flip_mask = bitcast(shot.unitary[0].x); - let phase_flip_mask = bitcast(shot.unitary[0].y); - - // If no flips to apply, early exit - if (bit_flip_mask == 0u && phase_flip_mask == 0u) { - return; - } - - var entry_index = params.thread_idx_in_shot; - - for (var i = 0; i < params.op_iterations; i++) { - // Get the target index to swap the state with by flipping the bits as indicated in the bit_flip_mask - let target_index = entry_index ^ i32(bit_flip_mask); - - // If there are an odd number of phase flips for the entry, we need to negate the amplitude - let negate_index: f32 = select(1.0, -1.0, (countOneBits(entry_index & i32(phase_flip_mask)) & 1) != 0); - - if (bit_flip_mask == 0u && negate_index == -1.0) { - // No bit flips to perform, but need to negate this entry (phase flip only) - stateVector[params.shot_state_vector_start + entry_index] = cplxNeg(stateVector[params.shot_state_vector_start + entry_index]); - } else if (entry_index < target_index) { - // Bit flips are happening (as the indices are different), but to avoid double swapping only handle the swap - // when entry_index < target_index (avoid reprocessing when later we encounter the target_index entry as the entry_index) - - let amp_entry: vec2f = stateVector[params.shot_state_vector_start + entry_index]; - let amp_target: vec2f = stateVector[params.shot_state_vector_start + target_index]; - - // If there are an odd number of phase flips for the target, we need to negate that amplitude too - let negate_target: f32 = select(1.0, -1.0, (countOneBits(target_index & i32(phase_flip_mask)) & 1) != 0); - - // Swap and apply any negations for phase flips. - // Note this only applies -1 & 1 to the phase, not -i and i as the 'canonical' Y gate does. - // However, this is sufficient for simulating noise, as the global phase doesn't matter. - stateVector[params.shot_state_vector_start + entry_index] = cplxMul(amp_target, vec2f(negate_index, 0.0)); - stateVector[params.shot_state_vector_start + target_index] = cplxMul(amp_entry, vec2f(negate_target, 0.0)); - } - - // Jump ahead to the next entry to process - entry_index += params.total_threads_per_shot; - } -} - -// For the state vector index and amplitude probability, update all the qubit probabilities for this thread -fn update_all_qubit_probs(stateVectorIndex: u32, amplitude: vec2f, tid: u32) { - var mask: u32 = 1u; - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - let is_one: bool = (stateVectorIndex & mask) != 0u; - let prob: f32 = cplxMag2(amplitude); - if (is_one) { - qubitProbabilities[tid].one[q] += prob; - } else { - qubitProbabilities[tid].zero[q] += prob; - } - mask = mask << 1u; - } -} - -fn sum_thread_totals_to_shot(q: u32, shot_idx: i32, wkg_collation_idx: i32) { - var total_zero: f32 = 0.0; - var total_one: f32 = 0.0; - for (var j = 0; j < THREADS_PER_WORKGROUP; j++) { - total_zero += qubitProbabilities[j].zero[q]; - total_one += qubitProbabilities[j].one[q]; - } - if (wkg_collation_idx >= 0) { - // Write to the workgroup collation buffer for later summation into the shot state - workgroup_collation.sums[wkg_collation_idx].qubits[q] = vec2f(total_zero, total_one); - } else { - // Single workgroup per shot case - write directly to the shot state - let within_threshold = abs(1.0 - (total_zero + total_one)) < PROB_THRESHOLD; - if !within_threshold { - // Populate the diagnostics buffer, if not already set - let old_value = atomicCompareExchangeWeak( - &diagnostics.error_code, - 0u, - ERR_INVALID_THREAD_TOTAL); - if old_value.exchanged { - // This is the first error - fill in the details - let shot = &shots[shot_idx]; - diagnostics.extra1 = q; - diagnostics.extra2 = total_zero; - diagnostics.extra3 = total_one; - // DX12 backend has issues copying structs. See https://github.com/gfx-rs/wgpu/issues/8552 - // DX12-start-strip - diagnostics.shot = *shot; - diagnostics.op = ops[shot.op_idx]; - // DX12-end-strip - } - let err_index = (shot_idx + 1) * i32(RESULT_COUNT) - 1; - atomicCompareExchangeWeak( - &results[err_index], - 0u, - ERR_INVALID_THREAD_TOTAL); - } else { - shots[shot_idx].qubit_state[q].zero_probability = total_zero; - shots[shot_idx].qubit_state[q].one_probability = total_one; - } - } -} - -// Complex number utilities - -// Get the magnitude squared of a complex number -fn cplxMag2(a: vec2f) -> f32 { - return (a.x * a.x + a.y * a.y); -} - -// Complex multiplication -fn cplxMul(a: vec2f, b: vec2f) -> vec2f { - return vec2f( - a.x * b.x - a.y * b.y, - a.x * b.y + a.y * b.x - ); -} - -// Complex negation -fn cplxNeg(a: vec2f) -> vec2f { - return vec2f(-a.x, -a.y); -} - -// Negate all elements in a 4-element row of complex numbers -fn rowNeg(a: array) -> array { - return array( - cplxNeg(a[0]), - cplxNeg(a[1]), - cplxNeg(a[2]), - cplxNeg(a[3])); -} - -// Compute the inner product of two 4-element rows of complex numbers -fn innerProduct(a: array, b: array) -> vec2f { - var result: vec2f = vec2f(0.0, 0.0); - for (var i: u32 = 0u; i < 4u; i++) { - result += cplxMul(a[i], b[i]); - } - return result; -} - -fn getOpRow(op_idx: u32, row: u32) -> array { - let op = &ops[op_idx]; - return array( - op.unitary[row * 4 + 0], - op.unitary[row * 4 + 1], - op.unitary[row * 4 + 2], - op.unitary[row * 4 + 3]); -} - -fn getUnitaryRow(shot_idx: i32, row: u32) -> array { - let shot = &shots[shot_idx]; - return array( - shot.unitary[row * 4 + 0], - shot.unitary[row * 4 + 1], - shot.unitary[row * 4 + 2], - shot.unitary[row * 4 + 3]); -} - -fn setUnitaryRow(shot_idx: u32, row: u32, newRow: array) { - let shot = &shots[shot_idx]; - shot.unitary[row * 4 + 0] = newRow[0]; - shot.unitary[row * 4 + 1] = newRow[1]; - shot.unitary[row * 4 + 2] = newRow[2]; - shot.unitary[row * 4 + 3] = newRow[3]; -} - -// Result of sampling which correlated noise entry (if any) to apply. -struct CorrelatedNoiseSample { - should_apply: u32, // 0 = no noise, 1 = apply noise - paulis_lo: u32, - paulis_hi: u32, -} - -// Samples the correlated noise table to determine whether noise should be applied, and if so, -// which Pauli string was selected. If no noise is applied, the shot is set to ID and the caller -// can return early. -fn sample_correlated_noise(shot_idx: u32, op_idx: u32, noise_table_idx: u32) -> CorrelatedNoiseSample { - let shot = &shots[shot_idx]; - let table = &batch_data.correlated_noise_tables[noise_table_idx]; - - // Generate a Q1.63 random number (two u32 values for lo and hi 32 bits) - // Mask off the high bit of rand_hi to ensure the value is in [0, 1) range - let rand_lo = next_rand_u32(shot_idx); - let rand_hi = next_rand_u32(shot_idx) & 0x7FFFFFFFu; - - // Get the total noise probability from the table metadata - let noise_prob_lo = table.noise_probability_lo; - let noise_prob_hi = table.noise_probability_hi; - - // Check if noise should be applied at all by comparing the random number against the total noise probability - // If rand >= noise_probability, then no noise is applied - if (rand_hi > noise_prob_hi || (rand_hi == noise_prob_hi && rand_lo >= noise_prob_lo)) { - // No noise to apply - set the op to ID - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - shot.qubits_updated_last_op_mask = 0u; - return CorrelatedNoiseSample(0u, 0u, 0u); - } - - // Noise should be applied - binary search to find which Pauli string to apply - let start = i32(table.start_offset); - let count = i32(table.entry_count); - let entry_idx = binary_search_noise_table(rand_lo, rand_hi, start, count); - let entry = &batch_data.correlated_noise_entries[start + entry_idx]; - - return CorrelatedNoiseSample(1u, entry.paulis_lo, entry.paulis_hi); -} - -// Extracts the 3-bit term value for qubit position `i` from a Pauli + loss string. -// Terms use the encoding I=0, X=1, Z=2, Y=3, L=4. The low two bits double as the -// bit-flip (0x1) and phase-flip (0x2) indicators, and 0x4 marks loss. -// The Rust parsing stores terms with the rightmost (last) character at the lowest -// bits, so for position i we read the 3 bits at (qubit_count - 1 - i) * 3. -fn get_pauli_bits(paulis_lo: u32, paulis_hi: u32, qubit_count: u32, i: u32) -> u32 { - let bit_position = (qubit_count - 1u - i) * 3u; - if (bit_position + 3u <= 32u) { - return (paulis_lo >> bit_position) & 0x7u; - } else if (bit_position >= 32u) { - return (paulis_hi >> (bit_position - 32u)) & 0x7u; - } else { - // The 3-bit term straddles the boundary between the lo and hi words. - let low_part = paulis_lo >> bit_position; - let high_part = paulis_hi << (32u - bit_position); - return (low_part | high_part) & 0x7u; - } -} - -// Commits correlated noise masks into the shot state: stores the masks, swaps probabilities and -// tracking bits for bit-flipped qubits, records any loss, and sets the shot up for the correlated -// noise execute stage. Qubits in `loss_mask` are scheduled for loss; following loss-commit ops -// perform the measure + reset. -fn commit_correlated_noise(shot_idx: u32, op_idx: u32, bit_flip_mask: u32, phase_flip_mask: u32, loss_mask: u32) { - let shot = &shots[shot_idx]; - - // Schedule loss for any qubit whose sampled term was loss. The actual - // measure + reset is performed by the loss-commit ops emitted after the - // correlated-noise op. - shot.pending_loss_mask |= loss_mask; - - // Store the masks in the shot buffer for the execute stage - // We use the unitary entries to store these masks (reinterpreted as floats) - shot.unitary[0] = vec2f(bitcast(bit_flip_mask), bitcast(phase_flip_mask)); - - // For bit-flipped qubits, we need to swap the 0 and 1 probabilities and masks - // This is done in prepare_op, not execute_op, since it's a simple swap - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - let qubit_mask = 1u << q; - if ((bit_flip_mask & qubit_mask) != 0u) { - // Swap the probabilities - let temp = shot.qubit_state[q].zero_probability; - shot.qubit_state[q].zero_probability = shot.qubit_state[q].one_probability; - shot.qubit_state[q].one_probability = temp; - - // Swap the bits in qubit_is_0_mask and qubit_is_1_mask - let was_0 = (shot.qubit_is_0_mask & qubit_mask) != 0u; - let was_1 = (shot.qubit_is_1_mask & qubit_mask) != 0u; - if (was_0) { - shot.qubit_is_0_mask &= ~qubit_mask; - shot.qubit_is_1_mask |= qubit_mask; - } else if (was_1) { - shot.qubit_is_1_mask &= ~qubit_mask; - shot.qubit_is_0_mask |= qubit_mask; - } - } - } - - // Set up the shot state for the correlated noise execution - shot.op_type = OPID_CORRELATED_NOISE; - shot.op_idx = op_idx; - // No probabilities need to be recomputed in execute_op since we've already swapped them here - shot.qubits_updated_last_op_mask = 0u; -} - -// Performas a binary search on a correlated noise probability table -// -// Preconditions: -// - table is sorted ascending, with every entry higher than the prior -// - table entries are cumulative probabilities totaling <= 1.0 -// - 'start' is the offset into the buffer array where this table's entries begin -// - 'count' is the number of entries in this table -// - 'rand_lo' and 'rand_hi' form a Q1.63 format random number in [0.0, 1.0) to use for the search -// - This will only called if a result should be found, i.e., -// - count > 0 -// - rand < table[start + count - 1].probability -// -// Returns the index of the found entry relative to 'start', which is the smallest index where "rand < table[start + index].probability" -fn binary_search_noise_table(rand_lo: u32, rand_hi: u32, start: i32, count: i32) -> i32 { - var low: i32 = 0; - var high: i32 = count; - - while (low < high) { - let mid: i32 = low + (high - low) / 2; - let p_lo = batch_data.correlated_noise_entries[start + mid].probability_lo; - let p_hi = batch_data.correlated_noise_entries[start + mid].probability_hi; - - if (rand_hi < p_hi || (rand_hi == p_hi && rand_lo < p_lo)) { - high = mid; - } else { - low = mid + 1; - } - } - return low; -} - -// Hash and random number generation functions - -// See https://www.reedbeta.com/blog/hash-functions-for-gpu-rendering/ -// Use PCG hash function to generate a well-distributed hash from a simple integer input (e.g., shot id) -fn hash_pcg(input: u32) -> u32 { - var state = input * 747796405u + 2891336453u; - var word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u; - return (word >> 22u) ^ word; -} - -// Returns a random u32 value based on the xorwow algorithm -fn next_rand_u32(shot_idx: u32) -> u32 { - // Based on https://en.wikipedia.org/wiki/Xorshift - let rng_state = &shots[shot_idx].rng_state; - - var t: u32 = rng_state.x[4]; - let s: u32 = rng_state.x[0]; - rng_state.x[4] = rng_state.x[3]; - rng_state.x[3] = rng_state.x[2]; - rng_state.x[2] = rng_state.x[1]; - rng_state.x[1] = s; - - t = t ^ (t >> 2u); - t = t ^ (t << 1u); - t = t ^ s ^ (s << 4u); - rng_state.x[0] = t; - rng_state.counter = rng_state.counter + 362437u; - return t + rng_state.counter; -} - -fn next_rand_f32(shot_idx: u32) -> f32 { - let rand_u32: u32 = next_rand_u32(shot_idx); - - // Convert the 32 random bits to a float in the [0.0, 1.0) range - - // Keep only the lower 23 bits (the fraction portion of a float) with a 0 exponent biased to 127 - let rand_f32_bits = (rand_u32 & 0x7FFFFF) | (127 << 23); - // Bitcast to an f32 in the [1.0, 2.0) range - let f: f32 = bitcast(rand_f32_bits); - // And decrement by 1 to return values from [0..1) - return f - 1.0; -} diff --git a/source/simulators/src/gpu_full_state_simulator/simulator_adaptive.wgsl b/source/simulators/src/gpu_full_state_simulator/simulator_adaptive.wgsl deleted file mode 100644 index c9d8331f93c..00000000000 --- a/source/simulators/src/gpu_full_state_simulator/simulator_adaptive.wgsl +++ /dev/null @@ -1,1772 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -// common.wgsl is appended to the beginning of this file at runtime. - -const ERR_CALL_STACK_OVERFLOW = 3u; -const ERR_CALL_STACK_UNDERFLOW = 4u; -const ERR_INVALID_INSTRUCTION = 5u; -const ERR_ALLOCA_OUT_OF_BOUNDS = 6u; -const ERR_MEMORY_OUT_OF_BOUNDS = 7u; - -@group(0) @binding(0) -var workgroup_collation: WorkgroupCollationBuffer; -// Around 128 max partitions times 27 qubits times 8 bytes = 27 KB max size - -struct QubitState { - zero_probability: f32, - one_probability: f32, - heat: f32, // -1.0 = lost - idle_since: f32, -} - -// Used to track state for the random number generator per shot. See `next_rand_f32` later for details. -struct xorwow_state { - counter: u32, - x: array -} - -// Buffer containing the state for each shot to execute per kernel dispatch -// An instance of this is tracked on the GPU for every active shot -struct ShotData { - shot_id: u32, - next_op_idx: u32, - - // The below random numbers will be initialized from the RNG per operation in the 'prepare_op' stage - // Then the 'execute_op' stage will read these precomputed random numbers for noise modeling - rng_state: xorwow_state, // 6 x u32 - rand_pauli: f32, - rand_damping: f32, - rand_dephase: f32, - rand_measure: f32, - // Bitmask of qubits the most recent noise sampler chose to lose. - pending_loss_mask: u32, - - // The type of the next operation to execute. This will be OPID_SHOT_BUFF_* if it should use the unitary from the op buffer - op_type: u32, - op_idx: u32, - - duration: f32, // Total duration of the shot so far, used for time-dependent noise modeling and shot estimations - renormalize: f32, // Value to renormalize the state vector by on next execute (1.0 = no renormalization needed) - - // For quick testing during execution to enable skipping blocks of entries - // TODO: Actually use these masks during execution to skip unneeded work - qubit_is_0_mask: u32, // Bitmask for which qubits are currently in |0> state - qubit_is_1_mask: u32, // Bitmask for which qubits are currently in |1> state - - // Track which qubit probabilities were updated in the last operation (to collate on next prepare_op) - qubits_updated_last_op_mask: u32, - // 20 x 4 bytes to this point = 80 bytes - - // Track the per-qubit probabilities for optimization of measurement sampling and noise modeling - qubit_state: array, // 27 x 16 bytes = 432 bytes - // 512 bytes to this point - - // Map this to the Op structure for ease of use - unitary: array, // For MAT1Q and MAT2Q ops. - - // Adaptive interpreter state (embedded to reduce storage buffer count). - // This is initialized by the host after the GPU init kernel runs. - interp: InterpreterState, -} -// See https://www.w3.org/TR/WGSL/#structure-member-layout for alignment rules - -@group(0) @binding(1) -var shots: array; - -// Buffer containing the list of operations (gates and noise) that make up the program to simulate -struct Op { - id: u32, - q1: u32, - q2: u32, - q3: u32, - policy: u32, - pad0: u32, - pad1: u32, - pad2: u32, - // Entries in the unitary are: 00, 01, 02, 03, 10, 11, 12, 13, 20, ..., 32, 33 - // 1q matrix elements are stored in: 00, 01, 10, 11 (i.e., indices 0, 1, 4, and 5) - unitary: array, -} // Struct size: 4 * 8 + 16 * 8 = 160 bytes (which is aligned to 16 bytes) - -@group(0) @binding(2) -var ops: array; - -// The one large buffer of state vector amplitudes. (Partitioned into multiple shots) -@group(0) @binding(3) -var stateVector: array; - -// Buffer for storing measurement results per shot -@group(0) @binding(4) -var results: array>; - -// When an error occurs, the below diagnostic data structure is used to store information about the error -struct DiagnosticData { - error_code: atomic, - termination_count: atomic, - extra1: u32, - extra2: f32, - extra3: f32, - _padding: u32, - shot: ShotData, // 640 bytes - op: Op, // 144 bytes - // Below is usually 6,912 bytes (size = THREADS_PER_WORKGROUP (32) * (8 * MAX_QUBIT_COUNT (27)) - workgroup_probabilities: array, - // Below is usually 27,648 bytes (1 << u32(MAX_QUBIT_COUNT - MAX_QUBITS_PER_WORKGROUP)) * (8 * MAX_QUBIT_COUNT) bytes - collation_buffer: WorkgroupCollationBuffer, -}; - -@group(0) @binding(5) -var diagnostics: DiagnosticData; - -struct Uniforms { - batch_start_shot_id: i32, - rng_seed: u32, -} - -@group(0) @binding(6) -var uniforms: Uniforms; - -struct NoiseTableMetadata { - /// The total probability of any noise (i.e. sum of all noise entries) in `Q1.63` format - noise_probability_lo: u32, - noise_probability_hi: u32, - /// The start offset of this table's entries in the global `NoiseTableEntry` array - start_offset: u32, - /// The number of entries in this noise table - entry_count: u32, -} - -struct NoiseTableEntry { - /// The correlated pauli string as bits (2 bits per qubit). If bit 0 is set, then it has bit-flip - /// noise, and if bit 1 is set then it has phase-flip noise. e.g., `110001 == "YIX"` - paulis_lo: u32, - paulis_hi: u32, - /// The probability of the noise occurring in `Q1_63` format. This is a float format where the high - /// order bit (bit 63) has the value 1.0 (`2^0 / 1`), bit 62 has the value 0.5 (`2^1 / 1`), etc. - /// all the way to bit 63 with a value of approx 1.0842e-19 (`2^63 / 1`). This gives a range of - /// values from [0..2) with equal spacing of 1.0842e-19 between values (unlike float or double), - /// which makes it more suitable for random numbers used to select between a large number of small - /// probability entries. - probability_lo: u32, - probability_hi: u32, -} - -// Template constants for noise table sizes (must be ≥ 1; host uses max(count,1)). -const NOISE_TABLE_COUNT: u32 = {{NOISE_TABLE_COUNT}}; -const NOISE_ENTRY_COUNT: u32 = {{NOISE_ENTRY_COUNT}}; - -// BatchData holds all the read-only data shared across all shots in a batch. -struct BatchData { - correlated_noise_tables: array, - correlated_noise_entries: array, - program: Program, -} - -@group(0) @binding(7) -var batch_data: BatchData; - - -/// GPU bytecode instruction. -/// -/// Layout: -/// - `opcode`: packed word — bits\[7:0\]=primary, bits\[15:8\]=sub/condition, bits\[23:16\]=flags -/// - `dst`: destination register or branch target -/// - `src0`, `src1`: source registers or immediates -/// - `aux0`-`aux3`: auxiliary fields (gate index, block ids, side-table offsets, etc.) -struct Instruction { - opcode: u32, - dst: u32, - src0: u32, - src1: u32, - aux0: u32, - aux1: u32, - aux2: u32, - aux3: u32, -} - -struct Block { - instr_offset: u32, - instr_count: u32, -} - -struct Function { - entry_block_id: u32, - param_count: u32, - param_base_reg: u32, - reserved: u32, -} - -struct PhiNodeEntry { - block_id: u32, - val_reg: u32, -} - -struct SwitchCase { - case_val: u32, - target_block: u32, -} - -const INSTRUCTIONS_SIZE: u32 = {{INSTRUCTIONS_SIZE}}; -const BLOCK_TABLE_SIZE: u32 = {{BLOCK_TABLE_SIZE}}; -const FUNCTION_TABLE_SIZE: u32 = {{FUNCTION_TABLE_SIZE}}; -const PHI_TABLE_SIZE: u32 = {{PHI_TABLE_SIZE}}; -const SWITCH_CASES_SIZE: u32 = {{SWITCH_CASES_SIZE}}; -const CALL_ARGS_SIZE: u32 = {{CALL_ARGS_SIZE}}; -const CONSTANT_DATA_SIZE: u32 = {{CONSTANT_DATA_SIZE}}; - -struct Program { - /// Bytecode instructions. - instructions: array, - /// Block table: indexed by block ID. - block_table: array, - /// Function table. - function_table: array, - /// Phi entries table: `[predecessor_block_id, value_register]` entries. - phi_table: array, - /// Switch cases table: `[match_value, target_block]` entries. - switch_table: array, - /// Call argument register indices. - call_arg_table: array, - /// Constant data pool (flattened array constant values). - constant_data: array, -} - -struct CallStackFrame { - /// Resume on this block on return. - block_id: u32, - /// Instruction after the call. - return_pc: u32, - /// Where to write the return value. - return_reg: u32, - /// This is for alignment. - reserved: u32, -} - -// MAX_REGISTERS must be declared before InterpreterState which uses it. -const MAX_REGISTERS: u32 = {{MAX_REGISTERS}}; -const MAX_MEMORY: u32 = {{MAX_MEMORY}}; - -/// Per-shot interpreter state. -struct InterpreterState { - /// Instruction index (absolute), PC stands for Program Counter. - pc: u32, - /// Current block ID. - current_block_id: u32, - ///Previous block ID (for phi resolution). - previous_block_id: u32, - /// 0=running, 1=quantum_pending, 2=terminated, 3=error, 4=yield. - status: u32, - /// Quantum op table index. - pending_op_idx: u32, - /// 0=gate, 1=measure, 2=reset. - pending_op_type: u32, - /// From ret instruction - exit_code: u32, - /// Call stack pointer. - call_sp: u32, - /// Call stack frames (4 u32 per frame × 14 frames = 56). - call_stack_frames: array, - /// Per-shot register file. - registers: array, - /// Per-shot memory (constant_data + alloca'd values). - memory: array, -} - -// ----------------------------------------------------------------------------- -// Adaptive interpreter buffer bindings -// Termination counting is done via diagnostics.termination_count (binding 5). -// Interpreter state and registers are embedded in ShotData (binding 1). -// The program, noise tables, and noise entries are in batch_data (binding 7). -// ----------------------------------------------------------------------------- - -// ----------------------------------------------------------------------------- -// Adaptive interpreter constants -// ----------------------------------------------------------------------------- - -const MAX_CLASSICAL_STEPS: u32 = 4096u; - -// Status codes -const STATUS_RUNNING: u32 = 0u; -const STATUS_QUANTUM_PENDING: u32 = 1u; -const STATUS_TERMINATED: u32 = 2u; -const STATUS_ERROR: u32 = 3u; -const STATUS_YIELD: u32 = 4u; - -// pending_op_type values: 0 = gate, 1 = measure, 2 = reset, 3 = loss commit. -// A loss-commit pending op carries the lost qubit in pending_op_idx (not an -// ops-pool index) and is produced while draining pending_loss_mask. Its value -// must not collide with the gate/measure/reset types resolved in prepare_op. -const PENDING_OP_LOSS_COMMIT: u32 = 3u; - -// ----------------------------------------------------------------------------- -// Adaptive interpreter — opcodes -// ----------------------------------------------------------------------------- - -// Shared opcode constants for the Adaptive Profile QIR bytecode interpreter. -// -// These constants define the bytecode encoding used by the Python AdaptiveProfilePass -// (emitter). Values must stay in sync with the Python ``_adaptive_opcodes.py`` file. -// -// Opcode word layout:: -// -// bits [7:0] = primary opcode -// bits [15:8] = sub-opcode / condition code -// bits [23:16] = flags -// -// Compose via bitwise OR: ``opcode | (sub << 8) | flag`` -// Example: ``OP_ICMP | (ICMP_SLE << 8) | FLAG_SRC1_IMM`` - -// -- Flags (pre-shifted to bit 16+) ------------------------------------------ -const FLAG_SRC0_IMM: u32 = 1 << 16; // src0 field is an immediate value, not a register -const FLAG_SRC1_IMM: u32 = 1 << 17; // src1 field is an immediate value, not a register -const FLAG_DST_IMM: u32 = 1 << 18; // dst field is an immediate value, not a register -const FLAG_AUX0_IMM: u32 = 1 << 19; // aux0 field is an immediate value, not a register -const FLAG_AUX1_IMM: u32 = 1 << 20; // aux1 field is an immediate value, not a register -const FLAG_AUX2_IMM: u32 = 1 << 21; // aux2 field is an immediate value, not a register -const FLAG_AUX3_IMM: u32 = 1 << 22; // aux3 field is an immediate value, not a register - -// -- Control Flow ------------------------------------------------------------- -const OP_NOP: u32 = 0x00; -const OP_RET: u32 = 0x02; -const OP_JUMP: u32 = 0x04; -const OP_BRANCH: u32 = 0x05; -const OP_SWITCH: u32 = 0x06; -const OP_CALL: u32 = 0x07; -const OP_CALL_RETURN: u32 = 0x08; - -// -- Quantum ------------------------------------------------------------------ -const OP_QUANTUM_GATE: u32 = 0x10; -const OP_MEASURE: u32 = 0x11; -const OP_RESET: u32 = 0x12; -const OP_READ_RESULT: u32 = 0x13; -const OP_RECORD_OUTPUT: u32 = 0x14; -const OP_READ_LOSS: u32 = 0x15; - -// -- Integer Arithmetic ------------------------------------------------------- -const OP_ADD: u32 = 0x20; -const OP_SUB: u32 = 0x21; -const OP_MUL: u32 = 0x22; -const OP_UDIV: u32 = 0x23; -const OP_SDIV: u32 = 0x24; -const OP_UREM: u32 = 0x25; -const OP_SREM: u32 = 0x26; - -// -- Bitwise / Shift --------------------------------------------------------- -const OP_AND: u32 = 0x28; -const OP_OR: u32 = 0x29; -const OP_XOR: u32 = 0x2A; -const OP_SHL: u32 = 0x2B; -const OP_LSHR: u32 = 0x2C; -const OP_ASHR: u32 = 0x2D; - -// -- Comparison --------------------------------------------------------------- -const OP_ICMP: u32 = 0x30; -const OP_FCMP: u32 = 0x31; - -// -- Float Arithmetic --------------------------------------------------------- -const OP_FADD: u32 = 0x38; -const OP_FSUB: u32 = 0x39; -const OP_FMUL: u32 = 0x3A; -const OP_FDIV: u32 = 0x3B; -const OP_FREM: u32 = 0x3C; - -// -- Type Conversion ---------------------------------------------------------- -const OP_ZEXT: u32 = 0x40; -const OP_SEXT: u32 = 0x41; -const OP_TRUNC: u32 = 0x42; -const OP_FPEXT: u32 = 0x43; -const OP_FPTRUNC: u32 = 0x44; -const OP_INTTOPTR: u32 = 0x45; -const OP_FPTOSI: u32 = 0x46; -const OP_SITOFP: u32 = 0x47; -const OP_FPTOUI: u32 = 0x48; -const OP_UITOFP: u32 = 0x49; - -// -- SSA / Data Movement ----------------------------------------------------- -const OP_PHI: u32 = 0x50; -const OP_SELECT: u32 = 0x51; -const OP_MOV: u32 = 0x52; -const OP_CONST: u32 = 0x53; - -// -- Memory Operations -------------------------------------------------------- -const OP_ALLOCA: u32 = 0x60; -const OP_LOAD: u32 = 0x61; -const OP_STORE: u32 = 0x62; -const OP_GEP: u32 = 0x63; - -// -- ICmp condition codes (sub-opcode, placed in bits[15:8] via << 8) --------- -// Reference: https://llvm.org/docs/LangRef.html#icmp-instruction -const ICMP_EQ: u32 = 0; -const ICMP_NE: u32 = 1; -const ICMP_SLT: u32 = 2; -const ICMP_SLE: u32 = 3; -const ICMP_SGT: u32 = 4; -const ICMP_SGE: u32 = 5; -const ICMP_ULT: u32 = 6; -const ICMP_ULE: u32 = 7; -const ICMP_UGT: u32 = 8; -const ICMP_UGE: u32 = 9; - -// -- FCmp condition codes ----------------------------------------------------- -// Reference: https://llvm.org/docs/LangRef.html#fcmp-instruction -const FCMP_FALSE: u32 = 0; -const FCMP_OEQ: u32 = 1; -const FCMP_OGT: u32 = 2; -const FCMP_OGE: u32 = 3; -const FCMP_OLT: u32 = 4; -const FCMP_OLE: u32 = 5; -const FCMP_ONE: u32 = 6; -const FCMP_ORD: u32 = 7; -const FCMP_UNO: u32 = 8; -const FCMP_UEQ: u32 = 9; -const FCMP_UGT: u32 = 10; -const FCMP_UGE: u32 = 11; -const FCMP_ULT: u32 = 12; -const FCMP_ULE: u32 = 13; -const FCMP_UNE: u32 = 14; -const FCMP_TRUE: u32 = 15; - -// -- Sentinel values ---------------------------------------------------------- -const VOID_RETURN: u32 = 0xFFFFFFFF; // Function does not have a return value. - -// ----------------------------------------------------------------------------- -// Adaptive interpreter — register file access -// ----------------------------------------------------------------------------- - -fn read_reg(shot_idx: u32, reg: u32) -> u32 { - return shots[shot_idx].interp.registers[reg]; -} - -fn write_reg(shot_idx: u32, reg: u32, val: u32) { - shots[shot_idx].interp.registers[reg] = val; -} - -fn read_reg_i32(shot_idx: u32, reg: u32) -> i32 { - return bitcast(read_reg(shot_idx, reg)); -} - -fn write_reg_i32(shot_idx: u32, reg: u32, val: i32) { - write_reg(shot_idx, reg, bitcast(val)); -} - -fn read_reg_f32(shot_idx: u32, reg: u32) -> f32 { - return bitcast(read_reg(shot_idx, reg)); -} - -fn write_reg_f32(shot_idx: u32, reg: u32, val: f32) { - write_reg(shot_idx, reg, bitcast(val)); -} - -// ----------------------------------------------------------------------------- -// Adaptive interpreter — instruction fetch and opcode extraction -// ----------------------------------------------------------------------------- - -fn fetch_instr(pc: u32) -> Instruction { - return batch_data.program.instructions[pc]; -} - -fn get_opcode(packed: u32) -> u32 { return packed & 0xFFu; } -fn get_subcond(packed: u32) -> u32 { return (packed >> 8u) & 0xFFu; } -fn get_flags(packed: u32) -> u32 { return (packed >> 16u) & 0xFFu; } -fn is_src0_imm(flags: u32) -> bool { return (flags & 1u) != 0u; } -fn is_src1_imm(flags: u32) -> bool { return (flags & 2u) != 0u; } - -fn resolve_i32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> i32 { - if (flags & (1u << operand_idx)) != 0u { - return bitcast(operand); // immediate - } - return read_reg_i32(shot_idx, operand); // register -} - -fn resolve_u32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> u32 { - if (flags & (1u << operand_idx)) != 0u { - return operand; - } - return read_reg(shot_idx, operand); -} - -fn resolve_f32(shot_idx: u32, operand: u32, flags: u32, operand_idx: u32) -> f32 { - if (flags & (1u << operand_idx)) != 0u { - return bitcast(operand); // immediate (IEEE 754 bit pattern) - } - return read_reg_f32(shot_idx, operand); -} - -// Resolves q1 for the current quantum instruction. -fn resolve_q1(shot_idx: u32) -> u32 { - let state = shots[shot_idx].interp; - let instr = fetch_instr(state.pc - 1); - if (instr.opcode & FLAG_AUX1_IMM) != 0 { - return instr.aux1; - } - return read_reg(shot_idx, instr.aux1); -} - -// Resolves q2 for the current quantum instruction. -fn resolve_q2(shot_idx: u32) -> u32 { - let state = shots[shot_idx].interp; - let instr = fetch_instr(state.pc - 1); - if (instr.opcode & FLAG_AUX2_IMM) != 0 { - return instr.aux2; - } - return read_reg(shot_idx, instr.aux2); -} - -// Resolves the rotation angle for the current quantum instruction. -// The angle is stored in the instruction's src0 field (register or immediate). -fn resolve_gate_angle(shot_idx: u32) -> f32 { - let state = shots[shot_idx].interp; - let instr = fetch_instr(state.pc - 1); - let flags = get_flags(instr.opcode); - return resolve_f32(shot_idx, instr.src0, flags, 0u); -} - -fn get_measure_qubit(shot_idx: u32, op_idx: u32) -> u32 { - return resolve_q1(shot_idx); -} - -fn get_measure_result(shot_idx: u32, op_idx: u32) -> u32 { - return resolve_q2(shot_idx); -} - -// Read a measurement result from the existing results buffer. -// Results are stored as atomic at shot_idx * RESULT_COUNT + result_id. -fn read_measurement_result(shot_idx: u32, result_id: u32) -> bool { - return atomicLoad(&results[shot_idx * RESULT_COUNT + result_id]) == 1u; -} - -// Return true if the id corresponds to a rotation gate. -fn is_rotation_gate(id: u32) -> bool { - return (12 <= id && id <= 14) || (17 <= id && id <= 19); -} - -// Return true if the angle for the current rotation gate is dynamic. -fn is_dynamic_angle(shot_idx: u32) -> bool { - let state = shots[shot_idx].interp; - let instr = fetch_instr(state.pc - 1); - return (instr.opcode | FLAG_SRC0_IMM) != 0; -} - -// For every qubit, each 'execute' kernel thread will update its own workgroup storage location for accumulating probabilities -// The final probabilities will be reduced and written back to the shot state after the parallel execution completes. -struct QubitProbabilityPerThread { - zero: array, - one: array, -}; // size: 216 bytes - -var qubitProbabilities: array; -// Workgroup memory size: THREADS_PER_WORKGROUP (32) * 216 = 6,912 bytes. - -// Commit a sampled qubit loss on an explicitly given qubit (measure + reset to -// |0> and mark the qubit lost). The lost qubit is carried to the execute stage -// in `op_idx`, and `op_type` is set to OPID_LOSS_NOISE so execute applies the -// reset matrix to that explicit qubit. -fn prep_loss_commit(shot_idx: u32, qubit: u32) { - let shot = &shots[shot_idx]; - let result = select(1u, 0u, shot.rand_measure < shot.qubit_state[qubit].zero_probability); - shot.qubit_state[qubit].heat = -1.0; - prep_measure_reset_instrument(shot_idx, qubit, result, true /* resets_to_zero */); - shot.op_idx = qubit; // execute reads the lost qubit from op_idx - shot.op_type = OPID_LOSS_NOISE; -} - -// Prepare correlated noise for the adaptive path. -// Qubit IDs are read from call_arg_table (register indices), following the same -// pattern as OP_CALL argument passing. -fn prep_correlated_noise(shot_idx: u32, op_idx: u32, qubit_count: u32, arg_offset: u32) { - let noise_table_idx = ops[op_idx].q1; - - let sample = sample_correlated_noise(shot_idx, op_idx, noise_table_idx); - if (sample.should_apply == 0u) { return; } - - // Build bit-flip, phase-flip, and loss masks using qubit IDs from registers via call_arg_table - var bit_flip_mask: u32 = 0u; - var phase_flip_mask: u32 = 0u; - var loss_mask: u32 = 0u; - for (var i: u32 = 0u; i < qubit_count; i++) { - let pauli_bits = get_pauli_bits(sample.paulis_lo, sample.paulis_hi, qubit_count, i); - let arg_reg = batch_data.program.call_arg_table[arg_offset + i]; - let qubit_mask = 1u << read_reg(shot_idx, arg_reg); - if ((pauli_bits & 0x4u) != 0u) { - // Loss term (L = 4): the qubit is lost, no Pauli is applied to it. - loss_mask |= qubit_mask; - } else { - if ((pauli_bits & 0x1u) != 0u) { bit_flip_mask |= qubit_mask; } - if ((pauli_bits & 0x2u) != 0u) { phase_flip_mask |= qubit_mask; } - } - } - - commit_correlated_noise(shot_idx, op_idx, bit_flip_mask, phase_flip_mask, loss_mask); -} - -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn initialize( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - // Get the params - let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); - - // We want every thread to zero out its portion of the state vector for the shot - // We also want threads executing in lockstep to update adjacent entries for better memory access patterns - for (var i = 0; i < params.op_iterations; i++) { - let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; - stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); - } - - // NOTE: No need to synchronize here, as each thread is writing to unique locations - if (params.thread_idx_in_shot == 0) { - // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot - stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); - reset_all(params.shot_idx); - - // Zero the results buffer for this shot so stale exit codes from - // prior runs do not leak via atomicCompareExchangeWeak in OP_RET. - let results_base = u32(params.shot_idx) * RESULT_COUNT; - for (var r = 0u; r < RESULT_COUNT; r++) { - atomicStore(&results[results_base + r], 0u); - } - - // Initialize memory from constant_data - for (var m = 0u; m < CONSTANT_DATA_SIZE; m++) { - shots[params.shot_idx].interp.memory[m] = batch_data.program.constant_data[m]; - } - // Zero the alloca region for CPU-GPU parity - for (var m = CONSTANT_DATA_SIZE; m < MAX_MEMORY; m++) { - shots[params.shot_idx].interp.memory[m] = 0u; - } - } -} - -// ----------------------------------------------------------------------------- -// Adaptive interpreter — interpret_classical entry point -// ----------------------------------------------------------------------------- -// -// This is the main classical bytecode interpreter for the GPU-based adaptive -// quantum simulator. It implements a register-based virtual machine that -// executes classical (non-quantum) instructions on the GPU, one thread per -// shot. Each shot has its own independent interpreter state (program counter, -// registers, call stack) allowing many shots to run in parallel with -// potentially divergent control flow paths (e.g., after mid-circuit -// measurements). -// -// ## Execution Model -// -// The interpreter runs cooperatively with the quantum simulation pipeline: -// -// 1. The host dispatches `interpret_classical` for all shots. -// 2. Each shot executes classical instructions in a loop until one of: -// (a) A quantum operation is encountered → status = QUANTUM_PENDING, -// which tells the host to run the quantum simulation kernels -// (prepare_op → execute) before re-entering this function. -// (b) A `ret` instruction terminates the shot → status = TERMINATED. -// (c) The step limit (MAX_CLASSICAL_STEPS) is hit → status = YIELD, -// which prevents any single dispatch from running forever; the host -// simply re-dispatches to continue. -// (d) An unknown opcode is hit → status = ERROR. -// -// ## Instruction Encoding -// -// Each instruction occupies 2 × vec4 (8 u32 words) in the `bytecode` -// buffer, fetched by `fetch_instr(pc)` into the `Instr` struct with fields: -// -// opcode : packed opcode word (bits [7:0] = primary op, [15:8] = sub- -// condition for comparisons, [23:16] = flags for immediates) -// dst : destination register index (or immediate for RET) -// src0 : first source operand (register index or immediate) -// src1 : second source operand (register index or immediate) -// aux0–3 : auxiliary fields whose meaning varies per opcode (e.g., block -// IDs, function IDs, qubit indices, phi-table offsets, etc.) -// -// The `resolve_u32` / `resolve_i32` helpers read an operand as either a -// register value or an inline immediate based on the FLAG_SRC0_IMM / -// FLAG_SRC1_IMM bits in the flags byte. This lets the compiler embed small -// constants directly in the instruction stream without extra CONST ops. - - -@compute @workgroup_size(1) -fn interpret_classical(@builtin(global_invocation_id) gid: vec3) { - // Each GPU thread handles exactly one shot. The global invocation ID - // maps directly to the shot index. - let shot_idx = gid.x; - let state = shots[shot_idx].interp; - - // -- Early-exit for shots that already finished or errored -- - let status = state.status; - if status == STATUS_TERMINATED || status == STATUS_ERROR { - return; - } - - // -- Drain pending qubit losses before resuming classical execution -- - // The most recent noise op (per-gate Pauli/loss or correlated) may have - // sampled one or more qubits as lost, recorded in pending_loss_mask. Commit - // each as its own measure+reset quantum op (one per round) before running - // any more bytecode, so loss is applied with the correct correlation. - if shots[shot_idx].pending_loss_mask != 0u { - let q = firstTrailingBit(shots[shot_idx].pending_loss_mask); - shots[shot_idx].pending_loss_mask &= ~(1u << q); - shots[shot_idx].interp.pending_op_idx = q; - shots[shot_idx].interp.pending_op_type = PENDING_OP_LOSS_COMMIT; - shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; - return; - } - - // If we were paused (QUANTUM_PENDING after a quantum op, or YIELD after - // hitting the step limit), transition back to RUNNING so the main loop - // resumes executing instructions from where it left off. - if status != STATUS_RUNNING { - shots[shot_idx].interp.status = STATUS_RUNNING; - } - - // -- Load interpreter registers from GPU memory into local variables -- - // Using local vars for the hot-path state avoids repeated global memory - // loads/stores on every instruction. They are written back at the end. - var pc: u32 = state.pc; // program counter - var block_id: u32 = state.current_block_id; - var prev_block: u32 = state.previous_block_id; // for PHI - var steps: u32 = 0u; // counts instructions executed this dispatch - var should_break: bool = false; // set to true to exit the main loop - - // -- Main interpreter loop -- - // Fetches and executes one instruction per iteration. Exits when the - // shot terminates, yields for quantum work, hits the step limit, or - // encounters an error. - loop { - // Guard against infinite loops in classical code: after executing - // MAX_CLASSICAL_STEPS instructions, yield back to the host which - // will re-dispatch this kernel to continue. - if steps >= MAX_CLASSICAL_STEPS { - // Only yield if the shot hasn't already errored (an error - // status must not be overwritten by a yield). - if state.status != STATUS_ERROR { - shots[shot_idx].interp.status = STATUS_YIELD; - } - break; - } - - // Fetch the instruction at the current PC. Each instruction is - // 2 × vec4 (8 words) in the bytecode buffer. - let instr = fetch_instr(pc); - - // Unpack the opcode word into its three components: - // op — primary opcode (bits 7:0), determines which case below runs - // subcond — sub-condition code (bits 15:8), used only by ICMP/FCMP to - // select the specific comparison predicate (eq, ne, slt, etc.) - // flags — immediate-mode flags (bits 23:16), tells resolve_* whether - // src0/src1 are register indices or inline immediates - let op = get_opcode(instr.opcode); - let subcond = get_subcond(instr.opcode); - let flags = get_flags(instr.opcode); - - // -- Opcode dispatch -- - // The switch below implements every bytecode instruction. Instructions - // are grouped by category. Most follow a common pattern: - // 1. Read operands via resolve_u32/i32 (register or immediate) - // 2. Compute the result - // 3. Write back to the destination register via write_reg* - // 4. Advance pc++ - // - // Control-flow ops (JUMP, BRANCH, SWITCH, CALL) modify pc and - // block_id directly instead of incrementing pc. - // - // Quantum ops (QUANTUM_GATE, MEASURE, RESET) write pending-op - // metadata to the interpreter state and set should_break=true to - // pause execution and hand control back to the host for quantum - // kernel dispatch. - switch op { - - // ------------------------------------------------------------- - // CONTROL FLOW - // ------------------------------------------------------------- - - // NOP: No operation. Simply advances the program counter. - case OP_NOP { - pc++; - } - - // RET: Terminates this shot's execution. - // The exit code (from dst, which may be an immediate) is stored - // both in the per-shot interpreter state and atomically into the - // results buffer. The atomic-compare-exchange ensures only the - // first non-zero exit code is recorded for this shot (useful for - // error reporting). The termination count in the diagnostics - // buffer is incremented so the host can detect when all shots - // have finished. - case OP_RET { - let exit_code = resolve_u32(shot_idx, instr.dst, flags, 2u); - shots[shot_idx].interp.exit_code = exit_code; - // Atomically store exit code into the last slot of this shot's - // result region, but only if it has not already been set. - let err_index = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_index], 0u, exit_code); - shots[shot_idx].interp.status = STATUS_TERMINATED; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - } - - // JUMP: Unconditional branch to a target block. - // Encoding: dst = target block ID. - // Updates prev_block (needed by subsequent PHI instructions in - // the target block) and sets pc to the first instruction of the - // target block via the block_table lookup. - case OP_JUMP { - prev_block = block_id; - block_id = instr.dst; - pc = batch_data.program.block_table[instr.dst].instr_offset; - } - - // BRANCH: Conditional branch (if/else). - // Encoding: src0 = condition (register or immediate), - // aux0 = true-branch block ID, - // aux1 = false-branch block ID. - // Evaluates the condition: if non-zero, jumps to aux0; otherwise - // jumps to aux1. Like JUMP, updates prev_block for PHI nodes. - case OP_BRANCH { - let cond = resolve_u32(shot_idx, instr.src0, flags, 0u) != 0u; - prev_block = block_id; - if cond { - block_id = instr.aux0; - pc = batch_data.program.block_table[instr.aux0].instr_offset; - } else { - block_id = instr.aux1; - pc = batch_data.program.block_table[instr.aux1].instr_offset; - } - } - - // SWITCH: Multi-way branch (like a C switch statement). - // Encoding: src0 = value to match, - // aux0 = default block ID, - // aux1 = offset into switch_table, - // aux2 = number of case entries. - // Each switch_table entry is a vec2(match_value, target_block). - // Linearly scans the case table; if a match is found, jumps to - // that block. If no match, falls through to the default block. - case OP_SWITCH { - let val = resolve_u32(shot_idx, instr.src0, flags, 0u); - let default_block = instr.aux0; - let case_offset = instr.aux1; - let case_count = instr.aux2; - var target_block = default_block; - for (var i = 0u; i < case_count; i++) { - let entry = batch_data.program.switch_table[case_offset + i]; - if entry.case_val == val { - target_block = entry.target_block; - break; - } - } - prev_block = block_id; - block_id = target_block; - pc = batch_data.program.block_table[target_block].instr_offset; - } - - // CALL: Invokes a function. - // Encoding: dst = register to receive the return value, - // aux0 = function ID (index into function_table), - // aux1 = argument count, - // aux2 = offset into call_arg_table. - // - // The function_table entry is vec4(entry_block, param_count, - // param_base_reg, reserved). - // - // Steps: - // 1. Push a return frame onto the per-shot call stack. Each - // frame stores: (return_block, return_pc, return_reg, - // reserved) — 4 u32 words. The stack supports up to 8 frames. - // 2. Copy each argument from caller registers (looked up via - // call_arg_table) into callee parameter registers starting - // at param_base_reg. - // 3. Jump to the function's entry block. - case OP_CALL { - let func_id = instr.aux0; - let arg_count = instr.aux1; - let arg_offset = instr.aux2; - let func = batch_data.program.function_table[func_id]; - // Push return info onto the call stack - let sp = shots[shot_idx].interp.call_sp; - // Guard: prevent call stack overflow (max 8 frames) - if sp >= 8u { - shots[shot_idx].interp.exit_code = ERR_CALL_STACK_OVERFLOW; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_CALL_STACK_OVERFLOW); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - break; - } - shots[shot_idx].interp.call_stack_frames[sp].block_id = block_id; // return_block — resume here on return - shots[shot_idx].interp.call_stack_frames[sp].return_pc = pc + 1u; // return_pc — instruction after the CALL - shots[shot_idx].interp.call_stack_frames[sp].return_reg = instr.dst; // return_reg — where to write result - shots[shot_idx].interp.call_sp = sp + 1u; - // Copy caller arguments into the callee's parameter registers - let param_base = func.param_base_reg; - for (var i = 0u; i < arg_count; i++) { - let arg_reg = batch_data.program.call_arg_table[arg_offset + i]; - write_reg(shot_idx, param_base + i, read_reg(shot_idx, arg_reg)); - } - // Transfer control to the function entry block - block_id = func.entry_block_id; - pc = batch_data.program.block_table[block_id].instr_offset; - } - - // CALL_RETURN: Returns from a function call. - // Encoding: src0 = register holding the return value. - // - // Pops the top frame from the call stack to restore block_id and - // pc to the instruction after the CALL. If the caller specified a - // return register (not 0xFFFFFFFF), copies the return value into - // that register. - case OP_CALL_RETURN { - if shots[shot_idx].interp.call_sp == 0u { - shots[shot_idx].interp.exit_code = ERR_CALL_STACK_UNDERFLOW; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_CALL_STACK_UNDERFLOW); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - break; - } - - let sp = shots[shot_idx].interp.call_sp - 1; - shots[shot_idx].interp.call_sp = sp; - block_id = shots[shot_idx].interp.call_stack_frames[sp].block_id; - pc = shots[shot_idx].interp.call_stack_frames[sp].return_pc; - let return_reg = shots[shot_idx].interp.call_stack_frames[sp].return_reg; - if return_reg != VOID_RETURN { - write_reg(shot_idx, return_reg, read_reg(shot_idx, instr.src0)); - } - } - - // ------------------------------------------------------------- - // QUANTUM OPERATIONS — pause the interpreter, yield to the host - // ------------------------------------------------------------- - // When the interpreter hits a quantum instruction, it cannot - // execute it directly (quantum simulation runs in separate GPU - // kernels with parallel state-vector processing). Instead, it - // writes the pending operation details into the interpreter - // state for the host to read, sets status = QUANTUM_PENDING, - // advances pc past the instruction, and breaks out of the loop. - // - // The host then dispatches prepare_op (which reads the - // pending op metadata and configures the shot for the quantum - // kernel) followed by the execute kernel (which applies the - // gate/measurement/reset to the state vector). After that, the - // host re-dispatches interpret_classical to continue. - // - // Qubit IDs may be static (embedded in aux1/aux2 by the - // compiler) or dynamic (computed at runtime and stored in - // registers). - - // QUANTUM_GATE: Request a 1- or 2-qubit gate. - // Encoding: aux0 = quantum op table index, - // aux1 = qubit 1 (or register if not sentinel), - // aux2 = qubit 2 (or register if not sentinel). - case OP_QUANTUM_GATE { - shots[shot_idx].interp.pending_op_idx = instr.aux0; - shots[shot_idx].interp.pending_op_type = 0u; // type 0 = gate - // Qubit IDs are resolved in prepare_op via resolve_q1/resolve_q2, - // which use the FLAG_AUX1_IMM / FLAG_AUX2_IMM bits to decide - // between immediate values and register lookups. - shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; - pc++; - should_break = true; - } - - // MEASURE: Request a qubit measurement. - // Encoding: aux0 = quantum op table index, - // aux1 = qubit to measure (or register). - // Only q1 is used; q2 is set to sentinel (unused). - case OP_MEASURE { - shots[shot_idx].interp.pending_op_idx = instr.aux0; - shots[shot_idx].interp.pending_op_type = 1u; // type 1 = gate - // Qubit and result IDs are resolved in prepare_op via - // resolve_q1 (aux1) and resolve_q2 (aux2). - shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; - pc++; - should_break = true; - } - - // RESET: Request a qubit reset (measure + conditional X). - // Encoding: aux0 = quantum op table index, - // aux1 = qubit to reset (or register). - case OP_RESET { - shots[shot_idx].interp.pending_op_idx = instr.aux0; - shots[shot_idx].interp.pending_op_type = 2u; // type 2 = reset - // Qubit ID is resolved in prepare_op via resolve_q1 (aux1). - shots[shot_idx].interp.status = STATUS_QUANTUM_PENDING; - pc++; - should_break = true; - } - - // ------------------------------------------------------------- - // QUANTUM RESULT ACCESS - // ------------------------------------------------------------- - - // READ_RESULT: Load a prior measurement outcome into a register. - // Encoding: src0 = result ID (index into the results buffer), - // dst = destination register. - // The measurement result (0 or 1) was written by an earlier - // MEASURE quantum op. This reads it atomically from the shared - // results buffer and stores 0u or 1u into the destination - // register, allowing classical code to branch on measurement - // outcomes. - case OP_READ_RESULT { - let result_id = instr.src0; - let result_val = read_measurement_result(shot_idx, result_id); - write_reg(shot_idx, instr.dst, select(0u, 1u, result_val)); - pc++; - } - - // RECORD_OUTPUT: Marker for output recording. - // On the GPU this is a no-op — the host reads the results buffer - // directly after all shots terminate. The instruction exists to - // maintain compatibility with the QIR adaptive profile bytecode. - case OP_RECORD_OUTPUT { - pc++; - } - - // READ_LOSS: Reports whether the measurement that produced a - // result observed a lost qubit. The per-shot ``results`` buffer - // encodes loss as the value 2u (0u = Zero, 1u = One, 2u = Loss), - // so we compare against 2u and write 1u when the result was a loss, - // else 0u. - case OP_READ_LOSS { - let result_id = instr.src0; - let val = atomicLoad(&results[shot_idx * RESULT_COUNT + result_id]); - write_reg(shot_idx, instr.dst, select(0u, 1u, val == 2u)); - pc++; - } - - // ------------------------------------------------------------- - // INTEGER ARITHMETIC - // ------------------------------------------------------------- - // All integer arithmetic ops follow the pattern: - // dst = src0 src1 - // Operands are resolved via resolve_i32/u32, which checks the - // FLAG_SRC0_IMM / FLAG_SRC1_IMM bits to determine if the field - // is a register index or an inline immediate constant. - - // ADD: Signed integer addition. dst = src0 + src1. - case OP_ADD { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a + b); - pc++; - } - - // SUB: Signed integer subtraction. dst = src0 - src1. - case OP_SUB { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a - b); - pc++; - } - - // MUL: Signed integer multiplication. dst = src0 * src1. - case OP_MUL { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a * b); - pc++; - } - - // UDIV: Unsigned integer division. dst = src0 / src1. - case OP_UDIV { - let a = resolve_u32(shot_idx, instr.src0, flags, 0u); - let b = resolve_u32(shot_idx, instr.src1, flags, 1u); - write_reg(shot_idx, instr.dst, a / b); - pc++; - } - - // SDIV: Signed integer division (truncates toward zero). dst = src0 / src1. - case OP_SDIV { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a / b); - pc++; - } - - // UREM: Unsigned integer remainder. dst = src0 % src1. - case OP_UREM { - let a = resolve_u32(shot_idx, instr.src0, flags, 0u); - let b = resolve_u32(shot_idx, instr.src1, flags, 1u); - write_reg(shot_idx, instr.dst, a % b); - pc++; - } - - // SREM: Signed integer remainder. - // Computes a - b * trunc(a/b) manually rather than using the % - // operator, because WGSL i32 division truncates toward zero but - // the built-in % may not preserve the sign of the dividend on - // all GPU backends. This matches LLVM's srem semantics. - case OP_SREM { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a - b * (a / b)); - pc++; - } - - // ------------------------------------------------------------- - // BITWISE / SHIFT OPERATIONS - // ------------------------------------------------------------- - // Operate on the raw u32 bit pattern of the register values. - - // AND: Bitwise AND. dst = src0 & src1. - case OP_AND { - write_reg(shot_idx, instr.dst, - resolve_u32(shot_idx, instr.src0, flags, 0u) & resolve_u32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // OR: Bitwise OR. dst = src0 | src1. - case OP_OR { - write_reg(shot_idx, instr.dst, - resolve_u32(shot_idx, instr.src0, flags, 0u) | resolve_u32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // XOR: Bitwise exclusive OR. dst = src0 ^ src1. - case OP_XOR { - write_reg(shot_idx, instr.dst, - resolve_u32(shot_idx, instr.src0, flags, 0u) ^ resolve_u32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // SHL: Logical shift left. dst = src0 << src1. - case OP_SHL { - write_reg(shot_idx, instr.dst, - resolve_u32(shot_idx, instr.src0, flags, 0u) << resolve_u32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // LSHR: Logical shift right (zero-fill). dst = src0 >> src1. - case OP_LSHR { - write_reg(shot_idx, instr.dst, - resolve_u32(shot_idx, instr.src0, flags, 0u) >> resolve_u32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // ASHR: Arithmetic shift right (sign-extending). dst = src0 >> src1. - // Uses i32 to preserve the sign bit during the shift. - case OP_ASHR { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_u32(shot_idx, instr.src1, flags, 1u); - write_reg_i32(shot_idx, instr.dst, a >> b); - pc++; - } - - // ------------------------------------------------------------- - // INTEGER COMPARISON (ICMP) - // ------------------------------------------------------------- - // Compares two integer operands using the sub-condition code - // encoded in bits [15:8] of the opcode word. The result is - // written as 0u (false) or 1u (true) to the destination register. - // Signed comparisons (SLT, SLE, SGT, SGE) use i32 directly; - // unsigned comparisons (ULT, ULE, UGT, UGE) bitcast to u32. - // These mirror LLVM icmp predicates. - case OP_ICMP { - let a = resolve_i32(shot_idx, instr.src0, flags, 0u); - let b = resolve_i32(shot_idx, instr.src1, flags, 1u); - var result: bool = false; - switch subcond { - case ICMP_EQ { result = (a == b); } - case ICMP_NE { result = (a != b); } - case ICMP_SLT { result = (a < b); } - case ICMP_SLE { result = (a <= b); } - case ICMP_SGT { result = (a > b); } - case ICMP_SGE { result = (a >= b); } - case ICMP_ULT { result = (bitcast(a) < bitcast(b)); } - case ICMP_ULE { result = (bitcast(a) <= bitcast(b)); } - case ICMP_UGT { result = (bitcast(a) > bitcast(b)); } - case ICMP_UGE { result = (bitcast(a) >= bitcast(b)); } - default { - shots[shot_idx].interp.status = ERR_INVALID_INSTRUCTION; - shots[shot_idx].interp.exit_code = ERR_INVALID_INSTRUCTION; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_INVALID_INSTRUCTION); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - } - } - write_reg(shot_idx, instr.dst, select(0u, 1u, result)); - pc++; - } - - // ------------------------------------------------------------- - // FLOAT COMPARISON (FCMP) - // ------------------------------------------------------------- - // Compares two f32 operands using the sub-condition code. - // "O" prefix = ordered (both operands are not NaN). The result - // is written as 0u/1u. Mirrors LLVM fcmp ordered predicates. - case OP_FCMP { - let a = resolve_f32(shot_idx, instr.src0, flags, 0u); - let b = resolve_f32(shot_idx, instr.src1, flags, 1u); - var result: bool = false; - switch subcond { - case FCMP_OEQ { result = (a == b); } - case FCMP_ONE { result = (a != b); } - case FCMP_OLT { result = (a < b); } - case FCMP_OLE { result = (a <= b); } - case FCMP_OGT { result = (a > b); } - case FCMP_OGE { result = (a >= b); } - default { - shots[shot_idx].interp.exit_code = ERR_INVALID_INSTRUCTION; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_INVALID_INSTRUCTION); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - } - } - write_reg(shot_idx, instr.dst, select(0u, 1u, result)); - pc++; - } - - // ------------------------------------------------------------- - // FLOAT ARITHMETIC - // ------------------------------------------------------------- - // These operate on f32 values stored in registers via bitcast. - // Operands are always register-based (no immediate flags for - // float ops). - - // FADD: Float addition. dst = src0 + src1. - case OP_FADD { - write_reg_f32(shot_idx, instr.dst, - resolve_f32(shot_idx, instr.src0, flags, 0u) + resolve_f32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // FSUB: Float subtraction. dst = src0 - src1. - case OP_FSUB { - write_reg_f32(shot_idx, instr.dst, - resolve_f32(shot_idx, instr.src0, flags, 0u) - resolve_f32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // FMUL: Float multiplication. dst = src0 * src1. - case OP_FMUL { - write_reg_f32(shot_idx, instr.dst, - resolve_f32(shot_idx, instr.src0, flags, 0u) * resolve_f32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // FDIV: Float division. dst = src0 / src1. - case OP_FDIV { - write_reg_f32(shot_idx, instr.dst, - resolve_f32(shot_idx, instr.src0, flags, 0u) / resolve_f32(shot_idx, instr.src1, flags, 1u)); - pc++; - } - - // FREM: Float remainder. LLVM docs say this instruction has - // the same semantics as C's fmod, which is implemented as: - // dst = src0 - trunc(src0/src1) * src1 - case OP_FREM { - let a = resolve_f32(shot_idx, instr.src0, flags, 0u); - let b = resolve_f32(shot_idx, instr.src1, flags, 1u); - write_reg_f32(shot_idx, instr.dst, a - trunc(a / b) * b); - pc++; - } - - // ------------------------------------------------------------- - // TYPE CONVERSIONS - // ------------------------------------------------------------- - // Maps LLVM-style type conversion instructions. Many are - // identity ops on the GPU since all integer registers are 32-bit - // and all floats are f32. They exist to keep the bytecode in - // 1:1 correspondence with the compiled QIR instructions. - - // ZEXT: Zero-extend — identity on 32-bit GPU (values already u32). - case OP_ZEXT { - write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // SEXT: Sign-extend from a narrower bit width to i32. - // aux0 encodes the source bit width (e.g., 1 for i1→i32). - // The shift-left then arithmetic-shift-right trick propagates - // the sign bit from position (src_bits-1) into all higher bits. - case OP_SEXT { - let val = resolve_i32(shot_idx, instr.src0, flags, 0u); - let src_bits = instr.aux0; // source type bit width - if src_bits > 0u && src_bits < 32u { - let shift = 32u - src_bits; - write_reg_i32(shot_idx, instr.dst, (val << shift) >> shift); - } else { - write_reg_i32(shot_idx, instr.dst, val); - } - pc++; - } - - // TRUNC: Truncate — identity on 32-bit GPU (already the target width). - case OP_TRUNC { - write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // FPEXT: Float widen (e.g., f32→f64) — identity since GPU only uses f32. - case OP_FPEXT { - write_reg_f32(shot_idx, instr.dst, resolve_f32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // FPTRUNC: Float narrow (e.g., f64→f32) — identity since GPU only uses f32. - case OP_FPTRUNC { - write_reg_f32(shot_idx, instr.dst, resolve_f32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // INTTOPTR: Integer to pointer cast — identity, pointers are u32 on GPU. - case OP_INTTOPTR { - write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // FPTOSI: Float to signed integer conversion. dst = i32(src0). - case OP_FPTOSI { - write_reg_i32(shot_idx, instr.dst, i32(resolve_f32(shot_idx, instr.src0, flags, 0u))); - pc++; - } - - // SITOFP: Signed integer to float conversion. dst = f32(src0). - case OP_SITOFP { - write_reg_f32(shot_idx, instr.dst, f32(resolve_i32(shot_idx, instr.src0, flags, 0u))); - pc++; - } - - // FPTOUI: Float to unsigned integer conversion. dst = u32(src0). - case OP_FPTOUI { - write_reg(shot_idx, instr.dst, u32(resolve_f32(shot_idx, instr.src0, flags, 0u))); - pc++; - } - - // UITOFP: Unsigned integer to float conversion. dst = f32(src0). - case OP_UITOFP { - write_reg_f32(shot_idx, instr.dst, f32(resolve_u32(shot_idx, instr.src0, flags, 0u))); - pc++; - } - - // ------------------------------------------------------------- - // PHI NODE (SSA resolution at runtime) - // ------------------------------------------------------------- - // In SSA form, PHI nodes select a value based on which - // predecessor block the control flow came from. The compiler - // emits a phi_table with (predecessor_block_id, value_register) - // pairs for each PHI instruction. - // - // Encoding: dst = destination register, - // aux0 = offset into phi_table, - // aux1 = number of predecessor entries. - // - // At runtime, we scan the entries to find the one whose block - // ID matches prev_block, then copy that register's value into - // the destination. This is how the interpreter handles SSA - // control-flow merges without explicit move instructions on - // every edge. - case OP_PHI { - let offset = instr.aux0; - let count = instr.aux1; - for (var i = 0u; i < count; i++) { - let entry = batch_data.program.phi_table[offset + i]; - if entry.block_id == prev_block { - write_reg(shot_idx, instr.dst, read_reg(shot_idx, entry.val_reg)); - break; - } - } - pc++; - } - - // ------------------------------------------------------------- - // DATA MOVEMENT - // ------------------------------------------------------------- - - // SELECT: Conditional move (ternary operator). - // Encoding: src0 = condition, aux0 = true-value, - // aux1 = false-value, dst = destination. - // dst = cond ? aux0 : aux1 - case OP_SELECT { - let cond = resolve_u32(shot_idx, instr.src0, flags, 0u) != 0u; - let true_val = resolve_u32(shot_idx, instr.aux0, flags, 3u); - let false_val = resolve_u32(shot_idx, instr.aux1, flags, 4u); - write_reg(shot_idx, instr.dst, select(false_val, true_val, cond)); - pc++; - } - - // MOV: Register-to-register move (or immediate-to-register if flagged). - // dst = src0 (resolved through flags for possible immediate). - case OP_MOV { - write_reg(shot_idx, instr.dst, resolve_u32(shot_idx, instr.src0, flags, 0u)); - pc++; - } - - // CONST: Load an immediate constant into a register. - // dst = src0 (always treated as a literal value, not a register). - case OP_CONST { - write_reg(shot_idx, instr.dst, instr.src0); - pc++; - } - - // ------------------------------------------------------------- - // MEMORY OPERATIONS - // ------------------------------------------------------------- - - // ALLOCA: Reserve memory and write the address to dst. - // Encoding: src0 = number of words, src1 = compile-time assigned address. - case OP_ALLOCA { - let num_words = resolve_u32(shot_idx, instr.src0, flags, 0u); - let addr = resolve_u32(shot_idx, instr.src1, flags, 1u); - if addr + num_words > MAX_MEMORY { - shots[shot_idx].interp.exit_code = ERR_ALLOCA_OUT_OF_BOUNDS; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_ALLOCA_OUT_OF_BOUNDS); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - break; - } - write_reg(shot_idx, instr.dst, addr); - pc++; - } - - // LOAD: Read a value from memory at the given address. - // Encoding: src0 = memory address, dst = destination register. - case OP_LOAD { - let addr = resolve_u32(shot_idx, instr.src0, flags, 0u); - if addr >= MAX_MEMORY { - shots[shot_idx].interp.exit_code = ERR_MEMORY_OUT_OF_BOUNDS; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_MEMORY_OUT_OF_BOUNDS); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - break; - } - let val = shots[shot_idx].interp.memory[addr]; - write_reg(shot_idx, instr.dst, val); - pc++; - } - - // STORE: Write a value to memory at the given address. - // Encoding: src0 = value to store, src1 = memory address. - case OP_STORE { - let val = resolve_u32(shot_idx, instr.src0, flags, 0u); - let addr = resolve_u32(shot_idx, instr.src1, flags, 1u); - if addr >= MAX_MEMORY { - shots[shot_idx].interp.exit_code = ERR_MEMORY_OUT_OF_BOUNDS; - let err_idx = (shot_idx + 1) * RESULT_COUNT - 1; - atomicCompareExchangeWeak(&results[err_idx], 0u, ERR_MEMORY_OUT_OF_BOUNDS); - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - break; - } - shots[shot_idx].interp.memory[addr] = val; - pc++; - } - - // GEP: Get element pointer — compute address from base + index * elem_size. - // Encoding: src0 = base address, src1 = index, aux0 = element size. - case OP_GEP { - let base = resolve_u32(shot_idx, instr.src0, flags, 0u); - let index = resolve_u32(shot_idx, instr.src1, flags, 1u); - let elem_size = resolve_u32(shot_idx, instr.aux0, flags, 3u); - let addr = base + index * elem_size; - write_reg(shot_idx, instr.dst, addr); - pc++; - } - - // Unknown opcode — flag the shot as errored. - default { - shots[shot_idx].interp.status = STATUS_ERROR; - atomicAdd(&diagnostics.termination_count, 1u); - should_break = true; - } - } - steps++; - if should_break { break; } - } - - // -- Persist interpreter state back to GPU memory -- - // Write the local variables back so the next dispatch (after quantum ops - // or a yield) can resume exactly where this invocation left off. - shots[shot_idx].interp.pc = pc; - shots[shot_idx].interp.current_block_id = block_id; - shots[shot_idx].interp.previous_block_id = prev_block; -} - -// ----------------------------------------------------------------------------- -// Adaptive interpreter — prepare_op entry point -// ----------------------------------------------------------------------------- -// Prepares a quantum operation for shots that have STATUS_QUANTUM_PENDING. -// Shots not in that state are set to OPID_ID so execute is a no-op. - -@compute @workgroup_size(1) -fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { - let shot_idx = globalId.x; - let shot = &shots[shot_idx]; - let state = shots[shot_idx].interp; - let status = state.status; - - // Only process shots that are quantum-pending - if status != STATUS_QUANTUM_PENDING { - // Set op_type to ID so execute is a no-op for this shot - shot.op_type = OPID_ID; - shot.renormalize = 1.0; - shot.qubits_updated_last_op_mask = 0u; - return; - } - - // Update shot state from prior op execution - if shot.qubits_updated_last_op_mask != 0 { - update_qubit_state(shot_idx); - } - shot_init_per_op(shot_idx); - - let op_idx = state.pending_op_idx; - let op_type = state.pending_op_type; - - // Loss commit: pending_op_idx holds the lost qubit (not an ops-pool index). - // Measure + reset that qubit; the execute stage applies it via op_idx. - if op_type == PENDING_OP_LOSS_COMMIT { - prep_loss_commit(shot_idx, op_idx); - return; - } - - let op = &ops[op_idx]; - - // Correlated noise: qubit IDs are stored as register indices in - // call_arg_table; read aux1 (qubit count) and aux2 (arg offset) - // from the instruction that triggered this quantum op. - if op_type == 0u && op.id == OPID_CORRELATED_NOISE { - let pc = state.pc; - let noise_instr = fetch_instr(pc - 1u); - let qubit_count = noise_instr.aux1; - let arg_offset = noise_instr.aux2; - shot.op_idx = op_idx; - shot.op_type = op.id; - prep_correlated_noise(shot_idx, op_idx, qubit_count, arg_offset); - shots[shot_idx].interp.status = STATUS_RUNNING; - return; - } - - let q1 = resolve_q1(shot_idx); - let q2 = resolve_q2(shot_idx); - - shot.unitary = op.unitary; - - switch op_type { - case 0u { // Gate - // For rotation gates, recompute the unitary from the dynamic angle stored - // in the instruction's src0 field if needed. The op pool unitary was built - // at upload time and may not reflect a runtime-computed angle. - if is_rotation_gate(op.id) && is_dynamic_angle(shot_idx) { - if op.id == OPID_RX || op.id == OPID_RY || op.id == OPID_RZ { - let angle = resolve_gate_angle(shot_idx); - let half = angle * 0.5; - let c = cos(half); - let s = sin(half); - if op.id == OPID_RX { - // [[cos(θ/2), -i·sin(θ/2)], [-i·sin(θ/2), cos(θ/2)]] - shot.unitary[0] = vec2f(c, 0.0); - shot.unitary[1] = vec2f(0.0, -s); - shot.unitary[4] = vec2f(0.0, -s); - shot.unitary[5] = vec2f(c, 0.0); - } else if op.id == OPID_RY { - // [[cos(θ/2), -sin(θ/2)], [sin(θ/2), cos(θ/2)]] - shot.unitary[0] = vec2f(c, 0.0); - shot.unitary[1] = vec2f(-s, 0.0); - shot.unitary[4] = vec2f(s, 0.0); - shot.unitary[5] = vec2f(c, 0.0); - } else { - // RZ: [[1, 0], [0, e^(iθ)]] - shot.unitary[0] = vec2f(1.0, 0.0); - shot.unitary[1] = vec2f(0.0, 0.0); - shot.unitary[4] = vec2f(0.0, 0.0); - shot.unitary[5] = vec2f(cos(angle), sin(angle)); - } - } else if op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_RZZ { - let angle = resolve_gate_angle(shot_idx); - let half = angle * 0.5; - let c = cos(half); - let s = sin(half); - if op.id == OPID_RXX { - // exp(-i·θ/2·X⊗X) - shot.unitary[0] = vec2f(c, 0.0); - shot.unitary[3] = vec2f(0.0, -s); - shot.unitary[5] = vec2f(c, 0.0); - shot.unitary[6] = vec2f(0.0, -s); - shot.unitary[9] = vec2f(0.0, -s); - shot.unitary[10] = vec2f(c, 0.0); - shot.unitary[12] = vec2f(0.0, -s); - shot.unitary[15] = vec2f(c, 0.0); - } else if op.id == OPID_RYY { - // exp(-i·θ/2·Y⊗Y) - shot.unitary[0] = vec2f(c, 0.0); - shot.unitary[3] = vec2f(0.0, s); - shot.unitary[5] = vec2f(c, 0.0); - shot.unitary[6] = vec2f(0.0, -s); - shot.unitary[9] = vec2f(0.0, -s); - shot.unitary[10] = vec2f(c, 0.0); - shot.unitary[12] = vec2f(0.0, s); - shot.unitary[15] = vec2f(c, 0.0); - } else { - // RZZ: diag(1, e^(iθ), e^(iθ), 1) - shot.unitary[0] = vec2f(1.0, 0.0); - shot.unitary[5] = vec2f(cos(angle), sin(angle)); - shot.unitary[10] = vec2f(cos(angle), sin(angle)); - shot.unitary[15] = vec2f(1.0, 0.0); - } - } - } - - shot.op_idx = op_idx; - shot.op_type = op.id; - - // If any operand is lost, dispatch the gate's configured loss - // policy (stamped on op.policy). - let has_lost_operand = gate_has_lost_operand(shot_idx, op_idx, q1, q2); - if (has_lost_operand) { - handle_lost_operand_policy(shot_idx, op_idx, q1, q2); - } - - // Check for noise ops after this gate in the ops pool - let pauli_op_idx = get_pauli_noise_idx(op_idx); - - // Handle Pauli noise (loss, if sampled, is recorded in pending_loss_mask) - if pauli_op_idx != 0u { - if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { - // A 1-qubit gate has a single operand; if it is lost there - // is no surviving qubit to receive Pauli noise. - if (!has_lost_operand) { - apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1); - } - } else { - if (has_lost_operand) { - // The gate body was handled by the loss policy above; - // apply the noise to the surviving operand (if any). - apply_2q_pauli_noise_on_survivor(shot_idx, op_idx, pauli_op_idx, q1, q2); - } else { - apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1, q2); - } - } - shots[shot_idx].interp.status = STATUS_RUNNING; - return; - } - - // If the gate has any lost operands (and no attached noise), the gate - // logic was completely handled inside `handle_lost_operand_policy`. - if (has_lost_operand) { - shots[shot_idx].interp.status = STATUS_RUNNING; - return; - } - - // No noise — set up the op for execution - - // Turn multi-qubit matrix ops into shot buffer ops - if op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP { - shot.op_type = OPID_SHOT_BUFF_2Q; - } - - // Turn 1Q matrix ops into shot buffer ops - if op.id >= OPID_X && op.id < OPID_CX { - shot.op_type = OPID_SHOT_BUFF_1Q; - } - - // Phase gates all execute as RZ - if is_1q_phase_gate(op.id) { - shot.op_type = OPID_RZ; - } - - // Set qubits_updated mask so next round knows which probabilities to update - switch shot.op_type { - case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { - shot.qubits_updated_last_op_mask = 0u; - } - case OPID_SHOT_BUFF_1Q { - shot.qubits_updated_last_op_mask = 1u << q1; - } - case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { - shot.qubits_updated_last_op_mask = (1u << q1) | (1u << q2); - } - default {} - } - } - case 1u { // Measure - // Check for noise ops before the measure op - // (noise is applied as Id+noise, then original measure, matching non-adaptive pattern) - let pauli_op_idx = get_pauli_noise_idx(op_idx); - - if pauli_op_idx != 0u { - // Apply noise to the Id gate before measure, then the measure itself - // The non-adaptive path inserts Id+noise before measure; here the Id - // is at op_idx and the original measure op follows after noise ops - if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { - apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1); - } else { - apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, q1, q2); - } - shots[shot_idx].interp.status = STATUS_RUNNING; - return; - } - - // No noise — standard measure - let resets = op.id == OPID_MRESETZ; - prep_measure_reset(shot_idx, op_idx, false, true, resets); - } - case 2u { // Reset - prep_measure_reset(shot_idx, op_idx, false, false, true); - } - default { - shot.op_type = OPID_ID; - } - } - - // Mark shot as running so interpret_classical resumes next round - shots[shot_idx].interp.status = STATUS_RUNNING; -} - -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn execute( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let shot_idx_u32: u32 = u32(shot_idx); - let shot = &shots[shot_idx]; - - // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating - // Correlated noise also updates probabilities in prepare_op, so can skip doing that here - let update_probs = shot.op_type != OPID_ID && shot.op_type != OPID_CORRELATED_NOISE && - shot.op_type != OPID_RZ && shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; - - if (shot.op_type == OPID_ID) { - // IGNORE - } else if (shot.op_type == OPID_CORRELATED_NOISE) { - apply_correlated_noise(workgroupId.x, tid); - } else if (shot.op_type == OPID_LOSS_NOISE) { - // Loss commit: the lost qubit is carried in op_idx (set by prep_loss_commit). - apply_1q_op(workgroupId.x, tid, shot.op_idx); - } else if (is_1q_op(shot.op_type)) { - let q1: u32 = resolve_q1(shot_idx_u32); - apply_1q_op(workgroupId.x, tid, q1); - } else /* 2 qubit op */ { - let q1: u32 = resolve_q1(shot_idx_u32); - let q2: u32 = resolve_q2(shot_idx_u32); - apply_2q_op(workgroupId.x, tid, q1, q2); - } - - // workgroupBarrier can't be conditional in DX12 backend, so we have to do an unconditional one here - // outside of the skip_work conditional above. - workgroupBarrier(); - - // If the workgroup is done updating, have the first thread reduce the per-thread probabilities into the - // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. - // Skip for correlated noise since probabilities were already updated in prepare_op. - if (tid == 0 && update_probs) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { - sum_thread_totals_to_shot(q, shot_idx, workgroup_collation_idx); - } - } - } -} diff --git a/source/simulators/src/gpu_full_state_simulator/simulator_base.wgsl b/source/simulators/src/gpu_full_state_simulator/simulator_base.wgsl deleted file mode 100644 index ba802929851..00000000000 --- a/source/simulators/src/gpu_full_state_simulator/simulator_base.wgsl +++ /dev/null @@ -1,449 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -// common.wgsl is appended to the beginning of this file at runtime. - -@group(0) @binding(0) -var workgroup_collation: WorkgroupCollationBuffer; -// Around 128 max partitions times 27 qubits times 8 bytes = 27 KB max size - -struct QubitState { - zero_probability: f32, - one_probability: f32, - heat: f32, // -1.0 = lost - idle_since: f32, -} - -// Used to track state for the random number generator per shot. See `next_rand_f32` later for details. -struct xorwow_state { - counter: u32, - x: array -} - -// Buffer containing the state for each shot to execute per kernel dispatch -// An instance of this is tracked on the GPU for every active shot -struct ShotData { - shot_id: u32, - next_op_idx: u32, - - // The below random numbers will be initialized from the RNG per operation in the 'prepare_op' stage - // Then the 'execute_op' stage will read these precomputed random numbers for noise modeling - rng_state: xorwow_state, // 6 x u32 - rand_pauli: f32, - rand_damping: f32, - rand_dephase: f32, - rand_measure: f32, - // Bitmask of qubits the most recent noise sampler chose to lose. A following - // loss-commit op consumes (and clears) its qubit's bit. - pending_loss_mask: u32, - - // The type of the next operation to execute. This will be OPID_SHOT_BUFF_* if it should use the unitary from the op buffer - op_type: u32, - op_idx: u32, - - duration: f32, // Total duration of the shot so far, used for time-dependent noise modeling and shot estimations - renormalize: f32, // Value to renormalize the state vector by on next execute (1.0 = no renormalization needed) - - // For quick testing during execution to enable skipping blocks of entries - // TODO: Actually use these masks during execution to skip unneeded work - qubit_is_0_mask: u32, // Bitmask for which qubits are currently in |0> state - qubit_is_1_mask: u32, // Bitmask for which qubits are currently in |1> state - - // Track which qubit probabilities were updated in the last operation (to collate on next prepare_op) - qubits_updated_last_op_mask: u32, - // 20 x 4 bytes to this point = 80 bytes - - // Track the per-qubit probabilities for optimization of measurement sampling and noise modeling - qubit_state: array, // 27 x 16 bytes = 432 bytes - // 512 bytes to this point - - // Map this to the Op structure for ease of use - unitary: array, // For MAT1Q and MAT2Q ops. -} -// Total struct size = 640 bytes (which is aligned to 128 bytes) -// See https://www.w3.org/TR/WGSL/#structure-member-layout for alignment rules - -@group(0) @binding(1) -var shots: array; - -// Buffer containing the list of operations (gates and noise) that make up the program to simulate -struct Op { - id: u32, - q1: u32, - q2: u32, - q3: u32, - policy: u32, - pad0: u32, - pad1: u32, - pad2: u32, - // Entries in the unitary are: 00, 01, 02, 03, 10, 11, 12, 13, 20, ..., 32, 33 - // 1q matrix elements are stored in: 00, 01, 10, 11 (i.e., indices 0, 1, 4, and 5) - unitary: array, -} // Struct size: 4 * 4 + 16 * 8 = 160 bytes (which is aligned to 16 bytes) - -@group(0) @binding(2) -var ops: array; - -// The one large buffer of state vector amplitudes. (Partitioned into multiple shots) -@group(0) @binding(3) -var stateVector: array; - -// Buffer for storing measurement results per shot -@group(0) @binding(4) -var results: array>; - -// When an error occurs, the below diagnostic data structure is used to store information about the error -struct DiagnosticData { - error_code: atomic, - termination_count: atomic, - extra1: u32, - extra2: f32, - extra3: f32, - _padding: u32, - shot: ShotData, // 640 bytes - op: Op, // 144 bytes - // Below is usually 6,912 bytes (size = THREADS_PER_WORKGROUP (32) * (8 * MAX_QUBIT_COUNT (27)) - workgroup_probabilities: array, - // Below is usually 27,648 bytes (1 << u32(MAX_QUBIT_COUNT - MAX_QUBITS_PER_WORKGROUP)) * (8 * MAX_QUBIT_COUNT) bytes - collation_buffer: WorkgroupCollationBuffer, -}; - -@group(0) @binding(5) -var diagnostics: DiagnosticData; - -struct Uniforms { - batch_start_shot_id: i32, - rng_seed: u32, -} - -@group(0) @binding(6) -var uniforms: Uniforms; - -struct NoiseTableMetadata { - /// The total probability of any noise (i.e. sum of all noise entries) in `Q1.63` format - noise_probability_lo: u32, - noise_probability_hi: u32, - /// The start offset of this table's entries in the global `NoiseTableEntry` array - start_offset: u32, - /// The number of entries in this noise table - entry_count: u32, -} - -struct NoiseTableEntry { - /// The correlated pauli string as bits (2 bits per qubit). If bit 0 is set, then it has bit-flip - /// noise, and if bit 1 is set then it has phase-flip noise. e.g., `110001 == "YIX"` - paulis_lo: u32, - paulis_hi: u32, - /// The probability of the noise occurring in `Q1_63` format. This is a float format where the high - /// order bit (bit 63) has the value 1.0 (`2^0 / 1`), bit 62 has the value 0.5 (`2^1 / 1`), etc. - /// all the way to bit 63 with a value of approx 1.0842e-19 (`2^63 / 1`). This gives a range of - /// values from [0..2) with equal spacing of 1.0842e-19 between values (unlike float or double), - /// which makes it more suitable for random numbers used to select between a large number of small - /// probability entries. - probability_lo: u32, - probability_hi: u32, -} - -// Template constants for noise table sizes (must be ≥ 1; host uses max(count,1)). -const NOISE_TABLE_COUNT: u32 = {{NOISE_TABLE_COUNT}}; -const NOISE_ENTRY_COUNT: u32 = {{NOISE_ENTRY_COUNT}}; - -// BatchData holds all the read-only data shared across all shots in a batch. -struct BatchData { - correlated_noise_tables: array, - correlated_noise_entries: array, -} - -@group(0) @binding(7) -var batch_data: BatchData; - - -// For every qubit, each 'execute' kernel thread will update its own workgroup storage location for accumulating probabilities -// The final probabilities will be reduced and written back to the shot state after the parallel execution completes. -struct QubitProbabilityPerThread { - zero: array, - one: array, -}; // size: 216 bytes - -var qubitProbabilities: array; -// Workgroup memory size: THREADS_PER_WORKGROUP (32) * 216 = 6,912 bytes. - -fn get_measure_qubit(shot_idx: u32, op_idx: u32) -> u32 { - return ops[op_idx].q1; -} - -fn get_measure_result(shot_idx: u32, op_idx: u32) -> u32 { - return ops[op_idx].q2; -} - -// Get the qubit id at the given index from the correlated noise op's qubit args -// Qubit args are stored in the unitary matrix elements as f32 values -fn get_correlated_noise_qubit(op_idx: u32, index: u32) -> u32 { - // Qubit ids are stored in the unitary as f32 values, starting at unitary[0].x, unitary[0].y, etc. - let vec_idx = index / 2u; - let component = index % 2u; - if (component == 0u) { - return u32(ops[op_idx].unitary[vec_idx].x); - } else { - return u32(ops[op_idx].unitary[vec_idx].y); - } -} - -// Prepare the shot state for executing a correlated noise operation. -// Resolves qubit IDs from the op's unitary matrix, samples the noise table, builds masks, and applies. -fn prep_correlated_noise(shot_idx: u32, op_idx: u32) { - let op = &ops[op_idx]; - let noise_table_idx = op.q1; - let qubit_count = op.q2; - - let sample = sample_correlated_noise(shot_idx, op_idx, noise_table_idx); - if (sample.should_apply == 0u) { return; } - - // Build bit-flip, phase-flip, and loss masks using qubit IDs from the op's unitary matrix - var bit_flip_mask: u32 = 0u; - var phase_flip_mask: u32 = 0u; - var loss_mask: u32 = 0u; - for (var i: u32 = 0u; i < qubit_count; i++) { - let pauli_bits = get_pauli_bits(sample.paulis_lo, sample.paulis_hi, qubit_count, i); - let qubit_mask = 1u << get_correlated_noise_qubit(op_idx, i); - if ((pauli_bits & 0x4u) != 0u) { - // Loss term (L = 4): the qubit is lost, no Pauli is applied to it. - loss_mask |= qubit_mask; - } else { - if ((pauli_bits & 0x1u) != 0u) { bit_flip_mask |= qubit_mask; } - if ((pauli_bits & 0x2u) != 0u) { phase_flip_mask |= qubit_mask; } - } - } - - commit_correlated_noise(shot_idx, op_idx, bit_flip_mask, phase_flip_mask, loss_mask); -} - - -// ******************************* -// PREPARE OP -// This stage prepares the shot state for the next operation to execute (and any updates needed from the prior op) -// -// Each op is prepared by one thread. This is how we deal with some of the challenges with synchronization -// when multiple workgroups with multiple threads are used for a shot in the EXECUTE stage. The 'execute_op' -// does work that is 'embarrassingly parallel' across the state vector amplitudes, but the PREPARE_OP stage -// deal with preparing for that work, and collating results back into the shot state afterwards. -// -// This allows us to use the GPU 'dispatch' mechanism to ensure consistencty across shots without complex, -// synchronization code, as the GPU guarantees that all threads in a dispatch complete before the next dispatch -// starts, and all buffer writes are visible to the next dispatch. -// ******************************* - -// NOTE: Run with workgroup size of 1 for now, as threads may diverge too much in prepare_op stage causing performance issues. -// TODO: Try to increase later if lack of parallelism is a bottleneck. (Update the dispatch call accordingly). -@compute @workgroup_size(1) -fn prepare_op(@builtin(global_invocation_id) globalId: vec3) { - // For the 'prepare_op' stage, each thread dispatched handles one shot, so the globalId.x is the shot index - let shot_idx = globalId.x; - let shot = &shots[shot_idx]; - - // WebGPU guarantees that buffers are zero-initialized, so next_op_idx will correctly be 0 on the first dispatch - let op_idx = shot.next_op_idx; - - // If we've gone past the end, set the op type to id and exit, so the execute stage is a no-op - if (op_idx >= u32(arrayLength(&ops))) { - // TODO: Set error/diagnostic info here - shot.op_type = OPID_ID; - shot.renormalize = 1.0; - shot.qubits_updated_last_op_mask = 0u; - return; - } - - let op = &ops[op_idx]; - - // Update the shot state based on the results of the last executed op (if needed) - if (shot.qubits_updated_last_op_mask != 0) { - update_qubit_state(shot_idx); - } - - shot_init_per_op(shot_idx); - shot.unitary = op.unitary; - - // Handle MResetZ, MZ, and ResetZ operations. These have unique handling and no associated noise ops, so prep and exit - if (op.id == OPID_MRESETZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, true /* resets_to_zero */); - shot.next_op_idx = op_idx + 1u; // No associated noise ops, so just advance by 1 - return; - } - if (op.id == OPID_MZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, true /* stores_result */, false /* resets_to_zero */); - shot.next_op_idx = op_idx + 1u; - return; - } - if (op.id == OPID_RESETZ) { - prep_measure_reset(shot_idx, op_idx, false /* is_loss */, false /* stores_result */, true /* resets_to_zero */); - shot.next_op_idx = op_idx + 1u; - return; - } - - // Loss-commit op: lose this qubit if and only if the preceding noise sampler - // set its bit in pending_loss_mask; otherwise act as identity. - if (op.id == OPID_LOSS_NOISE) { - shot.next_op_idx = op_idx + 1u; - let loss_bit = 1u << op.q1; - if ((shot.pending_loss_mask & loss_bit) != 0u) { - shot.pending_loss_mask &= ~loss_bit; - prep_measure_reset(shot_idx, op_idx, true /* is_loss */, false /* stores_result */, true /* resets_to_zero */); - } else { - shot.op_type = OPID_ID; - shot.op_idx = op_idx; - shot.qubits_updated_last_op_mask = 0u; - } - return; - } - - /* Handle noise: - - For the 1-qubit op case, there could be pauli and loss noise after the op itself. We want to check for loss first and - only apply pauli noise if the qubit wasn't lost. (If lost, the pauli noise and even the gate itself don't matter). - - For the 2-qubit op case, there will only be optional pauli noise after the op itself. (Loss is applied via separate - Id ops on each qubit after the 2-qubit op). - */ - - let pauli_op_idx = get_pauli_noise_idx(op_idx); - // Advance past this gate and its (optional) inline Pauli/loss noise op. Any - // loss-commit ops that follow are separate ops handled on later iterations. - shot.next_op_idx = max(op_idx, pauli_op_idx) + 1u; - - // Handle correlated noise operations - if (op.id == OPID_CORRELATED_NOISE) { - prep_correlated_noise(shot_idx, op_idx); - return; - } - - // Before doing further work, if any qubit for the gate is lost, dispatch - // the gate's configured loss policy (stamped on op.policy). - let has_lost_operand = gate_has_lost_operand(shot_idx, op_idx, op.q1, op.q2); - if (has_lost_operand) { - handle_lost_operand_policy(shot_idx, op_idx, op.q1, op.q2); - } - - if pauli_op_idx != 0 { - if ops[pauli_op_idx].id == OPID_PAULI_NOISE_1Q { - // A 1-qubit gate has a single operand; if it is lost there is no - // surviving qubit to receive Pauli noise, so skip the noise. - if (!has_lost_operand) { - apply_1q_pauli_noise(shot_idx, op_idx, pauli_op_idx, op.q1); - } - return; - } else { - if (has_lost_operand) { - // The gate body was handled by the loss policy above. Still apply - // the attached Pauli noise to the surviving operand (if any). - apply_2q_pauli_noise_on_survivor(shot_idx, op_idx, pauli_op_idx, op.q1, op.q2); - } else { - apply_2q_pauli_noise(shot_idx, op_idx, pauli_op_idx, op.q1, op.q2); - } - return; - } - } - - // If the gate has any lost operands (and no attached noise), the gate logic - // was completely handled inside `handle_lost_operand_policy`. - if (has_lost_operand) { - return; - } - - // No noise to apply, just set up the shot to execute the op as-is - shot.op_idx = op_idx; - shot.op_type = op.id; - - // Turn any Rxx, Ryy, or Rzz gates into a gate from the shot buffer - // NOTE: Should probably just do this for all gates - if (op.id == OPID_RXX || op.id == OPID_RYY || op.id == OPID_MAT2Q || op.id == OPID_SWAP) { - shot.op_type = OPID_SHOT_BUFF_2Q; // Indicate to use the matrix in the shot buffer - } - - if (op.id >= OPID_X && op.id < OPID_CX) { - shot.op_type = OPID_SHOT_BUFF_1Q; // Indicate to use the matrix in the shot buffer - } - - if (is_1q_phase_gate(op.id)) { - // For phase gates, treat everything as RZ for execution purposes - shot.op_type = OPID_RZ; - } - - // Set this so the next prepare_op stage knows which qubits to update probabilities for - switch shot.op_type { - case OPID_ID, OPID_CZ, OPID_RZ, OPID_RZZ { - shot.qubits_updated_last_op_mask = 0u; - } - case OPID_SHOT_BUFF_1Q { - shot.qubits_updated_last_op_mask = 1u << op.q1; - } - case OPID_CX, OPID_CY, OPID_SHOT_BUFF_2Q { - shot.qubits_updated_last_op_mask = (1u << op.q1) | (1u << op.q2); - } - default { - // TODO: Set error/diagnostic info here - } - } -} - -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn initialize( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - // Get the params - let params = get_shot_params(workgroupId.x, tid, 0 /* qubits per op */); - - // We want every thread to zero out its portion of the state vector for the shot - // We also want threads executing in lockstep to update adjacent entries for better memory access patterns - for (var i = 0; i < params.op_iterations; i++) { - let entry_index: i32 = params.thread_idx_in_shot + i * params.total_threads_per_shot; - stateVector[params.shot_state_vector_start + entry_index] = vec2f(0.0, 0.0); - } - - // NOTE: No need to synchronize here, as each thread is writing to unique locations - if (params.thread_idx_in_shot == 0) { - // Set the |0...0> amplitude to 1.0 from the first workgroup & thread for the shot - stateVector[params.shot_state_vector_start] = vec2f(1.0, 0.0); - reset_all(params.shot_idx); - } -} - -@compute @workgroup_size(THREADS_PER_WORKGROUP) -fn execute( - @builtin(workgroup_id) workgroupId: vec3, - @builtin(local_invocation_index) tid: u32) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let shot = &shots[shot_idx]; - - // If it's an ID gate, or a pure phase gate (including CZ) then probabilities don't need updating - // Correlated noise also updates probabilities in prepare_op, so can skip doing that here - let update_probs = shot.op_type != OPID_ID && shot.op_type != OPID_CORRELATED_NOISE && - shot.op_type != OPID_RZ && shot.op_type != OPID_CZ && shot.op_type != OPID_RZZ; - - if (shot.op_type == OPID_ID) { - // IGNORE - } else if (shot.op_type == OPID_CORRELATED_NOISE) { - apply_correlated_noise(workgroupId.x, tid); - } else if (is_1q_op(shot.op_type)) { - let q1: u32 = ops[shot.op_idx].q1; - apply_1q_op(workgroupId.x, tid, q1); - } else /* 2 qubit op */ { - let q1: u32 = ops[shot.op_idx].q1; - let q2: u32 = ops[shot.op_idx].q2; - apply_2q_op(workgroupId.x, tid, q1, q2); - } - - // workgroupBarrier can't be conditional in DX12 backend, so we have to do an unconditional one here - // outside of the skip_work conditional above. - workgroupBarrier(); - - // If the workgroup is done updating, have the first thread reduce the per-thread probabilities into the - // totals for this workgroup. The subsequent 'prepare_op' will sum the workgroup entries into the shot state. - // Skip for correlated noise since probabilities were already updated in prepare_op. - if (tid == 0 && update_probs) { - let shot_idx: i32 = i32(workgroupId.x) / WORKGROUPS_PER_SHOT; - let workgroup_collation_idx: i32 = select(-1, i32(workgroupId.x), WORKGROUPS_PER_SHOT > 1); - for (var q: u32 = 0u; q < u32(QUBIT_COUNT); q++) { - if (shot.qubits_updated_last_op_mask & (1u << q)) != 0u { - sum_thread_totals_to_shot(q, shot_idx, workgroup_collation_idx); - } - } - } -} From e2f24396ca79a9f6b9709935a595d5143063fe23 Mon Sep 17 00:00:00 2001 From: Bill Ticehurst Date: Mon, 3 Aug 2026 11:31:53 -0700 Subject: [PATCH 7/7] Rename wgsl file --- source/simulators/src/bytecode/runtime.rs | 2 +- .../src/gpu_full_state_simulator/gpu_resources.rs | 7 ++++--- .../{unified.wgsl => gpu_statevector_shaders.wgsl} | 0 3 files changed, 5 insertions(+), 4 deletions(-) rename source/simulators/src/gpu_full_state_simulator/{unified.wgsl => gpu_statevector_shaders.wgsl} (100%) diff --git a/source/simulators/src/bytecode/runtime.rs b/source/simulators/src/bytecode/runtime.rs index 5c83957f7ed..8288df7ba50 100644 --- a/source/simulators/src/bytecode/runtime.rs +++ b/source/simulators/src/bytecode/runtime.rs @@ -21,7 +21,7 @@ use crate::{ // --------------------------------------------------------------------------- // Opcode constants — must stay in sync with the Python `_adaptive_bytecode.py` -// and the WGSL `unified.wgsl` shader. +// and the WGSL `gpu_statevector_shaders.wgsl` shader. // --------------------------------------------------------------------------- // Flags (pre-shifted to bit 16+) diff --git a/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs b/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs index f929fe00aca..d64ba9c04cd 100644 --- a/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs +++ b/source/simulators/src/gpu_full_state_simulator/gpu_resources.rs @@ -399,7 +399,7 @@ impl GpuResources { /// Creates the compute pipelines from the unified shader source. /// - /// The same `unified.wgsl` source powers both the base (linear op-list) and + /// The same `gpu_statevector_shaders.wgsl` source powers both the base (linear op-list) and /// adaptive (QIR bytecode interpreter) simulators; `is_adaptive` selects /// which code paths are compiled in via the `IS_ADAPTIVE` constant. pub(crate) fn create_shaders( @@ -459,7 +459,8 @@ impl GpuResources { ("IS_ADAPTIVE", is_adaptive.to_string()), ]; - let mut shader_src = stamp_shader_consts(include_str!("unified.wgsl"), &replacements); + let mut shader_src = + stamp_shader_consts(include_str!("gpu_statevector_shaders.wgsl"), &replacements); // Strip out DX12-incompatible code sections if needed if adapter.get_info().backend == wgpu::Backend::Dx12 { @@ -857,7 +858,7 @@ impl GpuResources { /// Stamps compile-time constant values into the shader source. /// -/// `unified.wgsl` declares its host-substituted constants with a literal +/// `gpu_statevector_shaders.wgsl` declares its host-substituted constants with a literal /// default followed by a `// REPLACE` marker, e.g. /// `const QUBIT_COUNT: i32 = 8; // REPLACE`. Keeping a valid literal default /// (rather than a `{{PLACEHOLDER}}` token) lets editor tooling parse the file. diff --git a/source/simulators/src/gpu_full_state_simulator/unified.wgsl b/source/simulators/src/gpu_full_state_simulator/gpu_statevector_shaders.wgsl similarity index 100% rename from source/simulators/src/gpu_full_state_simulator/unified.wgsl rename to source/simulators/src/gpu_full_state_simulator/gpu_statevector_shaders.wgsl