kernel/contract.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/core/scalar_concepts.hpp>
#include <uni20/mdspan/mdspan.hpp>
#include <uni20/mdspan/strides.hpp>
 
#include <uni20/kernel/cpu/contract.hpp>
 
namespace uni20::kernel
{
 
template <typename T, StridedMdspanLike AType, StridedMdspanLike BType, std::size_t N, typename U,
          MutableStridedMdspanLike CType>
  requires(AType::rank() + BType::rank() == CType::rank() + 2 * N) && DefaultAccessorMdspanLike<AType> &&
          DefaultAccessorMdspanLike<BType> && DefaultAccessorMdspanLike<CType>
void contract(T const& alpha, AType A, BType B, std::array<std::pair<std::size_t, std::size_t>, N> const& contractDims,
              U const& beta, CType C)
{
  auto [Mgroup, Ngroup, Kgroup] = extract_strides(A, B, contractDims, C);
  contract_strided(Mgroup, Ngroup, Kgroup, alpha, A.data_handle(), B.data_handle(), beta, C.data_handle());
}
 
template <typename T, StridedMdspanLike AType, StridedMdspanLike BType, typename U, MutableStridedMdspanLike CType,
          std::size_t N>
  requires DefaultAccessorMdspanLike<AType> && DefaultAccessorMdspanLike<BType> && DefaultAccessorMdspanLike<CType>
void contract(T const& alpha, AType A, BType B,
              std::pair<std::size_t, std::size_t> const (&dims)[N], // array reference, size known at compile time
              U const& beta, CType C)
{
  contract(alpha, A, B, std::to_array(dims), beta, C);
}
 
} // namespace uni20::kernel