From 529e2d43efe155a31344f0ef2ae3e838781ee0c8 Mon Sep 17 00:00:00 2001 From: ttrenty <154608953+ttrenty@users.noreply.github.com> Date: Sun, 29 Jun 2025 13:35:12 -0600 Subject: [PATCH 1/2] feat: add support for GPU in abstraction + extend benchmarks with iterations --- benchmarks/bench_main.mojo | 10 +- benchmarks/bench_qubit_wise_multiply.mojo | 161 +++++---- benchmarks/bench_qubit_wise_multiply_gpu.mojo | 136 +++---- examples/circuit_level.mojo | 95 +++-- examples/main.mojo | 3 +- src/abstractions/simulator.mojo | 336 +++++++++++++++++- 6 files changed, 537 insertions(+), 204 deletions(-) diff --git a/benchmarks/bench_main.mojo b/benchmarks/bench_main.mojo index 4b0d110..d8524a9 100644 --- a/benchmarks/bench_main.mojo +++ b/benchmarks/bench_main.mojo @@ -16,13 +16,14 @@ def main(): # bench_qubit_wise_multiply() bench_qubit_wise_multiply_inplace[ min_number_qubits=1, - max_number_qubits=20, + max_number_qubits=25, number_qubits_step_size=1, min_number_layers=5, max_number_layers=3500, number_layers_step_size=400, fixed_number_qubits=13, - fixed_number_layers=5, + fixed_number_layers=15, + fixed_number_iterations=1, ]() @parameter @@ -31,13 +32,14 @@ def main(): else: bench_qubit_wise_multiply_inplace_gpu[ min_number_qubits=1, - max_number_qubits=26, # 29 is OOM for my 3070 Ti Laptop GPU + max_number_qubits=30, # 29 is OOM for my 3070 Ti Laptop GPU number_qubits_step_size=1, min_number_layers=5, max_number_layers=7000, number_layers_step_size=400, fixed_number_qubits=13, - fixed_number_layers=5, + fixed_number_layers=15, + fixed_number_iterations=1, ]() # bench_qubit_wise_multiply_extended() diff --git a/benchmarks/bench_qubit_wise_multiply.mojo b/benchmarks/bench_qubit_wise_multiply.mojo index ef6d3b7..07f385e 100644 --- a/benchmarks/bench_qubit_wise_multiply.mojo +++ b/benchmarks/bench_qubit_wise_multiply.mojo @@ -90,7 +90,7 @@ fn benchmark_qubit_wise_multiply[ @parameter @always_inline fn benchmark_qubit_wise_multiply_inplace[ - num_qubits: Int, number_layers: Int + num_qubits: Int, number_layers: Int, number_iterations: Int = 1 ](mut b: Bencher) raises: gates_list: List[Gate] = [Hadamard, PauliX, PauliY, PauliZ] @@ -99,7 +99,7 @@ fn benchmark_qubit_wise_multiply_inplace[ ) random.seed() # Seed on current time random.randint( - indexes, number_layers * 2 * num_qubits, 0, len(gates_list) - 1 + indexes, number_layers * 1 * num_qubits, 0, len(gates_list) - 1 ) @parameter @@ -125,80 +125,86 @@ fn benchmark_qubit_wise_multiply_inplace[ quantum_state_1 = StateVector.from_bitstring("0" * num_qubits) current_state = 0 - for layer in range(number_layers): - for i in range(num_qubits): - # NOTE Works but is slow with the dictionary - # qubit_wise_multiply_inplace( - # gates_list[Int(indexes[layer * num_qubits + i])].matrix, - # i, - # quantum_states[current_state], - # quantum_states[1 - current_state], - # ) - # NOTE: Fast buty doesn't actually use the next state for new operations - # qubit_wise_multiply_inplace( - # gates_list[Int(indexes[layer * num_qubits + i])].matrix, - # i, - # quantum_state_0, - # quantum_state_1, - # ) - if current_state == 0: - qubit_wise_multiply_inplace( - gates_list[Int(indexes[layer * num_qubits + i])].matrix, - i, - quantum_state_0, - quantum_state_1, - ) - current_state = 1 - else: - qubit_wise_multiply_inplace( - gates_list[Int(indexes[layer * num_qubits + i])].matrix, - i, - quantum_state_1, - quantum_state_0, - ) - current_state = 0 - for i in range(num_qubits - 1): - # qubit_wise_multiply_inplace( - # gates_list[ - # Int(indexes[layer * num_qubits + num_qubits + i]) - # ].matrix, - # i, - # quantum_states[current_state], - # quantum_states[1 - current_state], - # [[(i + 1) % num_qubits, 1]], - # ) - # current_state = 1 - current_state - # qubit_wise_multiply_inplace( - # gates_list[ - # Int(indexes[layer * num_qubits + num_qubits + i]) - # ].matrix, - # i, - # quantum_state_0, - # quantum_state_1, - # [[(i + 1) % num_qubits, 1]], - # ) - if current_state == 0: - qubit_wise_multiply_inplace( - gates_list[ - Int(indexes[layer * num_qubits + num_qubits + i]) - ].matrix, - i, - quantum_state_0, - quantum_state_1, - [[(i + 1) % num_qubits, 1]], - ) - current_state = 1 - else: - qubit_wise_multiply_inplace( - gates_list[ - Int(indexes[layer * num_qubits + num_qubits + i]) - ].matrix, - i, - quantum_state_1, - quantum_state_0, - [[(i + 1) % num_qubits, 1]], - ) - current_state = 0 + for iteration in range(number_iterations): + for layer in range(number_layers): + for i in range(num_qubits): + # NOTE Works but is slow with the dictionary + # qubit_wise_multiply_inplace( + # gates_list[Int(indexes[layer * num_qubits + i])].matrix, + # i, + # quantum_states[current_state], + # quantum_states[1 - current_state], + # ) + # NOTE: Fast but doesn't actually use the next state for new operations + # qubit_wise_multiply_inplace( + # gates_list[Int(indexes[layer * num_qubits + i])].matrix, + # i, + # quantum_state_0, + # quantum_state_1, + # ) + if current_state == 0: + qubit_wise_multiply_inplace( + gates_list[ + Int(indexes[layer * num_qubits + i]) + ].matrix, + i, + quantum_state_0, + quantum_state_1, + ) + current_state = 1 + else: + qubit_wise_multiply_inplace( + gates_list[ + Int(indexes[layer * num_qubits + i]) + ].matrix, + i, + quantum_state_1, + quantum_state_0, + ) + current_state = 0 + # Compare with GPU which currently doesn't support controls qubits yet + # for i in range(num_qubits - 1): + # # qubit_wise_multiply_inplace( + # # gates_list[ + # # Int(indexes[layer * num_qubits + num_qubits + i]) + # # ].matrix, + # # i, + # # quantum_states[current_state], + # # quantum_states[1 - current_state], + # # [[(i + 1) % num_qubits, 1]], + # # ) + # # current_state = 1 - current_state + # # qubit_wise_multiply_inplace( + # # gates_list[ + # # Int(indexes[layer * num_qubits + num_qubits + i]) + # # ].matrix, + # # i, + # # quantum_state_0, + # # quantum_state_1, + # # [[(i + 1) % num_qubits, 1]], + # # ) + # if current_state == 0: + # qubit_wise_multiply_inplace( + # gates_list[ + # Int(indexes[layer * num_qubits + num_qubits + i]) + # ].matrix, + # i, + # quantum_state_0, + # quantum_state_1, + # [[(i + 1) % num_qubits, 1]], + # ) + # current_state = 1 + # else: + # qubit_wise_multiply_inplace( + # gates_list[ + # Int(indexes[layer * num_qubits + num_qubits + i]) + # ].matrix, + # i, + # quantum_state_1, + # quantum_state_0, + # [[(i + 1) % num_qubits, 1]], + # ) + # current_state = 0 bench_ctx = DeviceContext() b.iter_custom[qubit_wise_multiply_inplace_workflow](bench_ctx) @@ -369,6 +375,7 @@ def bench_qubit_wise_multiply_inplace[ number_layers_step_size: Int = 200, fixed_number_qubits: Int = 5, fixed_number_layers: Int = 2, + fixed_number_iterations: Int = 1, ](): print("Running qubit_wise_multiply_inplace() Benchmarks...") print("-" * 80) @@ -382,7 +389,7 @@ def bench_qubit_wise_multiply_inplace[ ): bench.bench_function[ benchmark_qubit_wise_multiply_inplace[ - number_qubits, fixed_number_layers + number_qubits, fixed_number_layers, fixed_number_iterations ] ]( BenchId( diff --git a/benchmarks/bench_qubit_wise_multiply_gpu.mojo b/benchmarks/bench_qubit_wise_multiply_gpu.mojo index 0d2d529..24d9732 100644 --- a/benchmarks/bench_qubit_wise_multiply_gpu.mojo +++ b/benchmarks/bench_qubit_wise_multiply_gpu.mojo @@ -38,7 +38,7 @@ alias NUMBER_CONTROL_BITS = 1 @parameter @always_inline fn benchmark_qubit_wise_multiply_inplace_gpu[ - num_qubits: Int, number_layers: Int + num_qubits: Int64, number_layers: Int, number_iterations: Int = 1 ](mut b: Bencher) raises: # gates_list: List[Gate] = [Hadamard, PauliX, PauliY, PauliZ] @@ -71,7 +71,7 @@ fn benchmark_qubit_wise_multiply_inplace_gpu[ gate_set_size, GATE_SIZE, GATE_SIZE ) - alias state_vector_size = 1 << num_qubits + alias state_vector_size = 1 << Int(num_qubits) alias state_vector_layout = Layout.row_major(state_vector_size) alias total_threads = state_vector_size @@ -140,7 +140,7 @@ fn benchmark_qubit_wise_multiply_inplace_gpu[ # -- Initialize the quantum circuit to the |000⟩ state -- # quantum_state: StateVector = StateVector.from_bitstring( - "0" * num_qubits + "0" * Int(num_qubits) ) # print("Initial quantum state:\n", quantum_state) @@ -240,68 +240,69 @@ fn benchmark_qubit_wise_multiply_inplace_gpu[ # -- Apply circuit operations -- # current_state = 0 - for layer in range(number_layers): - for qubit in range(num_qubits): - if current_state == 0: - ctx.enqueue_function[ - qubit_wise_multiply_inplace_gpu[ - state_vector_size=state_vector_size, - gate_set_size=gate_set_size, - circuit_number_control_gates=circuit_number_control_gates, - number_control_bits=0, - ] - ]( - gate_set_re_tensor, - gate_set_im_tensor, - gate_set_dic[Hadamard.symbol], - # gate_set_dic[ - # gates_list[ - # Int(indexes[layer * num_qubits + qubit]) - # ].symbol - # ], - GATE_SIZE, - qubit, # target_qubit - quantum_state_re_tensor, - quantum_state_im_tensor, - num_qubits, # number_qubits - quantum_state_out_re_tensor, - quantum_state_out_im_tensor, - control_bits_circuit_tensor, - current_control_gate_circuit_tensor, - grid_dim=blocks_per_grid, - block_dim=threads_per_block, - ) - current_state = 1 - else: - ctx.enqueue_function[ - qubit_wise_multiply_inplace_gpu[ - state_vector_size=state_vector_size, - gate_set_size=gate_set_size, - circuit_number_control_gates=circuit_number_control_gates, - number_control_bits=0, - ] - ]( - gate_set_re_tensor, - gate_set_im_tensor, - gate_set_dic[Hadamard.symbol], - # gate_set_dic[ - # gates_list[ - # Int(indexes[layer * num_qubits + qubit]) - # ].symbol - # ], - GATE_SIZE, - qubit, # target_qubit - quantum_state_out_re_tensor, - quantum_state_out_im_tensor, - num_qubits, # number_qubits - quantum_state_re_tensor, - quantum_state_im_tensor, - control_bits_circuit_tensor, - current_control_gate_circuit_tensor, - grid_dim=blocks_per_grid, - block_dim=threads_per_block, - ) - current_state = 0 + for iteration in range(number_iterations): + for layer in range(number_layers): + for qubit in range(Int(num_qubits)): + if current_state == 0: + ctx.enqueue_function[ + qubit_wise_multiply_inplace_gpu[ + state_vector_size=state_vector_size, + gate_set_size=gate_set_size, + circuit_number_control_gates=circuit_number_control_gates, + number_control_bits=0, + ] + ]( + gate_set_re_tensor, + gate_set_im_tensor, + gate_set_dic[Hadamard.symbol], + # gate_set_dic[ + # gates_list[ + # Int(indexes[layer * num_qubits + qubit]) + # ].symbol + # ], + GATE_SIZE, + qubit, # target_qubit + quantum_state_re_tensor, + quantum_state_im_tensor, + Int(num_qubits), # number_qubits + quantum_state_out_re_tensor, + quantum_state_out_im_tensor, + control_bits_circuit_tensor, + current_control_gate_circuit_tensor, + grid_dim=blocks_per_grid, + block_dim=threads_per_block, + ) + current_state = 1 + else: + ctx.enqueue_function[ + qubit_wise_multiply_inplace_gpu[ + state_vector_size=state_vector_size, + gate_set_size=gate_set_size, + circuit_number_control_gates=circuit_number_control_gates, + number_control_bits=0, + ] + ]( + gate_set_re_tensor, + gate_set_im_tensor, + gate_set_dic[Hadamard.symbol], + # gate_set_dic[ + # gates_list[ + # Int(indexes[layer * num_qubits + qubit]) + # ].symbol + # ], + GATE_SIZE, + qubit, # target_qubit + quantum_state_out_re_tensor, + quantum_state_out_im_tensor, + Int(num_qubits), # number_qubits + quantum_state_re_tensor, + quantum_state_im_tensor, + control_bits_circuit_tensor, + current_control_gate_circuit_tensor, + grid_dim=blocks_per_grid, + block_dim=threads_per_block, + ) + current_state = 0 keep(quantum_state_re.unsafe_ptr()) keep(quantum_state_im.unsafe_ptr()) @@ -321,12 +322,13 @@ fn benchmark_qubit_wise_multiply_inplace_gpu[ def bench_qubit_wise_multiply_inplace_gpu[ min_number_qubits: Int = 15, max_number_qubits: Int = 25, - number_qubits_step_size: Int = 1, + number_qubits_step_size: Int64 = 1, min_number_layers: Int = 1, max_number_layers: Int = 2000, number_layers_step_size: Int = 200, fixed_number_qubits: Int = 5, fixed_number_layers: Int = 2, + fixed_number_iterations: Int = 1, ](): print("Running qubit_wise_multiply_inplace_gpu() Benchmarks...") print("-" * 80) @@ -341,7 +343,7 @@ def bench_qubit_wise_multiply_inplace_gpu[ ): bench.bench_function[ benchmark_qubit_wise_multiply_inplace_gpu[ - number_qubits, fixed_number_layers + number_qubits, fixed_number_layers, fixed_number_iterations ] ]( BenchId( diff --git a/examples/circuit_level.mojo b/examples/circuit_level.mojo index 06b0d31..13192f3 100644 --- a/examples/circuit_level.mojo +++ b/examples/circuit_level.mojo @@ -41,7 +41,7 @@ fn simulate_figure1_circuit_abstract() -> None: """ ) - num_qubits: Int = 3 + alias num_qubits: Int = 3 qc: GateCircuit = GateCircuit(num_qubits) @@ -59,10 +59,25 @@ fn simulate_figure1_circuit_abstract() -> None: qsimu = StateVectorSimulator( qc, initial_state=initial_state, + use_gpu_if_available=True, verbose=True, - verbose_step_size=ShowAfterEachGate, # ShowAfterEachGate, ShowOnlyEnd + verbose_step_size=ShowOnlyEnd, # ShowAfterEachGate, ShowAfterEachLayer ) + # # For GPU (NOTE: Doesn't work currently as GPU doesn't fully support control bits) + # qsimu = StateVectorSimulator[ + # gpu_num_qubits=num_qubits, + # gpu_gate_ordered_set= [Hadamard, PauliZ, NOT], + # gpu_control_gate_count=2, + # gpu_control_bits_list= [[[1, 1], [1, 1]]], + # ]( + # qc, + # initial_state=initial_state, + # use_gpu_if_available=True, + # verbose=True, + # verbose_step_size=ShowOnlyEnd, # ShowAfterEachGate, ShowAfterEachLayer + # ) + state = qsimu.run() print("Final quantum state:\n", state) @@ -107,85 +122,57 @@ fn simulate_figure4_circuit_abstract() -> None: print("Final quantum state:\n", final_state) -fn simulate_random_circuit(number_qubits: Int, number_layers: Int) -> None: +fn simulate_random_circuit[number_qubits: Int](number_layers: Int) -> None: """Simulates a random quantum circuit with the specified number of qubits and layers. - Args: + Parameters: number_qubits: The number of qubits in the circuit. + + Args: number_layers: The number of layers in the circuit. """ - qc: GateCircuit = GateCircuit(number_qubits) + alias gates_list: List[Gate] = [Hadamard, PauliX, PauliY, PauliZ] - gates_list: List[Gate] = [Hadamard, PauliX, PauliY, PauliZ] - - # index: UnsafePointer[Int8] = UnsafePointer[Int8].alloc(2*number_qubits) - # print("Creating random circuit...") - # random.seed() # Seed on current time - # for _ in range(400): - # random.randint(index, 2*number_qubits, 0, len(gates_list) - 1) - # for i in range(number_qubits): - # qc = qc.apply(gates_list[Int(index[i])], i) - # qc = qc.barrier() - # for i in range(number_qubits - 1): - # qc = qc.apply( - # gates_list[Int(index[number_qubits + i])], - # i, - # controls=[(i + 1) % number_qubits], - # is_anti_control=[False], - # ) - # qc = qc.barrier() + qc: GateCircuit = GateCircuit(number_qubits) print( - "Creating random circuit of", + "Creating random circuit with", number_qubits, "qubits and", number_layers, - "layers...", + "layers.", ) index: UnsafePointer[Int8] = UnsafePointer[Int8].alloc( number_layers * 2 * number_qubits ) - random.seed() # Seed on current time + random.seed(42) random.randint( - index, number_layers * 2 * number_qubits, 0, len(gates_list) - 1 + index, number_layers * 1 * number_qubits, 0, len(gates_list) - 1 ) - print("Random gates choices generated, applying them to the circuit...") - for iter in range(number_layers): - for i in range(number_qubits): - qc.apply(gates_list[Int(index[iter * number_qubits + i])](i)) - qc.barrier() - for i in range(number_qubits - 1): - qc.apply( - gates_list[ - Int(index[iter * number_qubits + number_qubits + i]) - ](i, controls=[(i + 1) % number_qubits]), - ) - qc.barrier() + for layer in range(number_layers): + for qubit in range(number_qubits): + gate = gates_list[Int(index[layer * number_qubits + qubit])] + qc.apply(gate(qubit)) print("Random circuit created has", qc.num_gates(), "total gates.") - # print(qc) + # Initial state |000...0⟩ + quantum_state: StateVector = StateVector.from_bitstring("0" * number_qubits) print("Running Simulations...") - initial_state_bitstring: String = ( - "0" * number_qubits - ) # Initial state |000...0⟩ - initial_state: StateVector = StateVector.from_bitstring( - initial_state_bitstring - ) - - qsimu = StateVectorSimulator( + qsimu = StateVectorSimulator[ + gpu_num_qubits=number_qubits, + gpu_gate_ordered_set=gates_list, + ]( qc, - initial_state=initial_state, + initial_state=quantum_state, + use_gpu_if_available=True, verbose=False, - # verbose_step_size=ShowAfterEachLayer, # ShowAfterEachGate, ShowOnlyEnd - verbose_step_size=ShowAfterEachGate, # ShowAfterEachGate, ShowOnlyEnd - # stop_at=ShowAfterEachGate, # ShowAfterEachGate, ShowOnlyEnd # TODO implement that instead of having access to manual methods ) - for i in range(1000): - print("Iteration:", i, end="\r") + for i in range(20): + print("Simulating Circuit:", i, end="\r") _ = qsimu.run() print("") diff --git a/examples/main.mojo b/examples/main.mojo index 9a4b1cc..551f098 100644 --- a/examples/main.mojo +++ b/examples/main.mojo @@ -27,6 +27,7 @@ from circuit_level import ( def main(): print("QLabs: Quantum Circuit Composer & Simulator in Mojo") + print("=========================================") print("low_level: simulate_figure1_circuit()") simulate_figure1_circuit() @@ -45,7 +46,7 @@ def main(): print("=========================================") print("circuit_level: simulate_random_circuit()") - simulate_random_circuit(number_qubits=10, number_layers=20) + simulate_random_circuit[number_qubits=20](number_layers=5) print("=========================================") print("circuit_level: simulate_figure4_circuit_abstract()") diff --git a/src/abstractions/simulator.mojo b/src/abstractions/simulator.mojo index cbfee9a..fa435de 100644 --- a/src/abstractions/simulator.mojo +++ b/src/abstractions/simulator.mojo @@ -2,6 +2,11 @@ # MARK: Imports # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +from sys import has_accelerator +from gpu import thread_idx, block_dim, block_idx +from gpu.host import DeviceContext +from layout import Layout, LayoutTensor, IntTuple, print_layout + from ..base.qubits_operations import ( qubit_wise_multiply, qubit_wise_multiply_extended, @@ -13,8 +18,30 @@ from ..base.state_and_matrix import ( ) from ..base.gate import _START, _SEPARATOR, SWAP +from ..base.gate import ( + Gate, + Hadamard, + PauliX, + PauliY, + PauliZ, + NOT, + H, + X, + Y, + Z, + iSWAP, +) + +from ..base.gpu import qubit_wise_multiply_inplace_gpu + +from ..local_stdlib.complex import ComplexFloat32 + from .gate_circuit import GateCircuit +alias dtype = DType.float32 + +alias GATE_SIZE = 2 +alias NUMBER_CONTROL_BITS = 1 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # # MARK: Aliases # @@ -30,9 +57,26 @@ alias ShowOnlyEnd = "ShowOnlyEnd" # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # -struct StateVectorSimulator(Copyable, Movable): +struct StateVectorSimulator[ + gpu_num_qubits: Int = 0, + gpu_gate_ordered_set: List[Gate] = [ + Hadamard, + PauliX, + PauliY, + PauliZ, + iSWAP, + ], + gpu_control_gate_count: Int = 0, + gpu_control_bits_list: List[List[List[Int]]] = [[[]]], +](Copyable, Movable): """A noiseless simulator for quantum circuits that uses State Vector Simulation. + Parameters: + gpu_num_qubits: The number of qubits in the quantum circuit for GPU simulation. + gpu_gate_ordered_set: The ordered set of gates to be used in the GPU simulation. + gpu_control_gate_count: The number of control gates in the circuit for GPU simulation. + gpu_control_bits_list: A list of control bits for each gate in the circuit for GPU simulation. + This simulator applies quantum gates to a state vector representing the quantum state. """ @@ -80,6 +124,37 @@ struct StateVectorSimulator(Copyable, Movable): self.use_gpu_if_available = use_gpu_if_available self.verbose = verbose self.verbose_step_size = verbose_step_size + if self.use_gpu_if_available: + if not has_accelerator(): + print( + "No compatible GPU found. Falling back to CPU simulation." + ) + self.use_gpu_if_available = ( + False # Disable GPU if not available + ) + else: + try: + ctx = DeviceContext() + print("Using GPU simulation on device:", ctx.name()) + except e: + print("Failed to create GPU context:", e) + print("Falling back to CPU simulation.") + self.use_gpu_if_available = False + if ( + self.use_gpu_if_available + and self.verbose + and self.verbose_step_size != ShowOnlyEnd + ): + print( + ( + "Cannot have verbose output other than ShowOnlyEnd when" + " using GPU simulation." + ), + "Defaulting to ShowOnlyEnd.", + ) + self.verbose_step_size = ( + ShowOnlyEnd # Force ShowOnlyEnd for GPU simulation + ) @always_inline fn next_gate( @@ -206,6 +281,15 @@ struct StateVectorSimulator(Copyable, Movable): ) print("Initial state:\n", self.initial_state) + if self.use_gpu_if_available: + try: + ctx = DeviceContext() + final_state = self.run_gpu(ctx) + return final_state + except e: + print("Failed to create GPU context:", e) + print("Falling back to CPU simulation.") + # Start with the initial state quantum_state: StateVector = self.initial_state i: Int = 0 @@ -258,3 +342,253 @@ struct StateVectorSimulator(Copyable, Movable): print("Final quantum state:\n", quantum_state) return quantum_state + + def run_gpu(self, ctx: DeviceContext) -> StateVector: + """Runs the quantum circuit simulation. + + Applies all gates in sequence to the initial state and computes the + final quantum state. Will print verbose output if enabled and at + the specified verbosity level. + + Returns: + The final `StateVector` after all gates have been applied. + """ + + quantum_state: StateVector = self.initial_state + + alias num_qubits = gpu_num_qubits + + alias circuit_number_control_gates = gpu_control_gate_count + alias circuit_control_bits_layout = Layout.row_major( + circuit_number_control_gates, NUMBER_CONTROL_BITS, 2 + ) + + gate_set_dic: Dict[String, Int] = {} + for i in range(len(gpu_gate_ordered_set)): + gate_set_dic[gpu_gate_ordered_set[i].symbol] = i + + alias gate_set_size = len(gpu_gate_ordered_set) + alias gate_set_1qubit_layout = Layout.row_major( + gate_set_size, GATE_SIZE, GATE_SIZE + ) + + alias state_vector_size = 1 << num_qubits + alias state_vector_layout = Layout.row_major(state_vector_size) + + alias total_threads = state_vector_size + + alias max_threads_per_block = 1024 # Maximum threads per block in CUDA + alias blocks_per_grid = ( + total_threads + max_threads_per_block - 1 + ) // max_threads_per_block + + alias threads_per_block = ( + max_threads_per_block, + 1, + 1, + ) + + @parameter + if total_threads < max_threads_per_block: + alias threads_per_block = ( + total_threads, + 1, + 1, + ) # 1D block of threads + + # var control_bits_list: List[List[List[Int]]] = [ + # [[1, 1]], # Control on qubit 1 and is control because flag=1 + # [[1, 1]], # Control on qubit 1 and is control because flag=1 + # ] + var control_bits_list = gpu_control_bits_list + + # -- Create GPU variables -- # + # These don't need to be initialized to zero, they will be filled later + + host_quantum_state_re = ctx.enqueue_create_host_buffer[dtype]( + state_vector_size + ) + host_quantum_state_im = ctx.enqueue_create_host_buffer[dtype]( + state_vector_size + ) + + host_gate_set_re = ctx.enqueue_create_host_buffer[dtype]( + gate_set_size * GATE_SIZE * GATE_SIZE + ) + host_gate_set_im = ctx.enqueue_create_host_buffer[dtype]( + gate_set_size * GATE_SIZE * GATE_SIZE + ) + + host_control_bits_circuit = ctx.enqueue_create_host_buffer[DType.int32]( + circuit_number_control_gates * NUMBER_CONTROL_BITS * 2 + ) + + # -- Initialize the quantum circuit to the |000⟩ state -- # + + # Wait for host buffers to be ready + ctx.synchronize() + + # -- Fill host buffers -- # + + for i in range(state_vector_size): + host_quantum_state_re[i] = quantum_state[i].re + host_quantum_state_im[i] = quantum_state[i].im + + for i in range(gate_set_size): + gate = gpu_gate_ordered_set[i] + for j in range(GATE_SIZE): + for k in range(GATE_SIZE): + index = gate_set_1qubit_layout( + IntTuple(i, j, k) + ) # Get the index in the 1D buffer + host_gate_set_re[index] = gate[j, k].re + host_gate_set_im[index] = gate[j, k].im + + for i in range(circuit_number_control_gates): + for j in range(NUMBER_CONTROL_BITS): + for k in range(2): + index = circuit_control_bits_layout(IntTuple(i, j, k)) + host_control_bits_circuit[index] = control_bits_list[i][j][ + k + ] + + # -- Copy host buffers to device buffers -- # + quantum_state_re = ctx.enqueue_create_buffer[dtype](state_vector_size) + quantum_state_im = ctx.enqueue_create_buffer[dtype](state_vector_size) + + gate_set_re = ctx.enqueue_create_buffer[dtype]( + gate_set_size * GATE_SIZE * GATE_SIZE + ) + gate_set_im = ctx.enqueue_create_buffer[dtype]( + gate_set_size * GATE_SIZE * GATE_SIZE + ) + + control_bits_circuit = ctx.enqueue_create_buffer[DType.int32]( + circuit_number_control_gates * NUMBER_CONTROL_BITS * 2 + ) + current_control_gate_circuit = ctx.enqueue_create_buffer[DType.int32](1) + + # Create other buffers for functions + + quantum_state_out_re = ctx.enqueue_create_buffer[dtype]( + state_vector_size + ) + quantum_state_out_im = ctx.enqueue_create_buffer[dtype]( + state_vector_size + ) + + quantum_state_re.enqueue_copy_from(host_quantum_state_re) + quantum_state_im.enqueue_copy_from(host_quantum_state_im) + + gate_set_re.enqueue_copy_from(host_gate_set_re) + gate_set_im.enqueue_copy_from(host_gate_set_im) + + control_bits_circuit.enqueue_copy_from(host_control_bits_circuit) + + ctx.enqueue_memset(current_control_gate_circuit, 0) + ctx.enqueue_memset(quantum_state_out_re, 0.0) + ctx.enqueue_memset(quantum_state_out_im, 0.0) + + # -- Create layout tensors for GPU operations -- # + gate_set_re_tensor = LayoutTensor[ + mut=False, dtype, gate_set_1qubit_layout + ](gate_set_re.unsafe_ptr()) + gate_set_im_tensor = LayoutTensor[ + mut=False, dtype, gate_set_1qubit_layout + ](gate_set_im.unsafe_ptr()) + + quantum_state_re_tensor = LayoutTensor[ + mut=True, dtype, state_vector_layout + ](quantum_state_re.unsafe_ptr()) + quantum_state_im_tensor = LayoutTensor[ + mut=True, dtype, state_vector_layout + ](quantum_state_im.unsafe_ptr()) + + quantum_state_out_re_tensor = LayoutTensor[ + mut=True, dtype, state_vector_layout + ](quantum_state_out_re.unsafe_ptr()) + quantum_state_out_im_tensor = LayoutTensor[ + mut=True, dtype, state_vector_layout + ](quantum_state_out_im.unsafe_ptr()) + + control_bits_circuit_tensor = LayoutTensor[ + mut=False, DType.int32, circuit_control_bits_layout + ](control_bits_circuit.unsafe_ptr()) + current_control_gate_circuit_tensor = LayoutTensor[ + mut=True, DType.int32, Layout.row_major(1) + ](current_control_gate_circuit.unsafe_ptr()) + + # -- Apply circuit operations -- # + + current_buffer: Int = 0 + for gate in self.circuit.gates: # Iterate over the gates in the circuit + if gate.symbol == _SEPARATOR.symbol: + continue + elif gate.symbol == SWAP.symbol: + print("SWAP gate not implemented in GPU version.") + else: + # Apply the next gate + if current_buffer == 0: + ctx.enqueue_function[ + qubit_wise_multiply_inplace_gpu[ + state_vector_size=state_vector_size, + gate_set_size=gate_set_size, + circuit_number_control_gates=circuit_number_control_gates, + number_control_bits=0, + ] + ]( + gate_set_re_tensor, + gate_set_im_tensor, + gate_set_dic[gate.symbol], + GATE_SIZE, # Assumming single target qubit + gate.target_qubits[0], # target_qubit + quantum_state_re_tensor, + quantum_state_im_tensor, + num_qubits, # number_qubits + quantum_state_out_re_tensor, + quantum_state_out_im_tensor, + control_bits_circuit_tensor, + current_control_gate_circuit_tensor, + grid_dim=blocks_per_grid, + block_dim=threads_per_block, + ) + current_buffer = 1 + else: + ctx.enqueue_function[ + qubit_wise_multiply_inplace_gpu[ + state_vector_size=state_vector_size, + gate_set_size=gate_set_size, + circuit_number_control_gates=circuit_number_control_gates, + number_control_bits=0, + ] + ]( + gate_set_re_tensor, + gate_set_im_tensor, + gate_set_dic[gate.symbol], + GATE_SIZE, # Assumming single target qubit + gate.target_qubits[0], # target_qubit + quantum_state_out_re_tensor, + quantum_state_out_im_tensor, + num_qubits, # number_qubits + quantum_state_re_tensor, + quantum_state_im_tensor, + control_bits_circuit_tensor, + current_control_gate_circuit_tensor, + grid_dim=blocks_per_grid, + block_dim=threads_per_block, + ) + current_buffer = 0 + + if current_buffer == 0: + with quantum_state_re.map_to_host() as host_re, quantum_state_im.map_to_host() as host_im: + for i in range(state_vector_size): + quantum_state[i] = ComplexFloat32(host_re[i], host_im[i]) + else: + with quantum_state_out_re.map_to_host() as host_re, quantum_state_out_im.map_to_host() as host_im: + for i in range(state_vector_size): + quantum_state[i] = ComplexFloat32(host_re[i], host_im[i]) + + if self.verbose and self.verbose_step_size == ShowOnlyEnd: + print("Final quantum state:\n", quantum_state) + + return quantum_state From ab216e091853d6209fde0e3e13a13c92a9055f6d Mon Sep 17 00:00:00 2001 From: ttrenty <154608953+ttrenty@users.noreply.github.com> Date: Sun, 29 Jun 2025 13:40:17 -0600 Subject: [PATCH 2/2] fix: add has_accelerator() before DeviceContext() call --- src/abstractions/simulator.mojo | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/abstractions/simulator.mojo b/src/abstractions/simulator.mojo index fa435de..8a5430f 100644 --- a/src/abstractions/simulator.mojo +++ b/src/abstractions/simulator.mojo @@ -125,6 +125,8 @@ struct StateVectorSimulator[ self.verbose = verbose self.verbose_step_size = verbose_step_size if self.use_gpu_if_available: + + @parameter if not has_accelerator(): print( "No compatible GPU found. Falling back to CPU simulation." @@ -282,13 +284,16 @@ struct StateVectorSimulator[ print("Initial state:\n", self.initial_state) if self.use_gpu_if_available: - try: - ctx = DeviceContext() - final_state = self.run_gpu(ctx) - return final_state - except e: - print("Failed to create GPU context:", e) - print("Falling back to CPU simulation.") + + @parameter + if has_accelerator(): + try: + ctx = DeviceContext() + final_state = self.run_gpu(ctx) + return final_state + except e: + print("Failed to create GPU context:", e) + print("Falling back to CPU simulation.") # Start with the initial state quantum_state: StateVector = self.initial_state