CudaGibbsPenalty.h

Go to the documentation of this file.

/*
    Copyright (C) 2025, University College London
    Copyright (C) 2025, University of Milano-Bicocca
 
    SPDX-License-Identifier: Apache-2.0
 
    See STIR/LICENSE.txt for details
*/
 
#ifndef __stir_recon_buildblock_CUDA_CudaGibbsPenalty_H__
#define __stir_recon_buildblock_CUDA_CudaGibbsPenalty_H__
 
#include "stir/Array.h"
#include "stir/DiscretisedDensity.h"
#include "stir/cuda_utilities.h"
#include "stir/recon_buildblock/GibbsPenalty.h"
 
#include "cuvec.cuh"
#include "stir/shared_ptr.h"
#include <string>
 
START_NAMESPACE_STIR

template <typename elemT, typename PotentialT>
class CudaGibbsPenalty : public GibbsPenalty<elemT, PotentialT>
{
private:
  typedef GibbsPenalty<elemT, PotentialT> base_type;
 
protected:
  // GPU block and grid dimensions
 
  cuda_dim3 block_dim;
  cuda_dim3 grid_dim;
 
  // Variables used for shared memory operations
  int threads_per_block;
  size_t shared_mem_bytes;
 
  mutable CuVec<elemT> d_image_data;
  // Currently stir:CartesianCoordinate3D<int> is not supported on GPU, we need a simple structure to store boundaries.
  cuda_int3 d_image_dim;
  cuda_int3 d_image_max_indices;
  cuda_int3 d_image_min_indices;
  cuda_int3 d_weight_max_indices;
  cuda_int3 d_weight_min_indices;
 
  // GPU pointers to weights and kappa data
  CuVec<float> d_weights_data;
  CuVec<elemT> d_kappa_data;
 
  // Buffers for GPU input/output to avoid reallocating memory on each call see usage in set_up() and ~CudaGibbsPenalty()
  mutable CuVec<double> d_scalar;
  // d_scalar is used for compute_value and compute_gradient_times_input as output variable
  mutable CuVec<elemT> d_input_data;
  // d_input_data is used for storing input image for compute_gradient_times_input and accumulate_Hessian_times_input
  mutable CuVec<elemT> d_output_data;
  // d_output_data is used for storing output image for compute_gradient, accumulate_Hessian_times_input and
  // compute_Hessian_diagonal
 
public:
  CudaGibbsPenalty();
  CudaGibbsPenalty(const bool only_2D, float penalization_factor);

  Succeeded set_up(shared_ptr<const DiscretisedDensity<3, elemT>> const& target_sptr) override;

  void set_weights(const Array<3, float>& w) override;

  void set_kappa_sptr(const shared_ptr<const DiscretisedDensity<3, elemT>>& k) override;

  double compute_value(const DiscretisedDensity<3, elemT>& current_image_estimate) override;

  void compute_gradient(DiscretisedDensity<3, elemT>& prior_gradient,
                        const DiscretisedDensity<3, elemT>& current_image_estimate) override;

  double compute_gradient_times_input(const DiscretisedDensity<3, elemT>& input,
                                      const DiscretisedDensity<3, elemT>& current_image_estimate) override;

  void compute_Hessian_diagonal(DiscretisedDensity<3, elemT>& Hessian_diagonal,
                                const DiscretisedDensity<3, elemT>& current_estimate) const override;

  void accumulate_Hessian_times_input(DiscretisedDensity<3, elemT>& output,
                                      const DiscretisedDensity<3, elemT>& current_image_estimate,
                                      const DiscretisedDensity<3, elemT>& input) const override;
};
 
END_NAMESPACE_STIR
 
#ifdef __CUDACC__
// CUDA compiler sees everything
#  include "stir/recon_buildblock/CUDA/CudaGibbsPenalty.cuh"
#endif
 
#endif // __stir_recon_buildblock_CUDA_CudaGibbsPenalty_H__