[Optimizer] Add an option to store each circuit transformation step (#174)

* [Optimizer] Add an option to store each circuit transformation step

* Update test_optimization to test optimization steps

* Fix ECC Set name

* Rename to optimization_steps and let it able to run from any dir

* Fix bugs

Author: xumingkuan

src/quartz/circuitseq/circuitseq.cpp

@@ -187,8 +187,8 @@ bool CircuitSeq::less_than(const CircuitSeq &other) const {
           // so this circuit is greater.
           return false;
         }
-        assert(this_ptr->output_gates->size() == 1);
-        assert(other_ptr->output_gates->size() == 1);
+        assert(this_ptr->output_gates.size() == 1);
+        assert(other_ptr->output_gates.size() == 1);
         auto this_gate = this_ptr->output_gates[0];
         auto other_gate = other_ptr->output_gates[0];
         if (!compare_outcome.has_value()) {

src/quartz/tasograph/tasograph.cpp

@@ -1557,7 +1557,8 @@ void Graph::draw_circuit(const std::string &src_file_name,
 std::shared_ptr<Graph>
 Graph::greedy_optimize(Context *ctx, const std::string &equiv_file_name,
                        bool print_message,
-                       std::function<float(Graph *)> cost_function) {
+                       std::function<float(Graph *)> cost_function,
+                       const std::string &store_all_steps_file_prefix) {
   if (cost_function == nullptr) {
     cost_function = [](Graph *graph) { return graph->total_cost(); };
   }
@@ -1607,6 +1608,11 @@ Graph::greedy_optimize(Context *ctx, const std::string &equiv_file_name,
   bool optimized_in_this_iteration;
   std::vector<Op> all_nodes;
   optimized_graph->topology_order_ops(all_nodes);
+  int step_count = 0;
+  if (!store_all_steps_file_prefix.empty()) {
+    to_qasm(store_all_steps_file_prefix + "0.qasm", /*print_result=*/false,
+            /*print_guid=*/false);
+  }
   do {
     optimized_in_this_iteration = false;
     for (auto xfer : xfers) {
@@ -1623,6 +1629,13 @@ Graph::greedy_optimize(Context *ctx, const std::string &equiv_file_name,
             optimized_graph->topology_order_ops(all_nodes);
             optimized_this_xfer = true;
             optimized_in_this_iteration = true;
+            if (!store_all_steps_file_prefix.empty()) {
+              step_count++;
+              optimized_graph->to_qasm(store_all_steps_file_prefix +
+                                           std::to_string(step_count) + ".qasm",
+                                       /*print_result=*/false,
+                                       /*print_guid=*/false);
+            }
             // Since |all_nodes| has changed, we cannot continue this loop.
             break;
           }
@@ -1638,6 +1651,14 @@ Graph::greedy_optimize(Context *ctx, const std::string &equiv_file_name,
               << " to " << optimized_cost << std::endl;
   }
 
+  if (!store_all_steps_file_prefix.empty()) {
+    // Store the number of steps.
+    std::ofstream fout(store_all_steps_file_prefix + ".txt");
+    assert(fout.is_open());
+    fout << step_count << std::endl;
+    fout.close();
+  }
+
   return optimized_graph;
 }
 
@@ -1901,7 +1922,8 @@ std::shared_ptr<Graph>
 Graph::optimize(Context *ctx, const std::string &equiv_file_name,
                 const std::string &circuit_name, bool print_message,
                 std::function<float(Graph *)> cost_function,
-                double cost_upper_bound, double timeout) {
+                double cost_upper_bound, double timeout,
+                const std::string &store_all_steps_file_prefix) {
   if (cost_function == nullptr) {
     cost_function = [](Graph *graph) { return graph->total_cost(); };
   }
@@ -1947,24 +1969,22 @@ Graph::optimize(Context *ctx, const std::string &equiv_file_name,
   if (cost_upper_bound == -1) {
     cost_upper_bound = total_cost() * 1.05;
   }
-  auto log_file_name =
-      equiv_file_name.substr(0, std::max(0, (int)equiv_file_name.size() - 21)) +
-      circuit_name + ".log";
   auto preprocessed_graph =
-      greedy_optimize(ctx, equiv_file_name, print_message, cost_function);
-  //   return preprocessed_graph->optimize(xfers, cost_upper_bound,
-  //   circuit_name,
-  //                                       log_file_name, print_message,
-  //                                       cost_function, timeout);
-  return preprocessed_graph->optimize(xfers, cost_upper_bound, circuit_name, "",
-                                      print_message, cost_function, timeout);
+      greedy_optimize(ctx, equiv_file_name, print_message, cost_function,
+                      store_all_steps_file_prefix);
+  return preprocessed_graph->optimize(
+      xfers, cost_upper_bound, circuit_name, /*log_file_name=*/"",
+      print_message, cost_function, timeout, store_all_steps_file_prefix,
+      /*continue_storing_all_steps=*/true);
 }
 
 std::shared_ptr<Graph>
 Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
                 const std::string &circuit_name,
                 const std::string &log_file_name, bool print_message,
-                std::function<float(Graph *)> cost_function, double timeout) {
+                std::function<float(Graph *)> cost_function, double timeout,
+                const std::string &store_all_steps_file_prefix,
+                bool continue_storing_all_steps) {
   if (cost_function == nullptr) {
     cost_function = [](Graph *graph) { return graph->total_cost(); };
   }
@@ -1985,11 +2005,27 @@ Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
   if (print_message) {
     if (!log_file_name.empty()) {
       fout = fopen(log_file_name.c_str(), "w");
+      assert(fout);
     } else {
       fout = stdout;
     }
   }
 
+  // Information necessary to store each step
+  std::unordered_map<Graph *, std::shared_ptr<Graph>> previous_graph;
+  int step_count = 0;
+  if (!store_all_steps_file_prefix.empty()) {
+    if (continue_storing_all_steps) {
+      std::ifstream fin(store_all_steps_file_prefix + ".txt");
+      assert(fin.is_open());
+      fin >> step_count;
+      fin.close();
+    } else {
+      to_qasm(store_all_steps_file_prefix + "0.qasm", /*print_result=*/false,
+              /*print_guid=*/false);
+    }
+  }
+
   // TODO: make these numbers configurable
   constexpr int kMaxNumCandidates = 2000;
   constexpr int kShrinkToNumCandidates = 1000;
@@ -2009,6 +2045,11 @@ Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
       cost_count[cost_function(candidate.get())]++;
       if (new_candidates.size() < kShrinkToNumCandidates) {
         new_candidates.push(candidate);
+      } else {
+        if (!store_all_steps_file_prefix.empty()) {
+          // no need to record history of removed graphs
+          previous_graph.erase(candidate.get());
+        }
       }
       candidates.pop();
     }
@@ -2030,6 +2071,7 @@ Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
     }
   };
 
+  bool hit_timeout = false;
   while (!candidates.empty()) {
     auto graph = candidates.top();
     candidates.pop();
@@ -2048,7 +2090,8 @@ Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
             timeout) {
           std::cout << "Timeout. Program terminated. Best cost is " << best_cost
                     << std::endl;
-          return best_graph;
+          hit_timeout = true;
+          break;
         }
         if (new_graph == nullptr)
           continue;
@@ -2057,33 +2100,66 @@ Graph::optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
         auto new_cost = cost_function(new_graph.get());
         if (new_cost > cost_upper_bound)
           continue;
-        if (hashmap.find(new_hash) == hashmap.end()) {
-          hashmap.insert(new_hash);
-          candidates.push(new_graph);
-          if (candidates.size() > kMaxNumCandidates) {
-            shrink_candidates();
-          }
-          if (new_cost < best_cost) {
-            best_cost = new_cost;
-            best_graph = new_graph;
-          }
-        } else
+        if (hashmap.find(new_hash) != hashmap.end()) {
           continue;
+        }
+        hashmap.insert(new_hash);
+        candidates.push(new_graph);
+        if (!store_all_steps_file_prefix.empty()) {
+          // record history
+          previous_graph[new_graph.get()] = graph;
+        }
+        if (candidates.size() > kMaxNumCandidates) {
+          shrink_candidates();
+        }
+        if (new_cost < best_cost) {
+          best_cost = new_cost;
+          best_graph = new_graph;
+        }
       }
+      if (hit_timeout) {
+        break;
+      }
+    }
+    if (hit_timeout) {
+      break;
     }
 
     auto end = std::chrono::steady_clock::now();
     if (print_message) {
-      fprintf(fout,
-              "[%s] Best cost: %f\tcandidate number: %d\tafter %.3f seconds.\n",
-              circuit_name.c_str(), best_cost, candidates.size(),
-              (double)std::chrono::duration_cast<std::chrono::milliseconds>(
-                  end - start)
-                      .count() /
-                  1000.0);
+      fprintf(
+          fout,
+          "[%s] Best cost: %f\tcandidate number: %zu\tafter %.3f seconds.\n",
+          circuit_name.c_str(), best_cost, candidates.size(),
+          (double)std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                        start)
+                  .count() /
+              1000.0);
       fflush(fout);
     }
   }
+
+  if (!store_all_steps_file_prefix.empty()) {
+    std::vector<Graph *> steps(1, best_graph.get());
+    while (previous_graph.count(steps.back()) > 0) {
+      // there is a previous graph
+      steps.push_back(previous_graph[steps.back()].get());
+    }
+    // no need to save the initial graph again
+    for (int i = (int)steps.size() - 2; i >= 0; i--) {
+      step_count++;
+      steps[i]->to_qasm(store_all_steps_file_prefix +
+                            std::to_string(step_count) + ".qasm",
+                        /*print_result=*/false,
+                        /*print_guid=*/false);
+    }
+
+    // Store the number of steps.
+    std::ofstream fout_step(store_all_steps_file_prefix + ".txt");
+    fout_step << step_count << std::endl;
+    fout_step.close();
+  }
+
   return best_graph;
 }
 

src/quartz/tasograph/tasograph.h

@@ -208,12 +208,14 @@ class Graph {
    * @param equiv_file_name The ECC set file name.
    * @param print_message To output the log to the screen or not.
    * @param cost_function The cost function.
+   * @param store_all_steps_file_prefix If not empty, store each circuit
+   * transformation step in a file with the corresponding file prefix.
    * @return The optimized circuit.
    */
-  std::shared_ptr<Graph>
-  greedy_optimize(Context *ctx, const std::string &equiv_file_name,
-                  bool print_message,
-                  std::function<float(Graph *)> cost_function = nullptr);
+  std::shared_ptr<Graph> greedy_optimize(
+      Context *ctx, const std::string &equiv_file_name, bool print_message,
+      std::function<float(Graph *)> cost_function = nullptr,
+      const std::string &store_all_steps_file_prefix = std::string());
   std::shared_ptr<Graph>
   optimize_legacy(float alpha, int budget, bool print_subst, Context *ctx,
                   const std::string &equiv_file_name,
@@ -223,7 +225,7 @@ class Graph {
                   int timeout = 86400 /*1 day*/);
 
   /**
-   * Optimize this circuit.
+   * Optimize this circuit with a greedy phase at the beginning.
    * @param ctx The context variable.
    * @param equiv_file_name The ECC set file name.
    * @param circuit_name The circuit name shown in the log.
@@ -232,17 +234,20 @@ class Graph {
    * file name.
    * @param cost_function The cost function for the search.
    * @param cost_upper_bound The maximum cost of the circuits to be searched.
-   * @param timeout Timeout in seconds.
+   * @param timeout Timeout in seconds, for the search phase.
+   * @param store_all_steps_file_prefix If not empty, store each circuit
+   * transformation step in a file with the corresponding file prefix.
    * @return The optimized circuit.
    */
   std::shared_ptr<Graph>
   optimize(Context *ctx, const std::string &equiv_file_name,
            const std::string &circuit_name, bool print_message,
            std::function<float(Graph *)> cost_function = nullptr,
            double cost_upper_bound = -1 /*default = current cost * 1.05*/,
-           double timeout = 3600 /*1 hour*/);
+           double timeout = 3600 /*1 hour*/,
+           const std::string &store_all_steps_file_prefix = std::string());
   /**
-   * Optimize this circuit.
+   * Optimize this circuit without a greedy phase.
    * @param xfers The circuit transformations.
    * @param cost_upper_bound The maximum cost of the circuits to be searched.
    * @param circuit_name The circuit name shown in the log.
@@ -251,19 +256,26 @@ class Graph {
    * @param print_message To output the log or not.
    * @param cost_function The cost function for the search.
    * @param timeout Timeout in seconds.
+   * @param store_all_steps_file_prefix If not empty, store each circuit
+   * transformation step in a file with the corresponding file prefix.
+   * @param continue_storing_all_steps If true, there was a greedy phase
+   * before calling this function with the same |store_all_steps_file_prefix|.
+   * We should continue the numbering in this case.
    * @return The optimized circuit.
    */
   std::shared_ptr<Graph>
   optimize(const std::vector<GraphXfer *> &xfers, double cost_upper_bound,
            const std::string &circuit_name, const std::string &log_file_name,
            bool print_message,
            std::function<float(Graph *)> cost_function = nullptr,
-           double timeout = 3600 /*1 hour*/);
+           double timeout = 3600 /*1 hour*/,
+           const std::string &store_all_steps_file_prefix = std::string(),
+           bool continue_storing_all_steps = false);
   void constant_and_rotation_elimination();
   void rotation_merging(GateType target_rotation);
-  std::string to_qasm(bool print_result = false, bool print_id = false) const;
+  std::string to_qasm(bool print_result = false, bool print_guid = false) const;
   void to_qasm(const std::string &save_filename, bool print_result,
-               bool print_id) const;
+               bool print_guid) const;
   template <class _CharT, class _Traits>
   static std::shared_ptr<Graph>
   _from_qasm_stream(Context *ctx,

src/test/CMakeLists.txt

@@ -10,8 +10,7 @@ file(GLOB_RECURSE TEST_CONTEXT_SHIFT "test_context_shift.cpp")
 file(GLOB_RECURSE TEST_EQUIV_SET "test_equivalence_set.cpp")
 file(GLOB_RECURSE TEST_GENERATOR "test_generator.cpp")
 file(GLOB_RECURSE TEST_GRAPH_TO_QASM "test_graph_to_qasm.cpp")
-file(GLOB_RECURSE TEST_OPTIMIZATION "test_optimization.cpp")
-file(GLOB_RECURSE TEST_OPTIMIZATION_SA "test_optimization_sa.cpp")
+file(GLOB_RECURSE TEST_OPTIMIZATION_STEPS "test_optimization_steps.cpp")
 file(GLOB_RECURSE TEST_PHASE_SHIFT "test_phase_shift.cpp")
 file(GLOB_RECURSE TEST_PRUNING "test_pruning.cpp")
 file(GLOB_RECURSE TEST_QUARTZ "test_quartz.cpp")
@@ -46,8 +45,7 @@ add_executable(test_context_shift ${TEST_CONTEXT_SHIFT} )
 add_executable(test_equiv_set ${TEST_EQUIV_SET} )
 add_executable(test_generator ${TEST_GENERATOR} )
 add_executable(test_graph_to_qasm ${TEST_GRAPH_TO_QASM} )
-add_executable(test_optimization ${TEST_OPTIMIZATION} )
-add_executable(test_optimization_sa ${TEST_OPTIMIZATION_SA} )
+add_executable(test_optimization_steps ${TEST_OPTIMIZATION_STEPS} )
 add_executable(test_phase_shift ${TEST_PHASE_SHIFT} )
 add_executable(test_pruning ${TEST_PRUNING} )
 add_executable(test_quartz ${TEST_QUARTZ} )