getrf_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 "./getrf.hpp"
#include "./interface/cxx_interface.hpp"
#include "../basic_array.hpp"
#include "../basic_functions.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 <type_traits>
#include <utility>
 
namespace nda::lapack {
 
  namespace detail {
 
    // Implementation of the batched getrf routine.
    template <bool run_on_device>
    auto getrf_batch_impl(auto &&a, auto &&ipiv, [[maybe_unused]] auto &&work) {
      // check the dimensions of the input/output arrays/views and resize if necessary
      auto const m   = a.extent(0);
      auto const n   = a.extent(1);
      auto const n_b = a.extent(2);
      resize_or_check_if_view(ipiv, {std::min(m, n), n_b});
 
      // arrays/views must be LAPACK compatible
      EXPECTS(a.indexmap().min_stride() == 1);
      EXPECTS(ipiv.indexmap().min_stride() == 1);
 
#if defined(__has_feature)
#if __has_feature(memory_sanitizer)
      ipiv = 0;
#endif
#endif
 
      // loop over batches and call getrf for each matrix
      auto loop_getrf = [n_b, &a, &ipiv, &work](auto &info) {
        for (int i = 0; i < n_b; ++i) {
          auto a_i    = a(range::all, range::all, i);
          auto ipiv_i = ipiv(range::all, i);
          info(i)     = getrf(a_i, ipiv_i, work);
        }
      };
 
      // perform actual library call(s)
      auto info = array<int, 1>(n_b, 0);
      if constexpr (run_on_device) {
        if (m == n) {
          // for square matrices on the device use cuBLAS
          using arr_t = std::remove_cvref_t<decltype(a)>;
          auto ptr_d  = to_device(batch_ptrs(a));
          auto info_d = vector<int, heap<mem::get_addr_space<arr_t>>>(n_b, 0);
          blas::device::getrf_batch(n, ptr_d.data(), get_ld(a(range::all, range::all, 0)), ipiv.data(), info_d.data(), n_b);
          info = info_d;
        } else {
          // for rectangular matrices on the device, fall back to looping over batches
          loop_getrf(info);
        }
      } else {
        // for host, fall back to looping over batches
        loop_getrf(info);
      }
      return info;
    }
 
  } // namespace detail

  template <BlasArray<3> A, PivotArrayFor<A, 2> IPIV, BlasArrayFor<A, 1> W = vector_value_t<A>>
    requires(has_F_layout<A> or has_C_layout<A>)
  auto getrf_batch(A &&a, IPIV &&ipiv, W &&work = vector_value_t<A>{}) { // NOLINT (temporary views are allowed here)
    constexpr bool run_on_device = mem::have_device_compatible_addr_space<A, IPIV, W>;
 
    // transpose ipiv array/view if necessary
    if constexpr (has_C_layout<IPIV>) return getrf_batch(a, transpose(ipiv), work);
 
    // transpose A array/view if necessary and call the implementation
    if constexpr (has_C_layout<A>) {
      return detail::getrf_batch_impl<run_on_device>(transpose(a), std::forward<IPIV>(ipiv), std::forward<W>(work));
    } else {
      return detail::getrf_batch_impl<run_on_device>(std::forward<A>(a), std::forward<IPIV>(ipiv), std::forward<W>(work));
    }
  }

  template <BlasArray<3> A, PivotArrayFor<A, 2> IPIV, BlasArrayFor<A, 1> W = vector_value_t<A>>
    requires(has_F_layout<A> or has_C_layout<A>)
  auto getrf(A &&a, IPIV &&ipiv, W &&work = vector_value_t<A>{}) { // NOLINT (temporary views are allowed here)
    return getrf_batch(std::forward<A>(a), std::forward<IPIV>(ipiv), std::forward<W>(work));
  }
 
} // namespace nda::lapack