Add gate-batched cache-local runner across SIMD backends

apps/BUILD

@@ -32,6 +32,16 @@ cc_binary(
     ],
 )
 
+cc_binary(
+    name = "qsim_gate_batch",
+    srcs = ["qsim_gate_batch.cc"],
+    copts = gsframe_copts,
+    deps = [
+        "//lib:run_qsim_gate_batch",
+        "//lib:run_qsim_lib",
+    ],
+)
+
 cc_binary(
     name = "qsim_von_neumann",
     srcs = ["qsim_von_neumann.cc"],

apps/Makefile

@@ -1,6 +1,8 @@
 CXX_TARGETS = $(shell find . -maxdepth 1 -name '*.cc')
 CXX_TARGETS := $(CXX_TARGETS:%.cc=%.x)
 
+OUTPUT_DIR ?= .
+
 CUDA_TARGETS = $(shell find . -maxdepth 1 -name '*cuda.cu')
 CUDA_TARGETS := $(CUDA_TARGETS:%cuda.cu=%cuda.x)
 
@@ -16,6 +18,12 @@ HIP_TARGETS := $(HIP_TARGETS:%cuda.cu=%hip.x)
 .PHONY: qsim
 qsim: $(CXX_TARGETS)
 
+.PHONY: gate-batch
+gate-batch:
+	mkdir -p "$(OUTPUT_DIR)"
+	$(CXX) -o "$(OUTPUT_DIR)/qsim_gate_batch.x" qsim_gate_batch.cc $(CXXFLAGS) $(ARCHFLAGS)
+	$(CXX) -o "$(OUTPUT_DIR)/qsim_base.x" qsim_base.cc $(CXXFLAGS) $(ARCHFLAGS)
+
 .PHONY: qsim-cuda
 qsim-cuda: $(CUDA_TARGETS)
 

apps/qsim_gate_batch.cc

@@ -0,0 +1,150 @@
+// Benchmark app for the proposal-faithful gate-batch runner.
+// -f is the per-batch max_fused_size (0 = apply
+// raw gates without fusing, 2-3 = proposal's suggestion).
+
+#include <unistd.h>
+
+#ifdef _OPENMP
+# include <omp.h>
+#endif
+
+#include <limits>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "../lib/circuit_qsim_parser.h"
+#include "../lib/cooperative_for.h"
+#include "../lib/cpu_thread_topology.h"
+#include "../lib/formux.h"
+#include "../lib/fuser_mqubit.h"
+#include "../lib/gates_qsim.h"
+#include "../lib/io_file.h"
+#include "../lib/operation.h"
+#include "../lib/run_qsim_gate_batch.h"
+#include "../lib/seqfor.h"
+#include "../lib/simmux.h"
+#include "../lib/util_cpu.h"
+
+struct Options {
+  std::string circuit_file;
+  unsigned maxtime = std::numeric_limits<unsigned>::max();
+  unsigned seed = 1;
+  unsigned num_threads = 1;
+  unsigned inner_threads = 1;
+  unsigned max_fused_size = 3;
+  unsigned block_qubits = 19;
+  unsigned verbosity = 0;
+};
+
+Options GetOptions(int argc, char* argv[]) {
+  constexpr char usage[] = "usage:\n  ./qsim_gate_batch -c circuit "
+                           "-d maxtime -s seed -t threads "
+                           "-f max_fused_size -l block_qubits "
+                           "-i inner_threads -v verbosity\n";
+  Options opt;
+  int k;
+  while ((k = getopt(argc, argv, "c:d:s:t:f:l:i:v:")) != -1) {
+    switch (k) {
+      case 'c': opt.circuit_file = optarg; break;
+      case 'd': opt.maxtime = std::atoi(optarg); break;
+      case 's': opt.seed = std::atoi(optarg); break;
+      case 't': opt.num_threads = std::atoi(optarg); break;
+      case 'f': opt.max_fused_size = std::atoi(optarg); break;
+      case 'l': opt.block_qubits = std::atoi(optarg); break;
+      case 'i': opt.inner_threads = std::atoi(optarg); break;
+      case 'v': opt.verbosity = std::atoi(optarg); break;
+      default: qsim::IO::errorf(usage); exit(1);
+    }
+  }
+  if (opt.circuit_file.empty()) {
+    qsim::IO::errorf(usage);
+    exit(1);
+  }
+  return opt;
+}
+
+template <typename StateSpace, typename QubitMappedState>
+void PrintAmplitudes(unsigned num_qubits, const StateSpace& state_space,
+                     const QubitMappedState& state) {
+  static constexpr char const* bits[8] = {
+    "000", "001", "010", "011", "100", "101", "110", "111",
+  };
+  uint64_t size = std::min(uint64_t{8}, uint64_t{1} << num_qubits);
+  unsigned s = 3 - std::min(unsigned{3}, num_qubits);
+  for (uint64_t i = 0; i < size; ++i) {
+    auto a = state.GetAmpl(state_space, i);
+    qsim::IO::messagef("%s:%16.8g%16.8g%16.8g\n",
+                       bits[i] + s, std::real(a), std::imag(a), std::norm(a));
+  }
+}
+
+int main(int argc, char* argv[]) {
+  using namespace qsim;
+
+  auto opt = GetOptions(argc, argv);
+  std::vector<unsigned> team_thread_cpus;
+  if (opt.inner_threads > 1) {
+#ifndef _OPENMP
+    IO::errorf("cannot configure SMT teams: OpenMP is not enabled.\n");
+    return 1;
+#else
+    std::string topology_error;
+    const auto topology = CpuThreadTopology::Discover();
+    if (!topology.BuildTeamCpuOrder(opt.num_threads, opt.inner_threads,
+                                    team_thread_cpus, topology_error)) {
+      IO::errorf("cannot configure SMT teams: %s.\n",
+                 topology_error.c_str());
+      return 1;
+    }
+#endif
+  }
+
+#ifdef _OPENMP
+  omp_set_num_threads(opt.num_threads);
+#endif
+
+  Circuit<Operation<float>> circuit;
+  if (!CircuitQsimParser<IOFile>::FromFile(opt.maxtime, opt.circuit_file,
+                                           circuit)) {
+    return 1;
+  }
+
+  struct Factory {
+    Factory(unsigned num_threads) : num_threads(num_threads) {}
+    using Simulator = qsim::Simulator<For>;
+    using StateSpace = Simulator::StateSpace;
+    StateSpace CreateStateSpace() const { return StateSpace(num_threads); }
+    Simulator CreateSimulator() const { return Simulator(num_threads); }
+    unsigned num_threads;
+  };
+
+  using Simulator = Factory::Simulator;
+  using StateSpace = Simulator::StateSpace;
+  using State = StateSpace::State;
+  using Fuser = MultiQubitGateFuser<IO>;
+  using SeqSimulator = qsim::Simulator<CooperativeFor>;
+  using Runner = QSimGateBatchRunner<IO, Fuser, Factory, SeqSimulator>;
+
+  StateSpace state_space = Factory(opt.num_threads).CreateStateSpace();
+  QubitMappedState<State> state(state_space.Create(circuit.num_qubits));
+  if (state_space.IsNull(state.state)) {
+    IO::errorf("not enough memory: is the number of qubits too large?\n");
+    return 1;
+  }
+  state_space.SetStateZero(state.state);
+
+  Runner::Parameter param;
+  param.max_fused_size = opt.max_fused_size;
+  param.block_qubits = opt.block_qubits;
+  param.num_threads = opt.num_threads;
+  param.inner_threads = opt.inner_threads;
+  param.team_thread_cpus = std::move(team_thread_cpus);
+  param.verbosity = opt.verbosity;
+
+  if (Runner::Run(param, Factory(opt.num_threads), circuit, state)) {
+    PrintAmplitudes(circuit.num_qubits, state_space, state);
+  }
+
+  return 0;
+}

lib/BUILD

@@ -564,6 +564,23 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "run_qsim_gate_batch",
+    hdrs = [
+        "cooperative_for.h",
+        "cpu_thread_topology.h",
+        "qubit_mapped_state.h",
+        "qubit_remap.h",
+        "run_qsim_gate_batch.h",
+    ],
+    deps = [
+        ":gate",
+        ":matrix",
+        ":operation_base",
+        ":util",
+    ],
+)
+
 cc_library(
     name = "run_qsimh",
     hdrs = ["run_qsimh.h"],

lib/cooperative_for.h

@@ -0,0 +1,49 @@
+// Copyright 2026 Google LLC. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef COOPERATIVE_FOR_H_
+#define COOPERATIVE_FOR_H_
+
+#include <cstdint>
+
+namespace qsim {
+
+// Splits one simulator kernel between the SMT siblings assigned to the same
+// state block. ExecuteGateBatchOnBlocks configures the lane before invoking
+// the simulator and provides the barrier between consecutive gates.
+struct CooperativeFor {
+  explicit CooperativeFor(unsigned num_threads) { (void) num_threads; }
+
+  static void Configure(unsigned num_threads, unsigned thread_id) {
+    team_size_ = num_threads;
+    team_thread_id_ = thread_id;
+  }
+
+  template <typename Function, typename... Args>
+  static void Run(uint64_t size, Function&& func, Args&&... args) {
+    const auto begin = size * team_thread_id_ / team_size_;
+    const auto end = size * (team_thread_id_ + 1) / team_size_;
+    for (uint64_t i = begin; i < end; ++i) {
+      func(team_size_, team_thread_id_, i, args...);
+    }
+  }
+
+ private:
+  inline static thread_local unsigned team_size_ = 1;
+  inline static thread_local unsigned team_thread_id_ = 0;
+};
+
+}  // namespace qsim
+
+#endif  // COOPERATIVE_FOR_H_

lib/cpu_thread_topology.h

@@ -0,0 +1,134 @@
+// Copyright 2026 Google LLC. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CPU_THREAD_TOPOLOGY_H_
+#define CPU_THREAD_TOPOLOGY_H_
+
+#include <algorithm>
+#include <fstream>
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#ifdef __linux__
+#include <sched.h>
+#endif
+
+namespace qsim {
+
+struct CpuCore {
+  unsigned package_id;
+  unsigned core_id;
+  std::vector<unsigned> logical_cpus;
+};
+
+class CpuThreadTopology {
+ public:
+  static CpuThreadTopology Discover() {
+    CpuThreadTopology topology;
+
+#ifdef __linux__
+    cpu_set_t allowed_cpus;
+    CPU_ZERO(&allowed_cpus);
+    if (sched_getaffinity(0, sizeof(allowed_cpus), &allowed_cpus) != 0) {
+      topology.error_ = "cannot read the process CPU affinity";
+      return topology;
+    }
+
+    std::map<std::pair<unsigned, unsigned>, std::vector<unsigned>> cores;
+    for (unsigned cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
+      if (!CPU_ISSET(cpu, &allowed_cpus)) continue;
+
+      unsigned package_id;
+      unsigned core_id;
+      if (!ReadTopologyValue(cpu, "physical_package_id", package_id) ||
+          !ReadTopologyValue(cpu, "core_id", core_id)) {
+        topology.error_ = "cannot read Linux CPU topology from sysfs";
+        return topology;
+      }
+      cores[{package_id, core_id}].push_back(cpu);
+    }
+
+    for (auto& [core, logical_cpus] : cores) {
+      std::sort(logical_cpus.begin(), logical_cpus.end());
+      topology.cores_.push_back(
+          CpuCore{core.first, core.second, std::move(logical_cpus)});
+    }
+#else
+    topology.error_ = "automatic CPU topology discovery requires Linux";
+#endif
+
+    return topology;
+  }
+
+  bool BuildTeamCpuOrder(unsigned num_threads, unsigned threads_per_team,
+                         std::vector<unsigned>& thread_cpus,
+                         std::string& error) const {
+    thread_cpus.clear();
+    error.clear();
+    if (!error_.empty()) {
+      error = error_;
+      return false;
+    }
+    if (threads_per_team == 0 || num_threads % threads_per_team != 0) {
+      error = "thread count must be divisible by threads per team";
+      return false;
+    }
+
+    const auto num_teams = num_threads / threads_per_team;
+    for (const auto& core : cores_) {
+      if (core.logical_cpus.size() < threads_per_team) continue;
+      for (unsigned lane = 0; lane < threads_per_team; ++lane) {
+        thread_cpus.push_back(core.logical_cpus[lane]);
+      }
+      if (thread_cpus.size() == num_threads) return true;
+    }
+
+    error = "not enough physical cores with the requested SMT siblings";
+    thread_cpus.clear();
+    return false;
+  }
+
+ private:
+#ifdef __linux__
+  static bool ReadTopologyValue(unsigned cpu, const char* name,
+                                unsigned& value) {
+    const auto path = "/sys/devices/system/cpu/cpu" +
+                      std::to_string(cpu) + "/topology/" + name;
+    std::ifstream input(path);
+    return bool(input >> value);
+  }
+#endif
+
+  std::vector<CpuCore> cores_;
+  std::string error_;
+};
+
+inline bool PinCurrentThreadToCpu(unsigned cpu) {
+#ifdef __linux__
+  if (cpu >= CPU_SETSIZE) return false;
+  cpu_set_t affinity;
+  CPU_ZERO(&affinity);
+  CPU_SET(cpu, &affinity);
+  return sched_setaffinity(0, sizeof(affinity), &affinity) == 0;
+#else
+  (void) cpu;
+  return false;
+#endif
+}
+
+}  // namespace qsim
+
+#endif  // CPU_THREAD_TOPOLOGY_H_