lanes.h

Go to the documentation of this file.

 
#ifndef CHEMISTRY_LANES_H
#define CHEMISTRY_LANES_H
 
#include <cstdint>
 
#include <Kokkos_Core.hpp>
 
namespace chemistry {
 
template <bool ColumnMajor>
KOKKOS_INLINE_FUNCTION constexpr uint32_t dense_index(const uint32_t n, const uint32_t r, const uint32_t c) {
    return ColumnMajor ? c * n + r : r * n + c;
}
 
struct SerialLanes {
    static constexpr uint32_t lanes = 1;
    static constexpr bool column_major = false;  // dense matrices: rows contiguous for one thread's inner loops
    static constexpr bool parallel = false;      // loops that balance work across lanes are not worth it
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION void for_each(const uint32_t n, const F & f) const {
        for (uint32_t i = 0; i < n; i++) f(i);
    }
 
    KOKKOS_INLINE_FUNCTION void sync() const {}
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION void single(const F & f) const {
        f();
    }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION double sum(const uint32_t n, const F & f) const {
        double s = 0.0;
        for (uint32_t i = 0; i < n; i++) s += f(i);
        return s;
    }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION double argmax_abs(const uint32_t first, const uint32_t n, const F & f,
                                             uint32_t & index) const {
        index = first;
        double largest = Kokkos::fabs(f(first));
        for (uint32_t i = first + 1; i < n; i++) {
            const double v = Kokkos::fabs(f(i));
            if (v > largest) {
                largest = v;
                index = i;
            }
        }
        return largest;
    }
};
 
template <typename Member>
struct TeamLanes {
    static constexpr bool column_major = true;  // dense matrices: a column's rows contiguous across lanes
    static constexpr bool parallel = true;      // balance work across lanes (e.g. Jacobian entries, not rows)
    static constexpr uint32_t SUM_SHARES = 32;
    const Member & member;
    uint32_t lanes;
    double * partial;    // (lanes)
    uint32_t * where;    // (lanes)
 
    KOKKOS_INLINE_FUNCTION static constexpr uint32_t shares(const uint32_t L) { return L > SUM_SHARES ? L : SUM_SHARES; }
    static constexpr size_t scratch_bytes(const uint32_t L) {
        return shares(L) * sizeof(double) + L * sizeof(uint32_t) + 16;
    }
 
    KOKKOS_INLINE_FUNCTION TeamLanes(const Member & m, const uint32_t L) : member(m), lanes(L) {
        partial = static_cast<double *>(m.team_scratch(0).get_shmem(shares(L) * sizeof(double)));
        where = static_cast<uint32_t *>(m.team_scratch(0).get_shmem(L * sizeof(uint32_t)));
    }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION void for_each(const uint32_t n, const F & f) const {
        Kokkos::parallel_for(Kokkos::TeamVectorRange(member, n), f);
    }
 
    KOKKOS_INLINE_FUNCTION void sync() const { member.team_barrier(); }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION void single(const F & f) const {
        Kokkos::single(Kokkos::PerTeam(member), f);
    }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION double sum(const uint32_t n, const F & f) const {
        constexpr uint32_t S = SUM_SHARES;
        Kokkos::parallel_for(Kokkos::TeamVectorRange(member, S), [&](const uint32_t l) {
            double s = 0.0;
            for (uint32_t i = l; i < n; i += S) s += f(i);
            partial[l] = s;
        });
        member.team_barrier();
        double s = 0.0;
        for (uint32_t l = 0; l < S; l++) s += partial[l];
        member.team_barrier();
        return s;
    }
 
    template <typename F>
    KOKKOS_INLINE_FUNCTION double argmax_abs(const uint32_t first, const uint32_t n, const F & f,
                                             uint32_t & index) const {
        const uint32_t L = lanes;
        Kokkos::parallel_for(Kokkos::TeamVectorRange(member, L), [&](const uint32_t l) {
            double largest = -1.0;
            uint32_t at = n;
            for (uint32_t i = first + l; i < n; i += L) {
                const double v = Kokkos::fabs(f(i));
                if (v > largest) {
                    largest = v;
                    at = i;
                }
            }
            partial[l] = largest;
            where[l] = at;
        });
        member.team_barrier();
        double largest = -1.0;
        index = n;
        for (uint32_t l = 0; l < L; l++) {
            if (partial[l] > largest || (partial[l] == largest && where[l] < index)) {
                largest = partial[l];
                index = where[l];
            }
        }
        member.team_barrier();
        return largest;
    }
};
 
} // namespace chemistry
 
#endif // CHEMISTRY_LANES_H