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>
@@ -224,6 +225,9 @@ class cpu_optimization_problem_t : public optimization_problem_interface_t<i_t,
   std::string problem_name_;
   std::vector<std::string> var_names_{};
   std::vector<std::string> row_names_{};
+
+  /** Reused while adding quadratic constraints. */
+  mutable io::coo_canonicalization_workspace_t<i_t, f_t> qc_coo_workspace_{};
 };
 
 }  // namespace cuopt::mathematical_optimization

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

@@ -15,6 +15,72 @@
 
 namespace cuopt::mathematical_optimization::io {
 
+/**
+ * @brief Reusable workspace for COO canonicalization.
+ *
+ * Lookup vectors grow to fit the largest variable index seen so far. Each call resets only
+ * entries that were touched. A workspace must not be shared across concurrent calls.
+ */
+template <typename i_t, typename f_t>
+class coo_canonicalization_workspace_t {
+ public:
+  static_assert(std::is_integral_v<i_t> && std::is_signed_v<i_t>,
+                "coo_canonicalization_workspace_t: i_t must be signed integral");
+
+  [[nodiscard]] i_t max_index_capacity() const noexcept { return row_perm_.size(); }
+
+ private:
+  static constexpr i_t invalid_index = -1;
+
+  struct active_row_t {
+    i_t row;
+    i_t count;
+  };
+
+  struct entry_t {
+    i_t row;
+    i_t col;
+    f_t val;
+  };
+
+  std::vector<i_t> row_perm_;
+  std::vector<i_t> col_position_;
+  std::vector<active_row_t> active_rows_;
+  std::vector<i_t> offsets_;
+  std::vector<i_t> cursor_;
+  std::vector<i_t> permutation_;
+  std::vector<entry_t> entries_;
+  std::vector<i_t> output_rows_;
+  std::vector<i_t> output_cols_;
+  std::vector<f_t> output_vals_;
+
+  void ensure_index_capacity(i_t max_index)
+  {
+    const size_t needed = static_cast<size_t>(max_index) + 1;
+    if (row_perm_.size() < needed) {
+      row_perm_.resize(needed, invalid_index);
+      col_position_.resize(needed, invalid_index);
+    }
+  }
+
+  template <typename index_t, typename value_t>
+  friend void canonicalize_coo_matrix(
+    std::vector<index_t>& rows,
+    std::vector<index_t>& cols,
+    std::vector<value_t>& vals,
+    coo_canonicalization_workspace_t<index_t, value_t>& workspace);
+};
+
+/**
+ * @brief Canonicalize a symmetric matrix in COO form to upper-triangular storage
+ *        using reusable workspace storage. Defined in mps_parser.cpp.
+ */
+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 +344,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 +464,15 @@ 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)
+{
+  using i_t = typename decltype(qc_t::rows)::value_type;
+  using f_t = typename decltype(qc_t::vals)::value_type;
+  coo_canonicalization_workspace_t<i_t, f_t> workspace;
+  for (auto& qc : constraints) {
+    canonicalize_coo_matrix(qc.rows, qc.cols, qc.vals, workspace);
+  }
+}
 
 }  // 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>
 
@@ -423,6 +424,9 @@ class optimization_problem_t : public optimization_problem_interface_t<i_t, f_t>
   std::string problem_name_;
   std::vector<std::string> var_names_{};
   std::vector<std::string> row_names_{};
+
+  /** Reused while adding quadratic constraints. */
+  mutable io::coo_canonicalization_workspace_t<i_t, f_t> qc_coo_workspace_{};
 };
 
 }  // namespace cuopt::mathematical_optimization

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

@@ -1471,6 +1471,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 +1482,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);
   }
 }
 

cpp/src/io/mps_data_model.cpp

@@ -142,15 +142,17 @@ void mps_data_model_t<i_t, f_t>::set_quadratic_objective_matrix(std::span<const
 }
 
 template <typename i_t, typename f_t>
-void mps_data_model_t<i_t, f_t>::append_quadratic_constraint(i_t constraint_row_index,
-                                                             const std::string& constraint_row_name,
-                                                             char constraint_row_type,
-                                                             std::span<const f_t> linear_values,
-                                                             std::span<const i_t> linear_indices,
-                                                             f_t rhs_value,
-                                                             std::span<const f_t> vals,
-                                                             std::span<const i_t> rows,
-                                                             std::span<const i_t> cols)
+void mps_data_model_t<i_t, f_t>::append_quadratic_constraint(
+  i_t constraint_row_index,
+  const std::string& constraint_row_name,
+  char constraint_row_type,
+  std::span<const f_t> linear_values,
+  std::span<const i_t> linear_indices,
+  f_t rhs_value,
+  std::span<const f_t> vals,
+  std::span<const i_t> rows,
+  std::span<const i_t> cols,
+  coo_canonicalization_workspace_t<i_t, f_t>& workspace)
 {
   mps_parser_expects(constraint_row_index >= 0,
                      error_type_t::ValidationError,
@@ -199,7 +201,7 @@ void mps_data_model_t<i_t, f_t>::append_quadratic_constraint(i_t constraint_row_
     qc.rows.assign(rows.begin(), rows.end());
     qc.cols.assign(cols.begin(), cols.end());
     qc.vals.assign(vals.begin(), vals.end());
-    canonicalize_coo_matrix(qc.rows, qc.cols, qc.vals);
+    canonicalize_coo_matrix(qc.rows, qc.cols, qc.vals, workspace);
   }
 
   quadratic_constraints_.push_back(std::move(qc));
@@ -459,24 +461,11 @@ bool mps_data_model_t<i_t, f_t>::has_quadratic_constraints() const noexcept
   return !quadratic_constraints_.empty();
 }
 
-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)
-{
-  for (auto& qc : constraints) {
-    canonicalize_coo_matrix(qc.rows, qc.cols, qc.vals);
-  }
-}
-
 // NOTE: Explicitly instantiate all types here in order to avoid linker error
 template class mps_data_model_t<int, float>;
 
 template class mps_data_model_t<int, double>;
 
-template void canonicalize_quadratic_constraints<int, float>(
-  std::vector<mps_data_model_t<int, float>::quadratic_constraint_t>&);
-template void canonicalize_quadratic_constraints<int, double>(
-  std::vector<mps_data_model_t<int, double>::quadratic_constraint_t>&);
 //  TODO current raft to cusparse wrappers only support int64_t
 //  can be CUSPARSE_INDEX_16U, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_64I