IterativeDiagonalSolverCuda.h

Go to the documentation of this file.

// Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy)
// SPDX-License-Identifier: GPL-3.0-or-later
 
#ifndef DDECAL_GAIN_SOLVERS_ITERATIVE_DIAGONAL_SOLVER_CUDA_H_
#define DDECAL_GAIN_SOLVERS_ITERATIVE_DIAGONAL_SOLVER_CUDA_H_
 
#include <vector>
 
#include <cudawrappers/cu.hpp>
 
#include "IterativeDiagonalSolver.h"
#include "SolverBase.h"
#include "SolveData.h"
#include "common/Timer.h"
 
namespace dp3 {
namespace ddecal {
 
template <typename VisMatrix>
class IterativeDiagonalSolverCuda final : public SolverBase {
 public:
  IterativeDiagonalSolverCuda(bool keep_buffers = false);
  SolveResult Solve(const SolveData<VisMatrix>& data,
                    std::vector<std::vector<DComplex>>& solutions, double time,
                    std::ostream* stat_stream) override;
 
  size_t NSolutionPolarizations() const override { return 2; }
 
  bool SupportsDdSolutionIntervals() const override { return true; }
 
 private:
  void AllocateGPUBuffers(const SolveData<VisMatrix>& data);
  void DeallocateHostBuffers();
  void AllocateHostBuffers(const SolveData<VisMatrix>& data);
 
  void CopyHostToHost(size_t ch_block, bool first_iteration,
                      const SolveData<VisMatrix>& data,
                      const std::vector<DComplex>& solutions,
                      cu::Stream& stream);
 
  void CopyHostToDevice(size_t ch_block, size_t buffer_id, cu::Stream& stream,
                        cu::Event& event, const SolveData<VisMatrix>& data);
 
  void PostProcessing(size_t& iteration, double time,
                      bool has_previously_converged, bool& has_converged,
                      bool& constraints_satisfied, bool& done,
                      SolverBase::SolveResult& result,
                      std::vector<std::vector<DComplex>>& solutions,
                      SolutionSpan& next_solutions,
                      std::vector<double>& step_magnitudes,
                      std::ostream* stat_stream);
 
  bool gpu_buffers_initialized_ = false;
  bool host_buffers_initialized_ = false;
 
  bool keep_buffers_ = false;
 
  std::unique_ptr<cu::Device> device_;
  std::unique_ptr<cu::Context> context_;
  std::unique_ptr<cu::Stream> execute_stream_;
  std::unique_ptr<cu::Stream> host_to_device_stream_;
  std::unique_ptr<cu::Stream> device_to_host_stream_;
 
  struct GPUBuffers {
    // <2>[n_antennas][2], uint32_t
    std::vector<cu::DeviceMemory> antenna_pairs;
    // <2>[n_directions][n_visibilities], uint32_t
    std::vector<cu::DeviceMemory> solution_map;
    // <2>[n_visibilities], DComplex
    std::vector<cu::DeviceMemory> solutions;
    // <2>[n_visibilities], DComplex
    std::vector<cu::DeviceMemory> next_solutions;
    // <2>[n_directions][n_visibilities], MC2x2F
    std::vector<cu::DeviceMemory> model;
    // <3>[n_visibilities], MC2x2F
    std::vector<cu::DeviceMemory> residual;
    // [n_antennas][n_directions], MC2x2FDiag
    std::unique_ptr<cu::DeviceMemory> numerator;
    // [n_antennas][n_directions_solutions], float
    std::unique_ptr<cu::DeviceMemory> denominator;
  } gpu_buffers_;
 
  struct HostBuffers {
    // <n_channelblocks>[n_directions][n_visibilities], MC2x2F
    std::vector<cu::HostMemory> model;
    // <n_channelblocks>[n_visibilities], MC2x2F
    std::vector<cu::HostMemory> residual;
    // <n_channelblocks>[n_visibilities], DComplex
    std::vector<cu::HostMemory> solutions;
    // [n_channelblocks][n_antennas][n_polarizations], DComplex
    std::unique_ptr<cu::HostMemory> next_solutions;
    // <n_channelblocks>[n_visibilities], std::pair<uin32_t, uint32_t>
    std::vector<cu::HostMemory> antenna_pairs;
    // <n_channelblocks>[n_directions][n_visibilities], uint32_t
    std::vector<cu::HostMemory> solution_map;
  } host_buffers_;
};
 
}  // namespace ddecal
}  // namespace dp3
 
#endif  // DDECAL_GAIN_SOLVERS_ITERATIVE_DIAGONAL_SOLVER_CUDA_H_