hyhound.hpp

Go to the documentation of this file.

#pragma once
 
#include <batmat/linalg/structure.hpp>
#include <batmat/linalg/uview.hpp>
#include <batmat/lut.hpp>
#include <batmat/micro-kernels/hyhound/export.h>
#include <batmat/platform/platform.hpp>
#include <batmat/simd.hpp>
 
namespace batmat::linalg::micro_kernels::hyhound {
 
struct BATMAT_LINALG_HYHOUND_EXPORT KernelConfig {
    bool sign_only = false;
};
 
template <class T, class Abi, index_t R>
struct triangular_accessor {
    using value_type = T;
    value_type *data;
 
    using simd = datapar::simd<std::remove_const_t<T>, Abi>;
    static constexpr ptrdiff_t inner_stride =
        datapar::simd_size<std::remove_const_t<T>, Abi>::value;
 
    static constexpr index_t num_elem_per_layer() { return R * (R + 1) / 2; }
    static constexpr size_t size() {
        return simd::size() * static_cast<size_t>(num_elem_per_layer());
    }
    static constexpr size_t alignment() {
        return datapar::simd_align<std::remove_const_t<T>, Abi>::value;
    }
 
    [[gnu::always_inline]] value_type &operator()(index_t r, index_t c) const noexcept {
        assert(r <= c);
        return data[(r + c * (c + 1) / 2) * inner_stride];
    }
    [[gnu::always_inline]] simd load(index_t r, index_t c) const noexcept {
        return datapar::aligned_load<simd>(&operator()(r, c));
    }
    [[gnu::always_inline]] void store(simd x, index_t r, index_t c) const noexcept
        requires(!std::is_const_v<T>)
    {
        datapar::aligned_store(x, &operator()(r, c));
    }
 
    [[gnu::always_inline]] triangular_accessor(value_type *data) noexcept : data{data} {}
    operator triangular_accessor<const T, Abi, R>() const noexcept { return {data}; }
};
 
template <class T, class Abi>
inline constexpr index_t SizeR = gemm::RowsReg<T, Abi>; // TODO
template <class T, class Abi>
inline constexpr index_t SizeS = gemm::RowsReg<T, Abi>; // TODO
 
template <class T, class Abi, KernelConfig Conf, index_t R, StorageOrder OL, StorageOrder OA>
void hyhound_diag_diag_microkernel(index_t kA, triangular_accessor<T, Abi, SizeR<T, Abi>> W,
                                   uview<T, Abi, OL> L, uview<T, Abi, OA> A,
                                   uview<const T, Abi, StorageOrder::ColMajor> diag) noexcept;
 
template <class T, class Abi, KernelConfig Conf, index_t R, StorageOrder OL, StorageOrder OA>
void hyhound_diag_full_microkernel(index_t kA, uview<T, Abi, OL> L, uview<T, Abi, OA> A,
                                   uview<const T, Abi, StorageOrder::ColMajor> diag) noexcept;
 
enum class Structure {
    General = 0,
    Zero    = 1,
    Upper   = 2,
};
 
template <class T, class Abi, KernelConfig Conf, index_t R, index_t S, StorageOrder OL,
          StorageOrder OA, StorageOrder OB>
void hyhound_diag_tail_microkernel(index_t kA_in_offset, index_t kA_in, index_t k,
                                   triangular_accessor<const T, Abi, SizeR<T, Abi>> W,
                                   uview<T, Abi, OL> L, uview<const T, Abi, OA> A_in,
                                   uview<T, Abi, OA> A_out, uview<const T, Abi, OB> B,
                                   uview<const T, Abi, StorageOrder::ColMajor> diag,
                                   Structure struc_L, int rotate_A) noexcept;
 
// Helper function to compute size of the storage for the matrix W (part of the hyperbolic
// Householder representation).
template <class T, class Abi, StorageOrder OL>
constexpr std::pair<index_t, index_t> hyhound_W_size(view<T, Abi, OL> L) {
    static constexpr index_constant<SizeR<std::remove_const_t<T>, Abi>> R;
    using W_t = triangular_accessor<std::remove_const_t<T>, Abi, R>;
    return {W_t::num_elem_per_layer(), (L.cols() + R - 1) / R};
}
 
// Low-level register-blocked routines
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL = StorageOrder::ColMajor,
          StorageOrder OA = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void hyhound_diag_register(view<T, Abi, OL> L, view<T, Abi, OA> A,
                                                        view<const T, Abi> D) noexcept;
 
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL = StorageOrder::ColMajor,
          StorageOrder OA = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void hyhound_diag_register(view<T, Abi, OL> L, view<T, Abi, OA> A,
                                                        view<const T, Abi> D,
                                                        view<T, Abi> W) noexcept;
 
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL = StorageOrder::ColMajor,
          StorageOrder OA = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void
hyhound_diag_apply_register(view<T, Abi, OL> L, view<const T, Abi, OA> Ain, view<T, Abi, OA> Aout,
                            view<const T, Abi, OA> B, view<const T, Abi> D, view<const T, Abi> W,
                            index_t kA_in_offset = 0) noexcept;
 
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL1 = StorageOrder::ColMajor,
          StorageOrder OA1 = StorageOrder::ColMajor, StorageOrder OL2 = StorageOrder::ColMajor,
          StorageOrder OA2 = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void
hyhound_diag_2_register(view<T, Abi, OL1> L11, view<T, Abi, OA1> A1, view<T, Abi, OL2> L21,
                        view<T, Abi, OA2> A2, view<const T, Abi> D) noexcept;
 
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL = StorageOrder::ColMajor,
          StorageOrder OW = StorageOrder::ColMajor, StorageOrder OY = StorageOrder::ColMajor,
          StorageOrder OU = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void
hyhound_diag_cyclic_register(view<T, Abi, OL> L11, view<T, Abi, OW> A1, view<T, Abi, OY> L21,
                             view<const T, Abi, OW> A22, view<T, Abi, OW> A2_out,
                             view<T, Abi, OU> L31, view<const T, Abi, OW> A31,
                             view<T, Abi, OW> A3_out, view<const T, Abi> D) noexcept;
 
template <class T, class Abi, KernelConfig Conf = {}, StorageOrder OL = StorageOrder::ColMajor,
          StorageOrder OA = StorageOrder::ColMajor, StorageOrder OLu = StorageOrder::ColMajor,
          StorageOrder OAu = StorageOrder::ColMajor>
BATMAT_LINALG_HYHOUND_EXPORT void
hyhound_diag_riccati_register(view<T, Abi, OL> L11, view<T, Abi, OA> A1, view<T, Abi, OL> L21,
                              view<const T, Abi, OA> A2, view<T, Abi, OA> A2_out,
                              view<T, Abi, OLu> Lu1, view<T, Abi, OAu> Au_out, view<const T, Abi> D,
                              bool shift_A_out) noexcept;
 
} // namespace batmat::linalg::micro_kernels::hyhound