gemm.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 "../layout_transforms.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#ifndef NDA_HAVE_DEVICE
#include "../device.hpp"
#endif // NDA_HAVE_DEVICE
 
#include <utility>
 
namespace nda::blas {


  template <Matrix A, Matrix B, MemoryMatrix C>
    requires((MemoryMatrix<A> or is_conj_array_expr<A>) and (MemoryMatrix<B> or is_conj_array_expr<B>)
             and have_same_value_type_v<A, B, C> and mem::have_compatible_addr_space<A, B, C> and is_blas_lapack_v<get_value_t<A>>)
  void gemm(get_value_t<A> alpha, A const &a, B const &b, get_value_t<A> beta, C &&c) {
    // if C is in C-layout, compute the transpose of the product in Fortran order
    if constexpr (has_C_layout<C>) {
      gemm(alpha, transpose(b), transpose(a), beta, transpose(std::forward<C>(c)));
    } else {
      // get underlying matrix in case it is given as a conjugate expression
      auto &mat_a = get_array(a);
      auto &mat_b = get_array(b);
 
      // check the dimensions of the input/output arrays/views
      auto const [m, k] = mat_a.shape();
      auto const [l, n] = mat_b.shape();
      EXPECTS(k == l);
      EXPECTS(m == c.extent(0));
      EXPECTS(n == c.extent(1));
 
      // arrays/views must be BLAS compatible
      EXPECTS(mat_a.indexmap().min_stride() == 1);
      EXPECTS(mat_b.indexmap().min_stride() == 1);
      EXPECTS(c.indexmap().min_stride() == 1);
 
      // perform the actual library call
      if constexpr (mem::have_device_compatible_addr_space<A, B, C>) {
#if defined(NDA_HAVE_DEVICE)
        device::gemm(get_op<A>, get_op<B>, m, n, k, alpha, mat_a.data(), get_ld(mat_a), mat_b.data(), get_ld(mat_b), beta, c.data(), get_ld(c));
#else
        compile_error_no_gpu();
#endif
      } else {
        f77::gemm(get_op<A>, get_op<B>, m, n, k, alpha, mat_a.data(), get_ld(mat_a), mat_b.data(), get_ld(mat_b), beta, c.data(), get_ld(c));
      }
    }
  }

 
} // namespace nda::blas