Unified GPU shaders

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+)

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::<InterpreterState>()
-                + 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::<Instruction<u32>>()
+        + params.num_blocks.max(1) * size_of::<Block<u32>>()
+        + params.num_functions.max(1) * size_of::<Function<u32>>()
+        + params.num_phi_entries.max(1) * size_of::<PhiNodeEntry<u32>>()
+        + params.num_switch_cases.max(1) * size_of::<SwitchCase<u32>>()
+        + params.num_call_args.max(1) * size_of::<u32>()
+        + params.num_constant_data.max(1) * size_of::<u32>()
+}

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}}", &params.qubit_count.to_string())
-            .replace("{{RESULT_COUNT}}", &(params.result_count + 1).to_string()) // +1 for result code per shot
-            .replace(
-                "{{WORKGROUPS_PER_SHOT}}",
-                &params.workgroups_per_shot.to_string(),
-            )
-            .replace(
-                "{{ENTRIES_PER_THREAD}}",
-                &params.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}}", &params.num_registers.to_string())
-            .replace("{{MAX_MEMORY}}", &params.max_memory.to_string())
-            .replace(
-                "{{INSTRUCTIONS_SIZE}}",
-                &params.num_instructions.to_string(),
-            )
-            .replace("{{BLOCK_TABLE_SIZE}}", &params.num_blocks.to_string())
-            .replace("{{FUNCTION_TABLE_SIZE}}", &params.num_functions.to_string())
-            .replace("{{PHI_TABLE_SIZE}}", &params.num_phi_entries.to_string())
-            .replace(
-                "{{SWITCH_CASES_SIZE}}",
-                &params.num_switch_cases.to_string(),
-            )
-            .replace("{{CALL_ARGS_SIZE}}", &params.num_call_args.to_string())
-            .replace(
-                "{{CONSTANT_DATA_SIZE}}",
-                &params.num_constant_data.to_string(),
-            )
-            .replace(
-                "{{NOISE_TABLE_COUNT}}",
-                &params.noise_table_count.to_string(),
-            )
-            .replace(
-                "{{NOISE_ENTRY_COUNT}}",
-                &params.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: "<up to '='>= <value>; // REPLACE",
+            // preserving the original indentation and trailing marker.
+            let eq = line.find('=').expect("REPLACE const line must contain '='");
+            result.push_str(&line[..eq]);
+            result.push_str("= ");
+            result.push_str(value);
+            result.push_str("; // REPLACE");
+        } else {
+            result.push_str(line);
+        }
+        result.push('\n');
+    }
+
+    result
+}
+
 /// Strips out sections of code delimited by DX12-start-strip and DX12-end-strip comments
 fn strip_dx12_sections(source: &str) -> String {
     let mut result = String::new();