tools.hpp

Go to the documentation of this file.

// Copyright (c) 2019--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../macros.hpp"
#include "../map.hpp"
#include "../mapped_functions.hpp"
#include "../mem/address_space.hpp"
#include "../mem/policies.hpp"
#include "../traits.hpp"
 
#include <complex>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace nda {

  using dcomplex = std::complex<double>;
 
} // namespace nda
 
namespace nda::blas_lapack {


  template <typename A>
  static constexpr bool is_conj_array_expr = false;

  template <MemoryArray A>
  static constexpr bool is_conj_array_expr<expr_call<detail::conj_f, A>> = true;
 
  // Specialization of nda::blas_lapack::is_conj_array_expr for cvref types.
  template <typename A>
    requires(!std::is_same_v<A, std::remove_cvref_t<A>>)
  static constexpr bool is_conj_array_expr<A> = is_conj_array_expr<std::remove_cvref_t<A>>;

  template <Array A>
    requires(MemoryArray<A> or is_conj_array_expr<A>)
  MemoryArray decltype(auto) get_array(A &&a) {
    if constexpr (is_conj_array_expr<A>) {
      return std::get<0>(std::forward<A>(a).a);
    } else {
      return std::forward<A>(a);
    }
  }

  template <Array... As>
    requires((MemoryArray<As> or is_conj_array_expr<As>) and ...)
  static constexpr bool has_F_layout = ([]<typename A>() constexpr {
    if constexpr (is_conj_array_expr<A>)
      return has_F_layout<decltype(std::get<0>(std::declval<A>().a))>;
    else
      return std::remove_cvref_t<A>::is_stride_order_Fortran();
  }.template operator()<As>() and ...);

  template <Array... As>
    requires((MemoryArray<As> or is_conj_array_expr<As>) and ...)
  static constexpr bool has_C_layout = ([]<typename A>() constexpr {
    if constexpr (is_conj_array_expr<A>)
      return has_C_layout<decltype(std::get<0>(std::declval<A>().a))>;
    else
      return std::remove_cvref_t<A>::is_stride_order_C();
  }.template operator()<As>() and ...);

  template <Array A>
  static constexpr char get_op = []() {
    auto constexpr conj      = is_conj_array_expr<A>;
    auto constexpr transpose = has_C_layout<A>;
    static_assert(!(conj and not transpose), "Error in nda::blas_lapack::get_op: Cannot use conjugate operation alone in BLAS operations");
    if constexpr (conj and transpose)
      return 'C';
    else if constexpr (transpose)
      return 'T';
    else // !conj and !transpose
      return 'N';
  }();

  template <MemoryArray A>
    requires(get_rank<A> == 1 or get_rank<A> == 2)
  int get_ld(A const &a) {
    if constexpr (get_rank<A> == 1) {
      return a.size();
    } else {
      return a.indexmap().strides()[has_F_layout<A> ? 1 : 0];
    }
  }

  template <MemoryArray A>
    requires(get_rank<A> == 1 or get_rank<A> == 2)
  int get_ncols(A const &a) {
    if constexpr (get_rank<A> == 1) {
      return 1;
    } else {
      return a.shape()[has_F_layout<A> ? 1 : 0];
    }
  }

  template <MemoryArray A>
  using vector_value_t = vector<get_value_t<A>, heap<mem::get_addr_space<A>>>;

  template <MemoryArray A>
  using vector_fp_t = vector<get_fp_t<A>, heap<mem::get_addr_space<A>>>;

  template <Array A>
    requires(MemoryArrayOfRank<A, 3> or MemoryArrayOfRank<A, 2>)
  auto batch_ptrs(A &&a) { // NOLINT (temporary views are allowed here)
    using ptr_t           = decltype(a.data());
    auto const idx        = (has_C_layout<A> ? 0 : (get_rank<A> == 3 ? 2 : 1));
    auto const batch_size = a.shape()[idx];
    auto const stride     = a.indexmap().strides()[idx];
 
    auto ptrs = vector<ptr_t>(batch_size);
    for (int i = 0; auto &ptr : ptrs) ptr = a.data() + i++ * stride;
    return ptrs;
  }

  template <typename A>
    requires(is_regular_or_view_v<A> and get_rank<A> == 1)
  void resize_or_check_work_buffer(A &a, long min_size) {
    if (a.size() >= min_size) {
      EXPECTS(a.indexmap().min_stride() == 1);
      return;
    }
    if constexpr (is_regular_v<A>) {
      a.resize(min_size);
    } else {
      NDA_RUNTIME_ERROR << "Error in nda::blas_lapack::resize_or_check_work_buffer: Size mismatch: " << a.size() << " < " << min_size;
    }
  }

  template <typename A, int R = -1>
  concept BlasArray = (R == -1 ? MemoryArray<A> : MemoryArrayOfRank<A, R>) and is_blas_lapack_v<get_value_t<A>>;

  template <typename A, int R = -1>
  concept BlasArrayReal = BlasArray<A, R> and AnyOf<get_value_t<A>, float, double>;

  template <typename A, int R = -1>
  concept BlasArrayCplx = BlasArray<A, R> and AnyOf<get_value_t<A>, std::complex<float>, std::complex<double>>;

  template <typename A, int R = -1>
  concept BlasArrayOrConj =
     BlasArray<A, R> or ((R == -1 ? Array<A> : ArrayOfRank<A, R>) and is_conj_array_expr<A> and is_blas_lapack_v<get_value_t<A>>);

  template <typename A, typename B, int R = -1>
  concept BlasArrayFor = BlasArrayOrConj<B> and BlasArray<A, R> and have_same_value_type_v<A, B> and mem::have_compatible_addr_space<A, B>;

  template <typename A, typename B, int R = -1>
  concept BlasArrayOrConjFor =
     BlasArrayOrConj<B> and BlasArrayOrConj<A, R> and have_same_value_type_v<A, B> and mem::have_compatible_addr_space<A, B>;

  template <typename A, typename B, int R = -1>
  concept PivotArrayFor = BlasArrayOrConj<B> and (R == -1 ? MemoryArray<A> : MemoryArrayOfRank<A, R>)
     and std::is_same_v<get_value_t<A>, int> and mem::have_compatible_addr_space<A, B>;

  template <typename A, typename B, int R = -1>
  concept BlasArrayRealFor =
     BlasArrayOrConj<B> and BlasArray<A, R> and std::is_same_v<get_value_t<A>, get_fp_t<B>> and mem::have_compatible_addr_space<A, B>;

 
} // namespace nda::blas_lapack