gemv.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 "./interface/cxx_interface.hpp"
#include "./tools.hpp"
#include "../concepts.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#ifndef NDA_HAVE_DEVICE
#include "../device.hpp"
#endif
 
#include <tuple>
 
namespace nda::blas {


  template <typename A, typename X, typename Y>
  void gemv_generic(get_value_t<A> alpha, A const &a, X const &x, get_value_t<A> beta, Y &&y) { // NOLINT (temporary views are allowed here)
    EXPECTS(a.extent(1) == x.extent(0));
    EXPECTS(a.extent(0) == y.extent(0));
    for (int i = 0; i < a.extent(0); ++i) {
      y(i) = beta * y(i);
      for (int k = 0; k < a.extent(1); ++k) y(i) += alpha * a(i, k) * x(k);
    }
  }

  template <Matrix A, MemoryVector X, MemoryVector Y>
    requires((MemoryMatrix<A> or is_conj_array_expr<A>) and have_same_value_type_v<A, X, Y> and is_blas_lapack_v<get_value_t<A>>)
  void gemv(get_value_t<A> alpha, A const &a, X const &x, get_value_t<A> beta, Y &&y) { // NOLINT (temporary views are allowed here)
    // get underlying matrix in case it is given as a lazy expression
    auto to_mat = []<Matrix Z>(Z const &z) -> decltype(auto) {
      if constexpr (is_conj_array_expr<Z>)
        return std::get<0>(z.a);
      else
        return z;
    };
    auto &mat = to_mat(a);
 
    // compile-time checks
    using mat_type = decltype(mat);
    static_assert(mem::have_compatible_addr_space<mat_type, X, Y>);
 
    // runtime checks
    EXPECTS(mat.extent(1) == x.extent(0));
    EXPECTS(mat.extent(0) == y.extent(0));
    EXPECTS(mat.indexmap().min_stride() == 1);
    EXPECTS(x.indexmap().min_stride() == 1);
    EXPECTS(y.indexmap().min_stride() == 1);
 
    // gather parameters for gemv call
    static constexpr bool conj_A = is_conj_array_expr<A>;
    char op_a                    = get_op<conj_A, /* transpose = */ !has_F_layout<mat_type>>;
    auto [m, n]                  = mat.shape();
    if constexpr (has_C_layout<mat_type>) std::swap(m, n);
 
    if constexpr (mem::have_device_compatible_addr_space<mat_type, X, Y>) {
#if defined(NDA_HAVE_DEVICE)
      device::gemv(op_a, m, n, alpha, mat.data(), get_ld(mat), x.data(), x.indexmap().strides()[0], beta, y.data(), y.indexmap().strides()[0]);
#else
      compile_error_no_gpu();
#endif
    } else {
      f77::gemv(op_a, m, n, alpha, mat.data(), get_ld(mat), x.data(), x.indexmap().strides()[0], beta, y.data(), y.indexmap().strides()[0]);
    }
  }

 
} // namespace nda::blas