precision_dispatch.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_PRECISION_DISPATCH_HPP_
#define GKO_PUBLIC_CORE_BASE_PRECISION_DISPATCH_HPP_
 
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/temporary_conversion.hpp>
#include <ginkgo/core/distributed/vector.hpp>
#include <ginkgo/core/matrix/dense.hpp>
 
 
namespace gko {
 
 
template <typename ValueType, typename Ptr>
detail::temporary_conversion<std::conditional_t<
    std::is_const<detail::pointee<Ptr>>::value, const matrix::Dense<ValueType>,
    matrix::Dense<ValueType>>>
make_temporary_conversion(Ptr&& matrix)
{
    using Pointee = detail::pointee<Ptr>;
    using Dense = matrix::Dense<ValueType>;
    using NextDense = matrix::Dense<next_precision<ValueType>>;
    using NextNextDense =
        matrix::Dense<next_precision<next_precision<ValueType>>>;
    using MaybeConstDense =
        std::conditional_t<std::is_const<Pointee>::value, const Dense, Dense>;
    auto result = detail::temporary_conversion<
        MaybeConstDense>::template create<NextDense, NextNextDense>(matrix);
    if (!result) {
        GKO_NOT_SUPPORTED(matrix);
    }
    return result;
}
 
 
template <typename ValueType, typename Function, typename... Args>
void precision_dispatch(Function fn, Args*... linops)
{
    fn(make_temporary_conversion<ValueType>(linops).get()...);
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* in, LinOp* out)
{
    // do we need to convert complex Dense to real Dense?
    // all real dense vectors are intra-convertible, thus by casting to
    // ConvertibleTo<matrix::Dense<>>, we can check whether a LinOp is a real
    // dense matrix:
    auto complex_to_real =
        !(is_complex<ValueType>() ||
          dynamic_cast<const ConvertibleTo<matrix::Dense<>>*>(in));
    if (complex_to_real) {
        auto dense_in = make_temporary_conversion<to_complex<ValueType>>(in);
        auto dense_out = make_temporary_conversion<to_complex<ValueType>>(out);
        using Dense = matrix::Dense<ValueType>;
        // These dynamic_casts are only needed to make the code compile
        // If ValueType is complex, this branch will never be taken
        // If ValueType is real, the cast is a no-op
        fn(dynamic_cast<const Dense*>(dense_in->create_real_view().get()),
           dynamic_cast<Dense*>(dense_out->create_real_view().get()));
    } else {
        precision_dispatch<ValueType>(fn, in, out);
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* alpha,
                                     const LinOp* in, LinOp* out)
{
    // do we need to convert complex Dense to real Dense?
    // all real dense vectors are intra-convertible, thus by casting to
    // ConvertibleTo<matrix::Dense<>>, we can check whether a LinOp is a real
    // dense matrix:
    auto complex_to_real =
        !(is_complex<ValueType>() ||
          dynamic_cast<const ConvertibleTo<matrix::Dense<>>*>(in));
    if (complex_to_real) {
        auto dense_in = make_temporary_conversion<to_complex<ValueType>>(in);
        auto dense_out = make_temporary_conversion<to_complex<ValueType>>(out);
        auto dense_alpha = make_temporary_conversion<ValueType>(alpha);
        using Dense = matrix::Dense<ValueType>;
        // These dynamic_casts are only needed to make the code compile
        // If ValueType is complex, this branch will never be taken
        // If ValueType is real, the cast is a no-op
        fn(dense_alpha.get(),
           dynamic_cast<const Dense*>(dense_in->create_real_view().get()),
           dynamic_cast<Dense*>(dense_out->create_real_view().get()));
    } else {
        precision_dispatch<ValueType>(fn, alpha, in, out);
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* alpha,
                                     const LinOp* in, const LinOp* beta,
                                     LinOp* out)
{
    // do we need to convert complex Dense to real Dense?
    // all real dense vectors are intra-convertible, thus by casting to
    // ConvertibleTo<matrix::Dense<>>, we can check whether a LinOp is a real
    // dense matrix:
    auto complex_to_real =
        !(is_complex<ValueType>() ||
          dynamic_cast<const ConvertibleTo<matrix::Dense<>>*>(in));
    if (complex_to_real) {
        auto dense_in = make_temporary_conversion<to_complex<ValueType>>(in);
        auto dense_out = make_temporary_conversion<to_complex<ValueType>>(out);
        auto dense_alpha = make_temporary_conversion<ValueType>(alpha);
        auto dense_beta = make_temporary_conversion<ValueType>(beta);
        using Dense = matrix::Dense<ValueType>;
        // These dynamic_casts are only needed to make the code compile
        // If ValueType is complex, this branch will never be taken
        // If ValueType is real, the cast is a no-op
        fn(dense_alpha.get(),
           dynamic_cast<const Dense*>(dense_in->create_real_view().get()),
           dense_beta.get(),
           dynamic_cast<Dense*>(dense_out->create_real_view().get()));
    } else {
        precision_dispatch<ValueType>(fn, alpha, in, beta, out);
    }
}
 
 
template <typename ValueType, typename Function>
void mixed_precision_dispatch(Function fn, const LinOp* in, LinOp* out)
{
#ifdef GINKGO_MIXED_PRECISION
    using fst_type = matrix::Dense<ValueType>;
    using snd_type = matrix::Dense<next_precision<ValueType>>;
    using trd_type = matrix::Dense<next_precision<next_precision<ValueType>>>;
    auto dispatch_out_vector = [&](auto dense_in) {
        if (auto dense_out = dynamic_cast<fst_type*>(out)) {
            fn(dense_in, dense_out);
        } else if (auto dense_out = dynamic_cast<snd_type*>(out)) {
            fn(dense_in, dense_out);
        } else if (auto dense_out = dynamic_cast<trd_type*>(out)) {
            fn(dense_in, dense_out);
        } else {
            GKO_NOT_SUPPORTED(out);
        }
    };
    if (auto dense_in = dynamic_cast<const fst_type*>(in)) {
        dispatch_out_vector(dense_in);
    } else if (auto dense_in = dynamic_cast<const snd_type*>(in)) {
        dispatch_out_vector(dense_in);
    } else if (auto dense_in = dynamic_cast<const trd_type*>(in)) {
        dispatch_out_vector(dense_in);
    } else {
        GKO_NOT_SUPPORTED(in);
    }
#else
    precision_dispatch<ValueType>(fn, in, out);
#endif
}
 
 
template <typename ValueType, typename Function,
          std::enable_if_t<is_complex<ValueType>()>* = nullptr>
void mixed_precision_dispatch_real_complex(Function fn, const LinOp* in,
                                           LinOp* out)
{
#ifdef GINKGO_MIXED_PRECISION
    mixed_precision_dispatch<ValueType>(fn, in, out);
#else
    precision_dispatch<ValueType>(fn, in, out);
#endif
}
 
 
template <typename ValueType, typename Function,
          std::enable_if_t<!is_complex<ValueType>()>* = nullptr>
void mixed_precision_dispatch_real_complex(Function fn, const LinOp* in,
                                           LinOp* out)
{
#ifdef GINKGO_MIXED_PRECISION
    if (!dynamic_cast<const ConvertibleTo<matrix::Dense<>>*>(in)) {
        mixed_precision_dispatch<to_complex<ValueType>>(
            [&fn](auto dense_in, auto dense_out) {
                fn(dense_in->create_real_view().get(),
                   dense_out->create_real_view().get());
            },
            in, out);
    } else {
        mixed_precision_dispatch<ValueType>(fn, in, out);
    }
#else
    precision_dispatch_real_complex<ValueType>(fn, in, out);
#endif
}
 
 
namespace experimental {
 
 
#if GINKGO_BUILD_MPI
 
 
namespace distributed {
 
 
template <typename ValueType>
gko::detail::temporary_conversion<Vector<ValueType>> make_temporary_conversion(
    LinOp* matrix)
{
    auto result =
        gko::detail::temporary_conversion<Vector<ValueType>>::template create<
            Vector<next_precision_base<ValueType>>>(matrix);
    if (!result) {
        GKO_NOT_SUPPORTED(matrix);
    }
    return result;
}
 
 
template <typename ValueType>
gko::detail::temporary_conversion<const Vector<ValueType>>
make_temporary_conversion(const LinOp* matrix)
{
    auto result = gko::detail::temporary_conversion<const Vector<ValueType>>::
        template create<Vector<next_precision_base<ValueType>>>(matrix);
    if (!result) {
        GKO_NOT_SUPPORTED(matrix);
    }
    return result;
}
 
 
template <typename ValueType, typename Function, typename... Args>
void precision_dispatch(Function fn, Args*... linops)
{
    if constexpr (std::is_same_v<remove_complex<ValueType>, half>) {
        GKO_NOT_SUPPORTED(nullptr);
    } else {
        fn(distributed::make_temporary_conversion<ValueType>(linops).get()...);
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* in, LinOp* out)
{
    if constexpr (std::is_same_v<remove_complex<ValueType>, half>) {
        GKO_NOT_SUPPORTED(nullptr);
    } else {
        auto complex_to_real = !(
            is_complex<ValueType>() ||
            dynamic_cast<
                const ConvertibleTo<experimental::distributed::Vector<>>*>(in));
        if (complex_to_real) {
            auto dense_in =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    in);
            auto dense_out =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    out);
            using Vector = experimental::distributed::Vector<ValueType>;
            // These dynamic_casts are only needed to make the code compile
            // If ValueType is complex, this branch will never be taken
            // If ValueType is real, the cast is a no-op
            fn(dynamic_cast<const Vector*>(dense_in->create_real_view().get()),
               dynamic_cast<Vector*>(dense_out->create_real_view().get()));
        } else {
            distributed::precision_dispatch<ValueType>(fn, in, out);
        }
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* alpha,
                                     const LinOp* in, LinOp* out)
{
    if constexpr (std::is_same_v<remove_complex<ValueType>, half>) {
        GKO_NOT_SUPPORTED(nullptr);
    } else {
        auto complex_to_real = !(
            is_complex<ValueType>() ||
            dynamic_cast<
                const ConvertibleTo<experimental::distributed::Vector<>>*>(in));
        if (complex_to_real) {
            auto dense_in =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    in);
            auto dense_out =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    out);
            auto dense_alpha = gko::make_temporary_conversion<ValueType>(alpha);
            using Vector = experimental::distributed::Vector<ValueType>;
            // These dynamic_casts are only needed to make the code compile
            // If ValueType is complex, this branch will never be taken
            // If ValueType is real, the cast is a no-op
            fn(dense_alpha.get(),
               dynamic_cast<const Vector*>(dense_in->create_real_view().get()),
               dynamic_cast<Vector*>(dense_out->create_real_view().get()));
        } else {
            fn(gko::make_temporary_conversion<ValueType>(alpha).get(),
               distributed::make_temporary_conversion<ValueType>(in).get(),
               distributed::make_temporary_conversion<ValueType>(out).get());
        }
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex(Function fn, const LinOp* alpha,
                                     const LinOp* in, const LinOp* beta,
                                     LinOp* out)
{
    if constexpr (std::is_same_v<remove_complex<ValueType>, half>) {
        GKO_NOT_SUPPORTED(nullptr);
    } else {
        auto complex_to_real = !(
            is_complex<ValueType>() ||
            dynamic_cast<
                const ConvertibleTo<experimental::distributed::Vector<>>*>(in));
        if (complex_to_real) {
            auto dense_in =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    in);
            auto dense_out =
                distributed::make_temporary_conversion<to_complex<ValueType>>(
                    out);
            auto dense_alpha = gko::make_temporary_conversion<ValueType>(alpha);
            auto dense_beta = gko::make_temporary_conversion<ValueType>(beta);
            using Vector = experimental::distributed::Vector<ValueType>;
            // These dynamic_casts are only needed to make the code compile
            // If ValueType is complex, this branch will never be taken
            // If ValueType is real, the cast is a no-op
            fn(dense_alpha.get(),
               dynamic_cast<const Vector*>(dense_in->create_real_view().get()),
               dense_beta.get(),
               dynamic_cast<Vector*>(dense_out->create_real_view().get()));
        } else {
            fn(gko::make_temporary_conversion<ValueType>(alpha).get(),
               distributed::make_temporary_conversion<ValueType>(in).get(),
               gko::make_temporary_conversion<ValueType>(beta).get(),
               distributed::make_temporary_conversion<ValueType>(out).get());
        }
    }
}
 
 
}  // namespace distributed
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex_distributed(Function fn, const LinOp* in,
                                                 LinOp* out)
{
    if (dynamic_cast<const experimental::distributed::DistributedBase*>(in)) {
        experimental::distributed::precision_dispatch_real_complex<ValueType>(
            fn, in, out);
    } else {
        gko::precision_dispatch_real_complex<ValueType>(fn, in, out);
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex_distributed(Function fn,
                                                 const LinOp* alpha,
                                                 const LinOp* in, LinOp* out)
{
    if (dynamic_cast<const experimental::distributed::DistributedBase*>(in)) {
        experimental::distributed::precision_dispatch_real_complex<ValueType>(
            fn, alpha, in, out);
    } else {
        gko::precision_dispatch_real_complex<ValueType>(fn, alpha, in, out);
    }
}
 
 
template <typename ValueType, typename Function>
void precision_dispatch_real_complex_distributed(Function fn,
                                                 const LinOp* alpha,
                                                 const LinOp* in,
                                                 const LinOp* beta, LinOp* out)
{
    if (dynamic_cast<const experimental::distributed::DistributedBase*>(in)) {
        experimental::distributed::precision_dispatch_real_complex<ValueType>(
            fn, alpha, in, beta, out);
 
    } else {
        gko::precision_dispatch_real_complex<ValueType>(fn, alpha, in, beta,
                                                        out);
    }
}
 
 
#else
 
 
template <typename ValueType, typename Function, typename... Args>
void precision_dispatch_real_complex_distributed(Function fn, Args*... args)
{
    precision_dispatch_real_complex<ValueType>(fn, args...);
}
 
 
#endif
 
 
}  // namespace experimental
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_PRECISION_DISPATCH_HPP_