Broadcasting.hpp

#pragma once
 
#include <algorithm>
#include <cstddef>
#include <vector>
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/exceptions/Exceptions.hpp>
#include <stratax/core/Shape.hpp>
#include <stratax/core/ArrayTraits.hpp>
#include <stratax/core/dtypes/Promotion.hpp>
 
namespace stratax::core::broadcast_detail {
 
inline std::size_t dimension_from_right(
    const Shape& shape,
    std::size_t offset)
{
    return offset < shape.rank()
        ? shape[shape.rank() - 1 - offset]
        : 1;
}
 
inline std::size_t flat_operand_index(
    std::size_t result_index,
    const stratax::core::Shape& result_shape,
    const stratax::core::Shape& operand_shape)
{
    std::size_t operand_index = 0;
    std::size_t operand_stride = 1;
 
    for (std::size_t result_axis = result_shape.rank(); result_axis-- > 0;)
    {
        const std::size_t coordinate =
            result_index % result_shape[result_axis];
        result_index /= result_shape[result_axis];
 
        const std::size_t offset = result_shape.rank() - 1 - result_axis;
        if (offset >= operand_shape.rank())
        {
            continue;
        }
 
        const std::size_t operand_axis = operand_shape.rank() - 1 - offset;
        const std::size_t operand_dimension = operand_shape[operand_axis];
        const std::size_t operand_coordinate =
            operand_dimension == 1 ? 0 : coordinate;
 
        operand_index += operand_coordinate * operand_stride;
        operand_stride *= operand_dimension;
    }
 
    return operand_index;
}
 
} // namespace stratax::core::broadcast_detail
 
inline bool broadcastable(
    const stratax::core::Shape& shape1,
    const stratax::core::Shape& shape2
)
{
    const std::size_t result_rank = std::max(shape1.rank(), shape2.rank());
 
    for (std::size_t offset = 0; offset < result_rank; ++offset)
    {
        const std::size_t left =
            stratax::core::broadcast_detail::dimension_from_right(shape1, offset);
        const std::size_t right =
            stratax::core::broadcast_detail::dimension_from_right(shape2, offset);
 
        if (left != right && left != 1 && right != 1)
        {
            return false;
        }
    }
 
    return true;
}
 
inline stratax::core::Shape broadcasted_shape(
    const stratax::core::Shape& shape1,
    const stratax::core::Shape& shape2)
{
    if (!broadcastable(shape1, shape2))
    {
        throw Exceptions::BroadcastError("Shapes are not broadcastable.");
    }
 
    const std::size_t result_rank = std::max(shape1.rank(), shape2.rank());
    std::vector<std::size_t> result(result_rank);
 
    for (std::size_t offset = 0; offset < result_rank; ++offset)
    {
        const std::size_t left =
            stratax::core::broadcast_detail::dimension_from_right(shape1, offset);
        const std::size_t right =
            stratax::core::broadcast_detail::dimension_from_right(shape2, offset);
 
        result[result_rank - 1 - offset] = left == 1 ? right : left;
    }
 
    return stratax::core::Shape{result};
}
 
namespace stratax::core {
 
template<
    DType Result,
    Array L,
    Array R,
    typename Op>
auto broadcasted_op(
    const L& lhs,
    const R& rhs,
    Op op)
{
    using result_type =
        promote_array_t<L, R, Result>;
 
    const auto result_shape =
        broadcasted_shape(lhs.shape(), rhs.shape());
 
    result_type result(result_shape);
 
    for (std::size_t i = 0; i < result.size(); ++i)
    {
        const auto lhs_index =
            broadcast_detail::flat_operand_index(
                i,
                result_shape,
                lhs.shape());
 
        const auto rhs_index =
            broadcast_detail::flat_operand_index(
                i,
                result_shape,
                rhs.shape());
 
        result[i] = static_cast<Result>(
            op(lhs[lhs_index], rhs[rhs_index]));
    }
 
    return result;
}
 
template<Array L, Array R, typename Op>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto broadcasted_op(
    const L& lhs,
    const R& rhs,
    Op op)
{
    using result_value_type =
        promote_t<
            typename L::value_type,
            typename R::value_type>;
 
    return broadcasted_op<result_value_type>(
        lhs,
        rhs,
        op);
}
 
}
 
template<Array L, Numeric S, typename Op>
auto broadcasted_op(const L& lhs, const S& rhs, Op op)
{
    using result_value_type =
        stratax::core::promote_t<
            typename L::value_type,
            S>;
 
    using result_type =
        stratax::core::rebind_array_t<
            L,
            result_value_type>;
 
    result_type result(lhs.shape());
 
    for (std::size_t i = 0; i < result.size(); ++i)
    {
        result[i] = static_cast<result_value_type>(
            op(lhs[i], rhs));
    }
 
    return result;
}
 
template<Numeric S, Array R, typename Op>
auto broadcasted_op(const S& lhs, const R& rhs, Op op)
{
    using result_value_type =
        stratax::core::promote_t<
            S,
            typename R::value_type>;
 
    using result_type =
        stratax::core::rebind_array_t<
            R,
            result_value_type>;
 
    result_type result(rhs.shape());
 
    for (std::size_t i = 0; i < result.size(); ++i)
    {
        result[i] = static_cast<result_value_type>(
            op(lhs, rhs[i]));
    }
 
    return result;
}