threaded jacobi

Change-Id: I96a5d2f719545e14c03d73e71c8c0564e8c1c729

Author: sandwichmaker

internal/ceres/block_jacobi_preconditioner.cc

@@ -36,14 +36,15 @@
 #include "ceres/block_structure.h"
 #include "ceres/casts.h"
 #include "ceres/internal/eigen.h"
+#include "ceres/parallel_for.h"
 #include "ceres/small_blas.h"
 
 namespace ceres::internal {
 
 BlockSparseJacobiPreconditioner::BlockSparseJacobiPreconditioner(
-    const BlockSparseMatrix& A) {
-  m_ = std::make_unique<BlockRandomAccessDiagonalMatrix>(
-      A.block_structure()->cols);
+    Preconditioner::Options options, const BlockSparseMatrix& A)
+    : options_(std::move(options)) {
+  m_ = std::make_unique<BlockRandomAccessDiagonalMatrix>(A.block_structure());
 }
 
 BlockSparseJacobiPreconditioner::~BlockSparseJacobiPreconditioner() = default;
@@ -94,7 +95,8 @@ bool BlockSparseJacobiPreconditioner::UpdateImpl(const BlockSparseMatrix& A,
 }
 
 BlockCRSJacobiPreconditioner::BlockCRSJacobiPreconditioner(
-    const CompressedRowSparseMatrix& A) {
+    Preconditioner::Options options, const CompressedRowSparseMatrix& A)
+    : options_(std::move(options)), locks_(A.col_blocks().size()) {
   auto& col_blocks = A.col_blocks();
 
   // Compute the number of non-zeros in the preconditioner. This is needed so

internal/ceres/block_jacobi_preconditioner.h

@@ -52,7 +52,8 @@ class CERES_NO_EXPORT BlockSparseJacobiPreconditioner
     : public BlockSparseMatrixPreconditioner {
  public:
   // A must remain valid while the BlockJacobiPreconditioner is.
-  explicit BlockSparseJacobiPreconditioner(const BlockSparseMatrix& A);
+  explicit BlockSparseJacobiPreconditioner(Preconditioner::Options,
+                                           const BlockSparseMatrix& A);
   ~BlockSparseJacobiPreconditioner() override;
   void RightMultiplyAndAccumulate(const double* x, double* y) const final {
     return m_->RightMultiplyAndAccumulate(x, y);
@@ -64,6 +65,7 @@ class CERES_NO_EXPORT BlockSparseJacobiPreconditioner
  private:
   bool UpdateImpl(const BlockSparseMatrix& A, const double* D) final;
 
+  Preconditioner::Options options_;
   std::unique_ptr<BlockRandomAccessDiagonalMatrix> m_;
 };
 
@@ -73,7 +75,8 @@ class CERES_NO_EXPORT BlockCRSJacobiPreconditioner
     : public CompressedRowSparseMatrixPreconditioner {
  public:
   // A must remain valid while the BlockJacobiPreconditioner is.
-  explicit BlockCRSJacobiPreconditioner(const CompressedRowSparseMatrix& A);
+  explicit BlockCRSJacobiPreconditioner(Preconditioner::Options options,
+                                        const CompressedRowSparseMatrix& A);
   ~BlockCRSJacobiPreconditioner() override;
   void RightMultiplyAndAccumulate(const double* x, double* y) const final {
     m_->RightMultiplyAndAccumulate(x, y);
@@ -85,6 +88,8 @@ class CERES_NO_EXPORT BlockCRSJacobiPreconditioner
  private:
   bool UpdateImpl(const CompressedRowSparseMatrix& A, const double* D) final;
 
+  Preconditioner::Options options_;
+  std::vector<std::mutex> locks_;
   std::unique_ptr<CompressedRowSparseMatrix> m_;
 };
 

internal/ceres/cgnr_solver.cc

@@ -135,18 +135,20 @@ LinearSolver::Summary CgnrSolver::SolveImpl(
     double* x) {
   EventLogger event_logger("CgnrSolver::Solve");
   if (!preconditioner_) {
+    Preconditioner::Options preconditioner_options;
+    preconditioner_options.type = options_.preconditioner_type;
+    preconditioner_options.subset_preconditioner_start_row_block =
+        options_.subset_preconditioner_start_row_block;
+    preconditioner_options.sparse_linear_algebra_library_type =
+        options_.sparse_linear_algebra_library_type;
+    preconditioner_options.ordering_type = options_.ordering_type;
+    preconditioner_options.num_threads = options_.num_threads;
+    preconditioner_options.context = options_.context;
+
     if (options_.preconditioner_type == JACOBI) {
-      preconditioner_ = std::make_unique<BlockSparseJacobiPreconditioner>(*A);
+      preconditioner_ = std::make_unique<BlockSparseJacobiPreconditioner>(
+          preconditioner_options, *A);
     } else if (options_.preconditioner_type == SUBSET) {
-      Preconditioner::Options preconditioner_options;
-      preconditioner_options.type = SUBSET;
-      preconditioner_options.subset_preconditioner_start_row_block =
-          options_.subset_preconditioner_start_row_block;
-      preconditioner_options.sparse_linear_algebra_library_type =
-          options_.sparse_linear_algebra_library_type;
-      preconditioner_options.ordering_type = options_.ordering_type;
-      preconditioner_options.num_threads = options_.num_threads;
-      preconditioner_options.context = options_.context;
       preconditioner_ =
           std::make_unique<SubsetPreconditioner>(preconditioner_options, *A);
     } else {
@@ -238,9 +240,11 @@ class CERES_NO_EXPORT CudaIdentityPreconditioner final
 class CERES_NO_EXPORT CudaJacobiPreconditioner final
     : public CudaPreconditioner {
  public:
-  explicit CudaJacobiPreconditioner(ContextImpl* context,
+  explicit CudaJacobiPreconditioner(Preconditioner::Options options,
                                     const CompressedRowSparseMatrix& A)
-      : cpu_preconditioner_(A), m_(context, cpu_preconditioner_.matrix()) {}
+      : options_(std::move(options)),
+        cpu_preconditioner_(options_, A),
+        m_(options_->context, cpu_preconditioner_.matrix()) {}
   ~CudaJacobiPreconditioner() = default;
 
   void Update(const CompressedRowSparseMatrix& A, const double* D) final {
@@ -253,6 +257,7 @@ class CERES_NO_EXPORT CudaJacobiPreconditioner final
   }
 
  private:
+  Preconditioner::Options options_;
   BlockCRSJacobiPreconditioner cpu_preconditioner_;
   CudaSparseMatrix m_;
 };
@@ -297,9 +302,20 @@ void CudaCgnrSolver::CpuToGpuTransfer(const CompressedRowSparseMatrix& A,
     Atb_ = std::make_unique<CudaVector>(options_.context, A.num_cols());
     Ax_ = std::make_unique<CudaVector>(options_.context, A.num_rows());
     D_ = std::make_unique<CudaVector>(options_.context, A.num_cols());
+
+    Preconditioner::Options preconditioner_options;
+    preconditioner_options.type = options_.preconditioner_type;
+    preconditioner_options.subset_preconditioner_start_row_block =
+        options_.subset_preconditioner_start_row_block;
+    preconditioner_options.sparse_linear_algebra_library_type =
+        options_.sparse_linear_algebra_library_type;
+    preconditioner_options.ordering_type = options_.ordering_type;
+    preconditioner_options.num_threads = options_.num_threads;
+    preconditioner_options.context = options_.context;
+
     if (options_.preconditioner_type == JACOBI) {
       preconditioner_ =
-          std::make_unique<CudaJacobiPreconditioner>(options_.context, A);
+          std::make_unique<CudaJacobiPreconditioner>(preconditioner_options, A);
     } else {
       preconditioner_ = std::make_unique<CudaIdentityPreconditioner>();
     }