Improve QC COO canonicalization performance

cpp/include/cuopt/mathematical_optimization/cpu_optimization_problem.hpp

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <cuopt/mathematical_optimization/io/mps_data_model.hpp>
 #include <cuopt/mathematical_optimization/optimization_problem_interface.hpp>
 
 #include <raft/core/handle.hpp>
@@ -45,6 +46,11 @@ class cpu_optimization_problem_t : public optimization_problem_interface_t<i_t,
   using typename optimization_problem_interface_t<i_t, f_t>::quadratic_constraint_t;
 
   cpu_optimization_problem_t();
+  cpu_optimization_problem_t(const cpu_optimization_problem_t<i_t, f_t>& other);
+  cpu_optimization_problem_t(cpu_optimization_problem_t<i_t, f_t>&& other) noexcept;
+  cpu_optimization_problem_t& operator=(const cpu_optimization_problem_t<i_t, f_t>& other);
+  cpu_optimization_problem_t& operator=(cpu_optimization_problem_t<i_t, f_t>&& other) noexcept;
+  ~cpu_optimization_problem_t();
 
   // Setters
   void set_maximize(bool maximize) override;
@@ -212,6 +218,8 @@ class cpu_optimization_problem_t : public optimization_problem_interface_t<i_t,
   std::vector<f_t> Q_values_;
 
   std::vector<quadratic_constraint_t> quadratic_constraints_{};
+  // Scratch storage reused across add_quadratic_constraint calls. Not thread-safe.
+  std::unique_ptr<io::coo_canonicalization_workspace_t<i_t, f_t>> qc_coo_workspace_;
 
   std::vector<f_t> variable_lower_bounds_;
   std::vector<f_t> variable_upper_bounds_;

cpp/include/cuopt/mathematical_optimization/io/mps_data_model.hpp

@@ -15,6 +15,19 @@
 
 namespace cuopt::mathematical_optimization::io {
 
+template <typename i_t, typename f_t>
+class coo_canonicalization_workspace_t;
+
+/**
+ * @brief Canonicalize a symmetric matrix in COO form to upper-triangular storage
+ *        using reusable workspace storage.
+ */
+template <typename i_t, typename f_t>
+void canonicalize_coo_matrix(std::vector<i_t>& rows,
+                             std::vector<i_t>& cols,
+                             std::vector<f_t>& vals,
+                             coo_canonicalization_workspace_t<i_t, f_t>& workspace);
+
 /**
  * @brief A representation of a linear programming (LP) optimization problem
  *
@@ -278,7 +291,8 @@ class mps_data_model_t {
                                    f_t rhs_value,
                                    std::span<const f_t> vals,
                                    std::span<const i_t> rows,
-                                   std::span<const i_t> cols);
+                                   std::span<const i_t> cols,
+                                   coo_canonicalization_workspace_t<i_t, f_t>& workspace);
 
   const std::vector<quadratic_constraint_t>& get_quadratic_constraints() const;
 
@@ -397,8 +411,7 @@ class mps_data_model_t {
  *
  * @param constraints Quadratic constraints whose rows, cols, and vals are updated.
  */
-template <typename i_t, typename f_t>
-void canonicalize_quadratic_constraints(
-  std::vector<typename mps_data_model_t<i_t, f_t>::quadratic_constraint_t>& constraints);
+template <typename qc_t>
+void canonicalize_quadratic_constraints(std::vector<qc_t>& constraints);
 
 }  // namespace cuopt::mathematical_optimization::io

cpp/include/cuopt/mathematical_optimization/optimization_problem.hpp

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <cuopt/mathematical_optimization/io/mps_data_model.hpp>
 #include <cuopt/mathematical_optimization/optimization_problem_interface.hpp>
 #include <cuopt/mathematical_optimization/utilities/internals.hpp>
 
@@ -117,8 +118,9 @@ class optimization_problem_t : public optimization_problem_interface_t<i_t, f_t>
 
   explicit optimization_problem_t(raft::handle_t const* handle_ptr);
   optimization_problem_t(const optimization_problem_t<i_t, f_t>& other);
-  optimization_problem_t(optimization_problem_t<i_t, f_t>&&) noexcept            = default;
-  optimization_problem_t& operator=(optimization_problem_t<i_t, f_t>&&) noexcept = default;
+  optimization_problem_t(optimization_problem_t<i_t, f_t>&& other) noexcept;
+  optimization_problem_t& operator=(optimization_problem_t<i_t, f_t>&& other) noexcept;
+  ~optimization_problem_t();
 
   std::vector<internals::base_solution_callback_t*> mip_callbacks_;
 
@@ -411,6 +413,8 @@ class optimization_problem_t : public optimization_problem_interface_t<i_t, f_t>
 
   /** QCQP: quadratic constraints **/
   std::vector<quadratic_constraint_t> quadratic_constraints_{};
+  // Scratch storage reused across add_quadratic_constraint calls. Not thread-safe.
+  std::unique_ptr<io::coo_canonicalization_workspace_t<i_t, f_t>> qc_coo_workspace_;
 
   rmm::device_uvector<f_t> variable_lower_bounds_;
   rmm::device_uvector<f_t> variable_upper_bounds_;

cpp/include/cuopt/mathematical_optimization/optimization_problem_utils.hpp

@@ -294,7 +294,7 @@ void populate_from_data_model_view(
   // Raw Q COO from data_model_view is canonicalized here before solver storage.
   if (data_model->has_quadratic_constraints()) {
     auto qcs = data_model->get_quadratic_constraints();
-    io::canonicalize_quadratic_constraints<i_t, f_t>(qcs);
+    io::canonicalize_quadratic_constraints(qcs);
     problem->set_quadratic_constraints(std::move(qcs));
   }
 }

cpp/src/grpc/codegen/generate_conversions.py

@@ -2215,12 +2215,12 @@ def _gen_proto_to_problem(registry, indent="  "):
                 f'{b["name"]} size mismatch");'
             )
             lines.append(f"{ind}    }}")
+        lines.append(f"{ind}    _entries.push_back(std::move(_entry));")
+        lines.append(f"{ind}  }}")
         if name == "quadratic_constraints":
             lines.append(
-                f"{ind}    io::canonicalize_coo_matrix(_entry.rows, _entry.cols, _entry.vals);"
+                f"{ind}  io::canonicalize_quadratic_constraints(_entries);"
             )
-        lines.append(f"{ind}    _entries.push_back(std::move(_entry));")
-        lines.append(f"{ind}  }}")
         lines.append(f"{ind}  cpu_problem.{setter}(std::move(_entries));")
         lines.append(f"{ind}}}")
 
@@ -2815,12 +2815,12 @@ def _gen_chunked_arrays_to_problem(registry, indent="  "):
                     f'{b["name"]} size mismatch");'
                 )
                 lines.append(f"{ind}    }}")
+            lines.append(f"{ind}    _entries.push_back(std::move(_entry));")
+            lines.append(f"{ind}  }}")
             if name == "quadratic_constraints":
                 lines.append(
-                    f"{ind}    io::canonicalize_coo_matrix(_entry.rows, _entry.cols, _entry.vals);"
+                    f"{ind}  io::canonicalize_quadratic_constraints(_entries);"
                 )
-            lines.append(f"{ind}    _entries.push_back(std::move(_entry));")
-            lines.append(f"{ind}  }}")
             lines.append(f"{ind}  cpu_problem.{setter}(std::move(_entries));")
             lines.append(f"{ind}}}")
 

cpp/src/grpc/codegen/generated/generated_chunked_arrays_to_problem.inc

@@ -194,8 +194,8 @@
       if (_entry.cols.size() != _entry.vals.size()) {
         throw std::invalid_argument("set_quadratic_constraints: cols/vals size mismatch");
       }
-      io::canonicalize_coo_matrix(_entry.rows, _entry.cols, _entry.vals);
       _entries.push_back(std::move(_entry));
     }
+    io::canonicalize_quadratic_constraints(_entries);
     cpu_problem.set_quadratic_constraints(std::move(_entries));
   }

cpp/src/grpc/codegen/generated/generated_proto_to_problem.inc

@@ -96,8 +96,8 @@
       if (_entry.cols.size() != _entry.vals.size()) {
         throw std::invalid_argument("set_quadratic_constraints: cols/vals size mismatch");
       }
-      io::canonicalize_coo_matrix(_entry.rows, _entry.cols, _entry.vals);
       _entries.push_back(std::move(_entry));
     }
+    io::canonicalize_quadratic_constraints(_entries);
     cpu_problem.set_quadratic_constraints(std::move(_entries));
   }

cpp/src/io/lp_parser.cpp

@@ -9,6 +9,7 @@
 
 #include <file_to_string.hpp>
 #include <lp_parser.hpp>
+#include <mps_parser_internal.hpp>
 #include <utilities/error.hpp>
 
 #include <algorithm>
@@ -1471,6 +1472,7 @@ void flush_quadratic_constraints(mps_data_model_t<i_t, f_t>& problem,
                                  const lp_parser_t<i_t, f_t>& parser)
 {
   const i_t linear_row_count = static_cast<i_t>(parser.row_names.size());
+  coo_canonicalization_workspace_t<i_t, f_t> qc_workspace;
   for (i_t k = 0; k < static_cast<i_t>(parser.quadratic_constraint_blocks.size()); k++) {
     const auto& block = parser.quadratic_constraint_blocks[k];
     problem.append_quadratic_constraint(linear_row_count + k,
@@ -1481,7 +1483,8 @@ void flush_quadratic_constraints(mps_data_model_t<i_t, f_t>& problem,
                                         block.rhs_value,
                                         block.quad_triples.vals,
                                         block.quad_triples.rows,
-                                        block.quad_triples.cols);
+                                        block.quad_triples.cols,
+                                        qc_workspace);
   }
 }