linalg/ops/contract.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/linalg/backends/cpu/contract.hpp>
#include <uni20/linalg/backends/direct_gemm/contract.hpp>
#include <uni20/linalg/backends/looped_gemm/contract.hpp>
#include <uni20/linalg/contraction_axes.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/tensor/concepts.hpp>
#include <uni20/tensor/output.hpp>
 
#include <array>
#include <concepts>
#include <cstddef>
#include <memory>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
 
namespace detail
{
 
template <KernelBackendSelector ExecutionSelector>
[[nodiscard]] auto make_contraction_backend_selector(ExecutionSelector execution_selector)
{
  auto direct_selector = execution_selector;
  auto looped_selector = execution_selector;
  auto fallback_selector = normalize_backend_selector(std::move(execution_selector));
  return std::apply(
      [&]<class... FallbackBackends>(FallbackBackends&&... fallback_backends) {
        return backend_list{DirectGemmContractionBackend{std::move(direct_selector)},
                            LoopedGemmContractionBackend{std::move(looped_selector)},
                            std::forward<FallbackBackends>(fallback_backends)...};
      },
      std::move(fallback_selector.entries));
}
 
} // namespace detail
 
template <std::size_t LhsRank, std::size_t RhsRank, std::size_t ContractedRank, class StoragePolicy>
struct backend_selector_default<contract_op<LhsRank, RhsRank, ContractedRank>, StoragePolicy>
{
    template <class StorageSelector>
    static auto select(contract_op<LhsRank, RhsRank, ContractedRank> const&, StorageSelector storage_selector)
    {
      return detail::make_contraction_backend_selector(std::move(storage_selector));
    }
};
 
namespace detail
{
 
template <class OutputTensor, class InputTensor>
[[nodiscard]] constexpr bool is_obvious_contraction_alias(OutputTensor& output, InputTensor const& input) noexcept
{
  if constexpr (std::same_as<std::remove_cvref_t<OutputTensor>, std::remove_cvref_t<InputTensor>>)
  {
    return static_cast<void const*>(std::addressof(output)) == static_cast<void const*>(std::addressof(input));
  }
  return false;
}
 
template <class OutputTensor, class LhsTensor, class RhsTensor>
concept CompatibleContractionTensors =
    uni20::MutableTensorView<OutputTensor> && uni20::TensorView<LhsTensor> && uni20::TensorView<RhsTensor> &&
    uni20::Scalar<uni20::tensor_element_t<OutputTensor>> && uni20::Scalar<uni20::tensor_element_t<LhsTensor>> &&
    uni20::Scalar<uni20::tensor_element_t<RhsTensor>> &&
    requires(uni20::tensor_element_t<LhsTensor> lhs, uni20::tensor_element_t<RhsTensor> rhs) {
      static_cast<uni20::tensor_element_t<OutputTensor>>(lhs);
      static_cast<uni20::tensor_element_t<OutputTensor>>(rhs);
    };
 
template <class OutputTensor, class LhsTensor, class RhsTensor>
void validate_contraction_aliasing(OutputTensor& output, LhsTensor const& lhs, RhsTensor const& rhs)
{
  ERROR_IF(is_obvious_contraction_alias(output, lhs) || is_obvious_contraction_alias(output, rhs),
           "tensor contraction output must not alias an input tensor");
}
 
} // namespace detail
 
template <class BackendSelector, class OutputTensor, class LhsTensor, class RhsTensor, std::size_t LhsRank,
          std::size_t RhsRank, std::size_t ContractedRank>
  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&
           uni20::MutableRankedTensorView<OutputTensor,
                                          ContractionAxes<LhsRank, RhsRank, ContractedRank>::output_rank> &&
           uni20::RankedTensorView<LhsTensor, LhsRank> && uni20::RankedTensorView<RhsTensor, RhsRank>
void contract(BackendSelector&& selector, OutputTensor&& output, uni20::tensor_element_t<OutputTensor> alpha,
              LhsTensor const& lhs, RhsTensor const& rhs, ContractionAxes<LhsRank, RhsRank, ContractedRank> axes,
              uni20::tensor_element_t<OutputTensor> beta)
{
  detail::validate_contraction_aliasing(output, lhs, rhs);
  auto const required_extents = contraction_output_extents(lhs, rhs, axes);
  uni20::require_output(output, required_extents);
 
  auto output_descriptor = uni20::mdspec_of(output);
  auto lhs_descriptor = uni20::mdspec_of(lhs);
  auto rhs_descriptor = uni20::mdspec_of(rhs);
  auto operation = contract_op<LhsRank, RhsRank, ContractedRank>{.axes = std::move(axes)};
  dispatch_kernel(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor, alpha,
                  lhs_descriptor, rhs_descriptor, beta);
}
 
template <class OutputTensor, class LhsTensor, class RhsTensor, std::size_t LhsRank, std::size_t RhsRank,
          std::size_t ContractedRank>
  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&
           uni20::MutableRankedTensorView<OutputTensor,
                                          ContractionAxes<LhsRank, RhsRank, ContractedRank>::output_rank> &&
           uni20::RankedTensorView<LhsTensor, LhsRank> && uni20::RankedTensorView<RhsTensor, RhsRank>
void contract(OutputTensor&& output, uni20::tensor_element_t<OutputTensor> alpha, LhsTensor const& lhs,
              RhsTensor const& rhs, ContractionAxes<LhsRank, RhsRank, ContractedRank> axes,
              uni20::tensor_element_t<OutputTensor> beta)
{
  auto operation = contract_op<LhsRank, RhsRank, ContractedRank>{.axes = axes};
  auto selector = select_backend(operation, output, lhs, rhs);
  contract(std::move(selector), std::forward<OutputTensor>(output), alpha, lhs, rhs, std::move(axes), beta);
}
 
template <class BackendSelector, class OutputTensor, class LhsTensor, class RhsTensor, std::size_t ContractedRank>
  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&
           (ContractedRank <= uni20::tensor_mdspec_t<LhsTensor>::rank()) &&
           (ContractedRank <= uni20::tensor_mdspec_t<RhsTensor>::rank()) &&
           (uni20::tensor_mdspec_t<OutputTensor>::rank() ==
            uni20::tensor_mdspec_t<LhsTensor>::rank() + uni20::tensor_mdspec_t<RhsTensor>::rank() - 2 * ContractedRank)
void contract(BackendSelector&& selector, OutputTensor&& output, uni20::tensor_element_t<OutputTensor> alpha,
              LhsTensor const& lhs, RhsTensor const& rhs,
              std::array<std::pair<std::size_t, std::size_t>, ContractedRank> requested_axes,
              uni20::tensor_element_t<OutputTensor> beta)
{
  constexpr std::size_t lhs_rank = uni20::tensor_mdspec_t<LhsTensor>::rank();
  constexpr std::size_t rhs_rank = uni20::tensor_mdspec_t<RhsTensor>::rank();
  auto axes = make_contraction_axes<lhs_rank, rhs_rank>(std::move(requested_axes));
  contract(std::forward<BackendSelector>(selector), std::forward<OutputTensor>(output), alpha, lhs, rhs,
           std::move(axes), beta);
}
 
template <class OutputTensor, class LhsTensor, class RhsTensor, std::size_t ContractedRank>
  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&
           (ContractedRank <= uni20::tensor_mdspec_t<LhsTensor>::rank()) &&
           (ContractedRank <= uni20::tensor_mdspec_t<RhsTensor>::rank()) &&
           (uni20::tensor_mdspec_t<OutputTensor>::rank() ==
            uni20::tensor_mdspec_t<LhsTensor>::rank() + uni20::tensor_mdspec_t<RhsTensor>::rank() - 2 * ContractedRank)
void contract(OutputTensor&& output, uni20::tensor_element_t<OutputTensor> alpha, LhsTensor const& lhs,
              RhsTensor const& rhs, std::array<std::pair<std::size_t, std::size_t>, ContractedRank> requested_axes,
              uni20::tensor_element_t<OutputTensor> beta)
{
  constexpr std::size_t lhs_rank = uni20::tensor_mdspec_t<LhsTensor>::rank();
  constexpr std::size_t rhs_rank = uni20::tensor_mdspec_t<RhsTensor>::rank();
  auto axes = make_contraction_axes<lhs_rank, rhs_rank>(std::move(requested_axes));
  contract(std::forward<OutputTensor>(output), alpha, lhs, rhs, std::move(axes), beta);
}
 
} // namespace uni20::linalg