tile_operations.hpp

Go to the documentation of this file.

#ifndef TDLS_SOLVERS_TILE_OPERATIONS_HPP
#define TDLS_SOLVERS_TILE_OPERATIONS_HPP
 
 
 
 
 
 
#include <tdls/core/macros.hpp>
 
 
 
// clang reports a forced unrolling that the optimizer could not perform
// through -Wpass-failed. The unrolling requested by TDLS_UNROLL_FORCE is
// a performance hint: a failed hint does not affect correctness. The
// suppression is scoped to this header and to that warning only.
#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wpass-failed"
#endif
 
namespace tdls {
 
 
 
template<typename T, int tile_size, bool unroll_inner>
struct TileOperations {
 
    template<int k_extent>
    TDLS_HOST_DEVICE TDLS_FORCEINLINE static constexpr void swap_rows(T* TDLS_RESTRICT t, int k,
                                                                      int r) noexcept {
        if constexpr (unroll_inner) {
            TDLS_UNROLL_FORCE
            for (int row = 0; row < k_extent; ++row) {
                if (row == r) {
                    TDLS_UNROLL_FORCE
                    for (int j = 0; j < k_extent; ++j) {
                        const T tmp            = t[k * tile_size + j];
                        t[k * tile_size + j]   = t[row * tile_size + j];
                        t[row * tile_size + j] = tmp;
                    }
                }
            }
        } else {
            for (int row = 0; row < k_extent; ++row) {
                if (row == r) {
                    for (int j = 0; j < k_extent; ++j) {
                        const T tmp            = t[k * tile_size + j];
                        t[k * tile_size + j]   = t[row * tile_size + j];
                        t[row * tile_size + j] = tmp;
                    }
                }
            }
        }
    }
 
    template<int diag_extent, int col_extent>
    TDLS_HOST_DEVICE TDLS_FORCEINLINE static constexpr void
    trsm_left_unit(const T* TDLS_RESTRICT lu, T* TDLS_RESTRICT B) noexcept {
        if constexpr (unroll_inner) {
            TDLS_UNROLL_FORCE
            for (int k = 0; k < diag_extent; ++k) {
                TDLS_UNROLL_FORCE
                for (int i = k + 1; i < diag_extent; ++i) {
                    const T L_ik = lu[i * tile_size + k];
                    TDLS_UNROLL_FORCE
                    for (int j = 0; j < col_extent; ++j)
                        B[i * tile_size + j] -= L_ik * B[k * tile_size + j];
                }
            }
        } else {
            for (int k = 0; k < diag_extent; ++k) {
                for (int i = k + 1; i < diag_extent; ++i) {
                    const T L_ik = lu[i * tile_size + k];
                    for (int j = 0; j < col_extent; ++j)
                        B[i * tile_size + j] -= L_ik * B[k * tile_size + j];
                }
            }
        }
    }
 
    template<int diag_extent, int row_extent>
    TDLS_HOST_DEVICE TDLS_FORCEINLINE static constexpr void
    trsm_right(const T* TDLS_RESTRICT lu, T* TDLS_RESTRICT B) noexcept {
        if constexpr (unroll_inner) {
            TDLS_UNROLL_FORCE
            for (int k = 0; k < diag_extent; ++k) {
                const T U_kk_inv = lu[k * tile_size + k];
                TDLS_UNROLL_FORCE
                for (int i = 0; i < row_extent; ++i)
                    B[i * tile_size + k] *= U_kk_inv;
                TDLS_UNROLL_FORCE
                for (int j = k + 1; j < diag_extent; ++j) {
                    const T U_kj = lu[k * tile_size + j];
                    TDLS_UNROLL_FORCE
                    for (int i = 0; i < row_extent; ++i)
                        B[i * tile_size + j] -= B[i * tile_size + k] * U_kj;
                }
            }
        } else {
            for (int k = 0; k < diag_extent; ++k) {
                const T U_kk_inv = lu[k * tile_size + k];
                for (int i = 0; i < row_extent; ++i)
                    B[i * tile_size + k] *= U_kk_inv;
                for (int j = k + 1; j < diag_extent; ++j) {
                    const T U_kj = lu[k * tile_size + j];
                    for (int i = 0; i < row_extent; ++i)
                        B[i * tile_size + j] -= B[i * tile_size + k] * U_kj;
                }
            }
        }
    }
 
    template<int row_extent, int col_extent, int k_extent>
    TDLS_HOST_DEVICE TDLS_FORCEINLINE static constexpr void
    gemm_sub(T* TDLS_RESTRICT Ct, const T* TDLS_RESTRICT At, const T* TDLS_RESTRICT Bt) noexcept {
        if constexpr (unroll_inner) {
            TDLS_UNROLL_FORCE
            for (int i = 0; i < row_extent; ++i) {
                TDLS_UNROLL_FORCE
                for (int j = 0; j < col_extent; ++j) {
                    T sum = T(0);
                    TDLS_UNROLL_FORCE
                    for (int k = 0; k < k_extent; ++k)
                        sum += At[i * tile_size + k] * Bt[k * tile_size + j];
                    Ct[i * tile_size + j] -= sum;
                }
            }
        } else {
            for (int i = 0; i < row_extent; ++i) {
                for (int j = 0; j < col_extent; ++j) {
                    T sum = T(0);
                    for (int k = 0; k < k_extent; ++k)
                        sum += At[i * tile_size + k] * Bt[k * tile_size + j];
                    Ct[i * tile_size + j] -= sum;
                }
            }
        }
    }
};
 
 
 
} // namespace tdls
 
#if defined(__clang__)
#pragma clang diagnostic pop
#endif
 
 
 
#endif // TDLS_SOLVERS_TILE_OPERATIONS_HPP