getrs_batch.hpp

Go to the documentation of this file.

// Copyright (c) 2021--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 "./getrs.hpp"
#include "./interface/cxx_interface.hpp"
#include "../basic_array.hpp"
#include "../blas/interface/cxx_interface.hpp"
#include "../blas/tools.hpp"
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../device.hpp"
#include "../layout_transforms.hpp"
#include "../macros.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <algorithm>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace nda::lapack {
 
  namespace detail {
 
    // Implementation of the batched getrs routine.
    template <bool run_on_device>
    int getrs_batch_impl(auto const &a, auto &b, auto const &ipiv, char op) {
      // get underlying array in case it is given as a lazy conjugate expression
      auto &a_arr = get_array(a);
 
      // check the dimensions of the input/output arrays/views
      auto const [m, n, n_b]      = a_arr.shape();
      auto const [k, nrhs, n_b_2] = b.shape();
      EXPECTS(m == n);
      EXPECTS(n == k);
      EXPECTS(n_b == n_b_2);
      EXPECTS(ipiv.extent(0) == n);
      EXPECTS(ipiv.extent(1) == n_b);
 
      // arrays/views must be LAPACK compatible
      EXPECTS(a_arr.indexmap().min_stride() == 1);
      EXPECTS(b.indexmap().min_stride() == 1);
      EXPECTS(ipiv.indexmap().min_stride() == 1);
 
      // perform actual library call(s)
      int info = 0;
      if constexpr (run_on_device) {
        auto a_ptrs = to_device(batch_ptrs(a_arr));
        auto b_ptrs = to_device(batch_ptrs(b));
        blas::device::getrs_batch(op, n, nrhs, a_ptrs.data(), get_ld(a_arr(range::all, range::all, 0)), ipiv.data(), b_ptrs.data(),
                                  get_ld(b(range::all, range::all, 0)), info, n_b);
      } else {
        // for host, fall back to looping over batches
        for (int i = 0; i < n_b; ++i) {
          auto a_i       = a_arr(range::all, range::all, i);
          auto b_i       = b(range::all, range::all, i);
          auto ipiv_i    = ipiv(range::all, i);
          int local_info = 0;
          f77::getrs(op, get_ncols(a_i), get_ncols(b_i), a_i.data(), get_ld(a_i), ipiv_i.data(), b_i.data(), get_ld(b_i), local_info);
          if (local_info != 0 && info == 0) info = local_info;
        }
      }
      return info;
    }
 
  } // namespace detail

  template <BlasArrayOrConj<3> A, BlasArrayFor<A, 3> B, PivotArrayFor<A, 2> IPIV>
    requires((has_F_layout<A> or has_C_layout<A>) and has_F_layout<B> and (not is_conj_array_expr<A> or has_C_layout<A>))
  int getrs_batch(A const &a, B &&b, IPIV const &ipiv) { // NOLINT (temporary views are allowed here)
    constexpr bool run_on_device = mem::have_device_compatible_addr_space<A, B, IPIV>;
 
    // transpose ipiv array/view if necessary
    if constexpr (has_C_layout<IPIV>) return getrs_batch(a, b, transpose(ipiv));
 
    // transpose A array/view if necessary and call the implementation with the correct cuBLAS op flag
    constexpr char op = (has_C_layout<A> ? (is_conj_array_expr<A> ? 'C' : 'T') : 'N');
    if constexpr (has_C_layout<A>) {
      return detail::getrs_batch_impl<run_on_device>(transpose(a), std::forward<B>(b), ipiv, op);
    } else {
      return detail::getrs_batch_impl<run_on_device>(a, std::forward<B>(b), ipiv, op);
    }
  }

  template <BlasArrayOrConj<3> A, BlasArrayFor<A, 3> B, PivotArrayFor<A, 2> IPIV>
    requires((has_F_layout<A> or has_C_layout<A>) and has_F_layout<B> and (not is_conj_array_expr<A> or has_C_layout<A>))
  int getrs(A const &a, B &&b, IPIV const &ipiv) { // NOLINT (temporary views are allowed here)
    return getrs_batch(a, std::forward<B>(b), ipiv);
  }
 
} // namespace nda::lapack