dot.hpp

Go to the documentation of this file.

// Copyright (c) 2022--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 "../blas/dot.hpp"
#include "../concepts.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <complex>
#include <cstddef>
 
namespace nda::linalg {


  template <bool star = false, Vector X, Vector Y>
    requires(Scalar<get_value_t<X>> and Scalar<get_value_t<Y>> and mem::have_host_compatible_addr_space<X, Y>)
  auto dot_generic(X const &x, Y const &y) {
    // check the dimensions of the input arrays/views
    EXPECTS(x.size() == y.size());
 
    // conditional conjugation
    auto cond_conj = [](auto z) __attribute__((always_inline)) {
      if constexpr (star and is_complex_v<decltype(z)>) {
        return std::conj(z);
      } else {
        return z;
      }
    };
 
    // early return for zero-sized vectors
    long const N = x.size();
    if (N == 0) return decltype(cond_conj(x(0)) * y(0)){0};
 
    // loop over vectors and sum up element-wise products
    if constexpr (has_layout_smallest_stride_is_one<X> and has_layout_smallest_stride_is_one<Y>) {
      if constexpr (is_regular_or_view_v<X> and is_regular_or_view_v<Y>) {
        auto *__restrict px = x.data();
        auto *__restrict py = y.data();
        auto res            = cond_conj(px[0]) * py[0];
        for (size_t i = 1; i < N; ++i) res += cond_conj(px[i]) * py[i];
        return res;
      } else {
        auto res = cond_conj(x(_linear_index_t{0})) * y(_linear_index_t{0});
        for (long i = 1; i < N; ++i) res += cond_conj(x(_linear_index_t{i})) * y(_linear_index_t{i});
        return res;
      }
    } else {
      auto res = cond_conj(x(0)) * y(0);
      for (long i = 1; i < N; ++i) res += cond_conj(x(i)) * y(i);
      return res;
    }
  }

  template <typename X, typename Y>
    requires((Scalar<X> and Scalar<Y>) or (Vector<X> and Vector<Y>))
  auto dot(X const &x, Y const &y) {
    if constexpr (Scalar<X>) {
      return x * y;
    } else if constexpr (requires { nda::blas::dot(x, y); }) {
      return nda::blas::dot(x, y);
    } else {
      return dot_generic<false>(x, y);
    }
  }

  template <typename X, typename Y>
    requires((Scalar<X> and Scalar<Y>) or (Vector<X> and Vector<Y>))
  auto dotc(X const &x, Y const &y) {
    if constexpr (Scalar<X>) {
      if constexpr (is_complex_v<X>) return std::conj(x) * y;
      return x * y;
    } else if constexpr (requires { nda::blas::dotc(x, y); }) {
      return nda::blas::dotc(x, y);
    } else {
      return dot_generic<true>(x, y);
    }
  }

 
} // namespace nda::linalg