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 "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;

  elemT* d_image_data = nullptr;
  // 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
  float* d_weights_data = nullptr;
  elemT* d_kappa_data = nullptr;

  // Buffers for GPU input/output to avoid reallocating memory on each call see usage in set_up() and ~CudaGibbsPenalty()
  mutable double* d_scalar = nullptr;
  // d_scalar is used for compute_value and compute_gradient_times_input as output variable
  mutable elemT* d_input_data = nullptr;
  // d_input_data is used for storing input image for compute_gradient_times_input and accumulate_Hessian_times_input
  mutable elemT* d_output_data = nullptr;
  // 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);
  ~CudaGibbsPenalty();

  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__