Skip to content

Commit

Permalink
use double/complex<double> for fosi
Browse files Browse the repository at this point in the history
Co-authored-by: Tobias Ribizel <mail@ribizel.de>
  • Loading branch information
yhmtsai and upsj committed Feb 14, 2025
1 parent 385d784 commit 37d6ecc
Show file tree
Hide file tree
Showing 7 changed files with 136 additions and 63 deletions.
1 change: 0 additions & 1 deletion common/unified/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@ set(UNIFIED_SOURCES
solver/bicgstab_kernels.cpp
solver/cg_kernels.cpp
solver/cgs_kernels.cpp
solver/cgs_kernels.cpp
solver/chebyshev_kernels.cpp
solver/common_gmres_kernels.cpp
solver/fcg_kernels.cpp
Expand Down
81 changes: 65 additions & 16 deletions common/unified/solver/chebyshev_kernels.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

#include "core/solver/chebyshev_kernels.hpp"

#include <type_traits>

#include <ginkgo/core/matrix/dense.hpp>

#include "common/unified/base/kernel_launch.hpp"
Expand All @@ -15,48 +17,95 @@ namespace GKO_DEVICE_NAMESPACE {
namespace chebyshev {


template <typename ValueType, typename ScalarType>
template <typename T>
using coeff_type = gko::kernels::chebyshev::coeff_type<T>;


#if GINKGO_DPCPP_SINGLE_MODE


// we only change type in device code to keep the interface is the same as the
// other backend.
template <typename coeff_type>
using if_single_only_type =
std::conditional_t<std::is_same_v<coeff_type, double>, float,
std::complex<float>>;


#else


template <typename T>
struct type_identity {
using type = T;
};


template <typename coeff_type>
using if_single_only_type = typename type_identity<coeff_type>::type;


#endif


template <typename ValueType>
void init_update(std::shared_ptr<const DefaultExecutor> exec,
const ScalarType alpha,
const coeff_type<ValueType> alpha,
const matrix::Dense<ValueType>* inner_sol,
matrix::Dense<ValueType>* update_sol,
matrix::Dense<ValueType>* output)
{
using actual_coeff_type = if_single_only_type<coeff_type<ValueType>>;
using type = device_type<actual_coeff_type>;

auto alpha_val = static_cast<actual_coeff_type>(alpha);

run_kernel(
exec,
[] GKO_KERNEL(auto row, auto col, auto alpha, auto inner_sol,
auto update_sol, auto output) {
const auto inner_val = inner_sol(row, col);
update_sol(row, col) = inner_val;
output(row, col) += alpha * inner_val;
const auto inner_val = static_cast<type>(inner_sol(row, col));
update_sol(row, col) =
static_cast<device_type<ValueType>>(inner_val);
output(row, col) = static_cast<device_type<ValueType>>(
static_cast<type>(output(row, col)) +
static_cast<type>(alpha) * inner_val);
},
output->get_size(), alpha, inner_sol, update_sol, output);
output->get_size(), alpha_val, inner_sol, update_sol, output);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_SCALAR_TYPE(
GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);


template <typename ValueType, typename ScalarType>
template <typename ValueType>
void update(std::shared_ptr<const DefaultExecutor> exec, const ScalarType alpha,
const ScalarType beta, matrix::Dense<ValueType>* inner_sol,
matrix::Dense<ValueType>* update_sol,
matrix::Dense<ValueType>* output)
{
using actual_coeff_type = if_single_only_type<coeff_type<ValueType>>;
using type = device_type<actual_coeff_type>;

auto alpha_val = static_cast<actual_coeff_type>(alpha);
auto beta_val = static_cast<actual_coeff_type>(beta);

run_kernel(
exec,
[] GKO_KERNEL(auto row, auto col, auto alpha, auto beta, auto inner_sol,
auto update_sol, auto output) {
const auto val = inner_sol(row, col) + beta * update_sol(row, col);
inner_sol(row, col) = val;
update_sol(row, col) = val;
output(row, col) += alpha * val;
const auto val = static_cast<type>(inner_sol(row, col)) +
static_cast<type>(beta) *
static_cast<type>(update_sol(row, col));
inner_sol(row, col) = static_cast<device_type<ValueType>>(val);
update_sol(row, col) = static_cast<device_type<ValueType>>(val);
output(row, col) = static_cast<device_type<ValueType>>(
static_cast<type>(output(row, col)) +
static_cast<type>(alpha) * val);
},
output->get_size(), alpha, beta, inner_sol, update_sol, output);
output->get_size(), alpha_val, beta_val, inner_sol, update_sol, output);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_SCALAR_TYPE(
GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);


} // namespace chebyshev
Expand Down
4 changes: 2 additions & 2 deletions core/device_hooks/common_kernels.inc.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -678,8 +678,8 @@ GKO_STUB_CB_GMRES_CONST(GKO_DECLARE_CB_GMRES_SOLVE_KRYLOV_KERNEL);
namespace chebyshev {


GKO_STUB_VALUE_AND_SCALAR_TYPE(GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);
GKO_STUB_VALUE_AND_SCALAR_TYPE(GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);
GKO_STUB_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);
GKO_STUB_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);


} // namespace chebyshev
Expand Down
11 changes: 6 additions & 5 deletions core/solver/chebyshev.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ void Chebyshev<ValueType>::apply_dense_impl(const VectorType* dense_b,
{
using Vector = matrix::Dense<ValueType>;
using ws = workspace_traits<Chebyshev>;
using coeff_type = detail::coeff_type<ValueType>;

auto exec = this->get_executor();
this->setup_workspace();
Expand All @@ -183,8 +184,8 @@ void Chebyshev<ValueType>::apply_dense_impl(const VectorType* dense_b,

GKO_SOLVER_ONE_MINUS_ONE();

auto alpha_ref = ValueType{1} / center_;
auto beta_ref = ValueType{0.5} * (foci_direction_ * alpha_ref) *
auto alpha_ref = coeff_type{1} / center_;
auto beta_ref = coeff_type{0.5} * (foci_direction_ * alpha_ref) *
(foci_direction_ * alpha_ref);

auto& stop_status = this->template create_workspace_array<stopping_status>(
Expand Down Expand Up @@ -239,10 +240,10 @@ void Chebyshev<ValueType>::apply_dense_impl(const VectorType* dense_b,
}
// beta_ref for iter == 1 is initialized in the beginning
if (iter > 1) {
beta_ref = (foci_direction_ * alpha_ref / ValueType{2.0}) *
(foci_direction_ * alpha_ref / ValueType{2.0});
beta_ref = (foci_direction_ * alpha_ref / coeff_type{2.0}) *
(foci_direction_ * alpha_ref / coeff_type{2.0});
}
alpha_ref = ValueType{1.0} / (center_ - beta_ref / alpha_ref);
alpha_ref = coeff_type{1.0} / (center_ - beta_ref / alpha_ref);
// z = z + beta * p
// p = z
// x += alpha * p
Expand Down
36 changes: 21 additions & 15 deletions core/solver/chebyshev_kernels.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,25 +20,31 @@ namespace kernels {
namespace chebyshev {


#define GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL(ValueType, ScalarType) \
void init_update(std::shared_ptr<const DefaultExecutor> exec, \
const ScalarType alpha, \
const matrix::Dense<ValueType>* inner_sol, \
matrix::Dense<ValueType>* update_sol, \
template <typename T>
using coeff_type =
std::conditional_t<is_complex<T>, std::complex<double>, double>;


#define GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL(ValueType) \
void init_update(std::shared_ptr<const DefaultExecutor> exec, \
const coeff_type<ValueType> alpha, \
const matrix::Dense<ValueType>* inner_sol, \
matrix::Dense<ValueType>* update_sol, \
matrix::Dense<ValueType>* output)

#define GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL(ValueType, ScalarType) \
void update(std::shared_ptr<const DefaultExecutor> exec, \
const ScalarType alpha, const ScalarType beta, \
matrix::Dense<ValueType>* inner_sol, \
matrix::Dense<ValueType>* update_sol, \
#define GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL(ValueType) \
void update(std::shared_ptr<const DefaultExecutor> exec, \
const coeff_type<ValueType> alpha, \
const coeff_type<ValueType> beta, \
matrix::Dense<ValueType>* inner_sol, \
matrix::Dense<ValueType>* update_sol, \
matrix::Dense<ValueType>* output)

#define GKO_DECLARE_ALL_AS_TEMPLATES \
template <typename ValueType, typename ScalarType> \
GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL(ValueType, ScalarType); \
template <typename ValueType, typename ScalarType> \
GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL(ValueType, ScalarType)
#define GKO_DECLARE_ALL_AS_TEMPLATES \
template <typename ValueType> \
GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL(ValueType); \
template <typename ValueType> \
GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL(ValueType)


} // namespace chebyshev
Expand Down
22 changes: 15 additions & 7 deletions include/ginkgo/core/solver/chebyshev.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,16 @@

namespace gko {
namespace solver {
namespace detail {


template <typename T>
using coeff_type =
std::conditional_t<is_complex<T>, std::complex<double>, double>;


}

/**
* Chebyshev iteration is an iterative method for solving nonsymmetric problems
* based on some knowledge of the spectrum of the (preconditioned) system
Expand Down Expand Up @@ -121,8 +129,11 @@ class Chebyshev final
* preconditioned system. It is usually be {lower bound of eigval, upper
* bound of eigval} of preconditioned real matrices.
*/
std::pair<value_type, value_type> GKO_FACTORY_PARAMETER_VECTOR(
foci, value_type{0}, value_type{1});
std::pair<detail::coeff_type<value_type>,
detail::coeff_type<value_type>>
GKO_FACTORY_PARAMETER_VECTOR(foci,
detail::coeff_type<value_type>{0},
detail::coeff_type<value_type>{1});

/**
* Default initial guess mode. The available options are under
Expand Down Expand Up @@ -170,9 +181,6 @@ class Chebyshev final
const LinOp* beta, LinOp* x,
initial_guess_mode guess) const override;

void set_relaxation_factor(
std::shared_ptr<const matrix::Dense<ValueType>> new_factor);

explicit Chebyshev(std::shared_ptr<const Executor> exec)
: EnableLinOp<Chebyshev>(std::move(exec))
{}
Expand All @@ -182,8 +190,8 @@ class Chebyshev final

private:
std::shared_ptr<const LinOp> solver_{};
ValueType center_;
ValueType foci_direction_;
detail::coeff_type<value_type> center_;
detail::coeff_type<value_type> foci_direction_;
};


Expand Down
44 changes: 27 additions & 17 deletions reference/solver/chebyshev_kernels.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,45 +12,55 @@ namespace reference {
namespace chebyshev {


template <typename ValueType, typename ScalarType>
template <typename T>
using coeff_type = gko::kernels::chebyshev::coeff_type<T>;


template <typename ValueType>
void init_update(std::shared_ptr<const DefaultExecutor> exec,
const ScalarType alpha,
const coeff_type<ValueType> alpha,
const matrix::Dense<ValueType>* inner_sol,
matrix::Dense<ValueType>* update_sol,
matrix::Dense<ValueType>* output)
{
using type = coeff_type<ValueType>;
for (size_t row = 0; row < output->get_size()[0]; row++) {
for (size_t col = 0; col < output->get_size()[1]; col++) {
const auto inner_val = inner_sol->at(row, col);
update_sol->at(row, col) = inner_val;
output->at(row, col) += alpha * inner_val;
const auto inner_val = static_cast<type>(inner_sol->at(row, col));
update_sol->at(row, col) = static_cast<ValueType>(inner_val);
output->at(row, col) =
static_cast<ValueType>(static_cast<type>(output->at(row, col)) +
static_cast<type>(alpha) * inner_val);
}
}
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_SCALAR_TYPE(
GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_INIT_UPDATE_KERNEL);


template <typename ValueType, typename ScalarType>
void update(std::shared_ptr<const DefaultExecutor> exec, const ScalarType alpha,
const ScalarType beta, matrix::Dense<ValueType>* inner_sol,
template <typename ValueType>
void update(std::shared_ptr<const DefaultExecutor> exec,
const coeff_type<ValueType> alpha, const coeff_type<ValueType> beta,
matrix::Dense<ValueType>* inner_sol,
matrix::Dense<ValueType>* update_sol,
matrix::Dense<ValueType>* output)
{
using type = coeff_type<ValueType>;
for (size_t row = 0; row < output->get_size()[0]; row++) {
for (size_t col = 0; col < output->get_size()[1]; col++) {
const auto val =
inner_sol->at(row, col) + beta * update_sol->at(row, col);
inner_sol->at(row, col) = val;
update_sol->at(row, col) = val;
output->at(row, col) += alpha * val;
const auto val = static_cast<type>(inner_sol->at(row, col)) +
static_cast<type>(beta) *
static_cast<type>(update_sol->at(row, col));
inner_sol->at(row, col) = static_cast<ValueType>(val);
update_sol->at(row, col) = static_cast<ValueType>(val);
output->at(row, col) =
static_cast<ValueType>(static_cast<type>(output->at(row, col)) +
static_cast<type>(alpha) * val);
}
}
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_SCALAR_TYPE(
GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_CHEBYSHEV_UPDATE_KERNEL);


} // namespace chebyshev
Expand Down

0 comments on commit 37d6ecc

Please sign in to comment.