rabc_contraction.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <uni20/linalg/backend_selector.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/symmetry/block_tensor_concepts.hpp>
#include <uni20/tensor_network/backends/right_first_rabc.hpp>
#include <uni20/tensor_network/rabc_contraction_plan.hpp>
#include <uni20/tensor_network/rabc_operation.hpp>
 
#include <concepts>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
namespace detail
{
 
template <class StoragePolicy> [[nodiscard]] auto make_rabc_backend_selector()
{
  using leaf_storage_policy = typename StoragePolicy::leaf_storage_policy;
  auto contraction_selector = linalg::select_backend_for_storage<leaf_storage_policy>(contract_op<2, 2, 1>{});
  using contraction_selector_type = std::remove_cvref_t<decltype(contraction_selector)>;
  return backend_list{tensor_network::RightFirstRabcBackend<leaf_storage_policy, contraction_selector_type>{
      std::move(contraction_selector)}};
}
 
} // namespace detail
 
template <class StoragePolicy> struct backend_selector_default<tensor_network::rabc_contract_op, StoragePolicy>
{
    template <class StorageSelector>
    static auto select(tensor_network::rabc_contract_op const&, StorageSelector storage_selector)
    {
      static_cast<void>(storage_selector);
      return detail::make_rabc_backend_selector<StoragePolicy>();
    }
};
 
} // namespace uni20::linalg
 
namespace uni20::tensor_network
{
namespace detail
{
 
template <class Output, class B> [[nodiscard]] constexpr bool is_obvious_rabc_alias(Output& output, B const& b) noexcept
{
  if constexpr (std::same_as<std::remove_cvref_t<Output>, std::remove_cvref_t<B>>)
  {
    return static_cast<void const*>(std::addressof(output)) == static_cast<void const*>(std::addressof(b));
  }
  return false;
}
 
} // namespace detail
 
template <class IntermediateStorage, class ContractionSelector, detail::WritableRabcTensor Output, RabcPlan Plan,
          detail::ReadableRabcTensor A, detail::ReadableRabcTensor B, detail::ReadableRabcTensor C>
  requires detail::CompatibleRabcOperands<Output, Plan, A, B, C>
[[nodiscard]] auto
prepare_rabc_contract(linalg::backend_list<RightFirstRabcBackend<IntermediateStorage, ContractionSelector>> selector,
                      Output const& output, Plan plan, A const& a, B const& b, C const& c)
{
  auto backend = std::move(std::get<0>(selector.entries));
  return PreparedRightFirstRabcContraction<IntermediateStorage, std::remove_cvref_t<Plan>, ContractionSelector>(
      std::move(backend.contraction_selector), output, std::move(plan), a, b, c);
}
 
template <detail::WritableRabcTensor Output, RabcPlan Plan, detail::ReadableRabcTensor A, detail::ReadableRabcTensor B,
          detail::ReadableRabcTensor C>
  requires detail::CompatibleRabcOperands<Output, Plan, A, B, C>
[[nodiscard]] auto prepare_rabc_contract(Output const& output, Plan plan, A const& a, B const& b, C const& c)
{
  auto selector = linalg::select_backend(rabc_contract_op{}, output);
  return prepare_rabc_contract(std::move(selector), output, std::move(plan), a, b, c);
}
 
template <linalg::KernelBackendSelector BackendSelector, MutableBlockTensorView Output, RabcPlan Plan,
          BlockTensorView A, BlockTensorView B, BlockTensorView C>
void rabc_contract(BackendSelector&& selector, Output& output, Plan const& plan, A const& a, B const& b, C const& c)
{
  if (detail::is_obvious_rabc_alias(output, b))
    throw std::invalid_argument("R/A/B/C contraction output must not alias its input center");
  linalg::dispatch_kernel(std::forward<BackendSelector>(selector), rabc_contract_op{}, output, plan, a, b, c);
}
 
template <MutableBlockTensorView Output, RabcPlan Plan, BlockTensorView A, BlockTensorView B, BlockTensorView C>
void rabc_contract(Output& output, Plan const& plan, A const& a, B const& b, C const& c)
{
  auto selector = linalg::select_backend(rabc_contract_op{}, output);
  rabc_contract(std::move(selector), output, plan, a, b, c);
}
 
} // namespace uni20::tensor_network