kernel_dot.h

Go to the documentation of this file.

 
#ifndef KERNEL_DOT_H
#define KERNEL_DOT_H
 
#include "config.h"
#include "fused/microkernels/microkernel_base.h"
#include "fused/kernel_ops/kernel_helpers.h"
 
namespace detail
{
 
    template <typename T, my_size_t Bits, typename Arch>
    struct KernelDot
    {
        using K = Microkernel<T, Bits, Arch>;
        using Helpers = KernelHelpers<T, Bits, Arch>;
        static constexpr my_size_t simdWidth = K::simdWidth;
 
        // ========================================================================
        // Public API
        // ========================================================================
 
        template <typename Expr1, typename Expr2>
        FORCE_INLINE static T dot(
            const Expr1 &expr1, my_size_t base1, my_size_t stride1,
            const Expr2 &expr2, my_size_t base2, my_size_t stride2,
            my_size_t len) noexcept
        {
            if (stride1 == 1 && stride2 == 1)
            {
                // std::cout << "dot: dispatching to contiguous impl" << std::endl;
                return dot_contiguous_impl(expr1, expr2, base1, base2, len);
            }
            else
            {
                // std::cout << "dot: dispatching to strided impl" << std::endl;
                return dot_strided_impl(expr1, expr2, base1, base2, stride1, stride2, len);
            }
        }
 
        template <typename Expr1, typename Expr2>
        FORCE_INLINE static T naive_dot_physical(
            const Expr1 &expr1, my_size_t base1, my_size_t stride1,
            const Expr2 &expr2, my_size_t base2, my_size_t stride2,
            my_size_t len) noexcept
        {
            T sum = T{0};
            for (my_size_t i = 0; i < len; ++i)
                sum += expr1.data()[base1 + i * stride1] *
                       expr2.data()[base2 + i * stride2];
            return sum;
        }
 
    private:
        // ========================================================================
        // Contiguous dot — both fibers have stride 1
        // ========================================================================
 
        template <typename Expr1, typename Expr2>
        FORCE_INLINE static T dot_contiguous_impl(
            const Expr1 &expr1,
            const Expr2 &expr2,
            my_size_t base1,
            my_size_t base2,
            my_size_t len) noexcept
        {
            // std::cout << "dot_contiguous_impl" << std::endl;
            const T *ptr1 = expr1.data() + base1;
            const T *ptr2 = expr2.data() + base2;
 
            const my_size_t simdSteps = len / simdWidth;
            const my_size_t scalarStart = simdSteps * simdWidth;
 
            T result = T{0};
 
            if (simdSteps > 0)
            {
                typename K::VecType acc = K::set1(T{0});
 
                for (my_size_t i = 0; i < simdSteps; ++i)
                {
                    auto v1 = K::load(ptr1 + i * simdWidth);
                    auto v2 = K::load(ptr2 + i * simdWidth);
                    acc = Helpers::fmadd_safe(v1, v2, acc);
                }
 
                alignas(DATA_ALIGNAS) T tmp[simdWidth];
                K::store(tmp, acc);
 
                for (my_size_t i = 0; i < simdWidth; ++i)
                    result += tmp[i];
            }
 
            for (my_size_t i = scalarStart; i < len; ++i)
                result += ptr1[i] * ptr2[i];
 
            return result;
        }
 
        // ========================================================================
        // Strided dot — one or both fibers have stride > 1
        // ========================================================================
 
        template <typename Expr1, typename Expr2>
        FORCE_INLINE static T dot_strided_impl(
            const Expr1 &expr1,
            const Expr2 &expr2,
            my_size_t idx1,
            my_size_t idx2,
            my_size_t stride1,
            my_size_t stride2,
            my_size_t len) noexcept
        {
            // std::cout << "dot_strided_impl" << std::endl;
            const my_size_t simdSteps = len / simdWidth;
            const my_size_t scalarStart = simdSteps * simdWidth;
 
            T result = T{0};
 
            if (simdSteps > 0)
            {
                typename K::VecType acc = K::set1(T{0});
 
                for (my_size_t i = 0; i < simdSteps; ++i)
                {
                    // Build gather indices for this chunk
                    my_size_t idxList1[simdWidth];
                    my_size_t idxList2[simdWidth];
                    for (my_size_t j = 0; j < simdWidth; ++j)
                    {
                        idxList1[j] = idx1 + j * stride1;
                        idxList2[j] = idx2 + j * stride2;
                    }
 
                    auto v1 = K::gather(expr1.data(), idxList1);
                    auto v2 = K::gather(expr2.data(), idxList2);
                    acc = Helpers::fmadd_safe(v1, v2, acc);
 
                    idx1 += simdWidth * stride1;
                    idx2 += simdWidth * stride2;
                }
 
                alignas(DATA_ALIGNAS) T tmp[simdWidth];
                K::store(tmp, acc);
 
                for (my_size_t i = 0; i < simdWidth; ++i)
                    result += tmp[i];
            }
 
            // Scalar tail
            for (my_size_t i = scalarStart; i < len; ++i)
            {
                result += expr1.data()[idx1] * expr2.data()[idx2];
                idx1 += stride1;
                idx2 += stride2;
            }
 
            return result;
        }
    };
 
} // namespace detail
 
#endif // KERNEL_DOT_H