getri_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 "./getri.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 <type_traits>
#include <utility>
 
namespace nda::lapack {
 
  namespace detail {
 
    // Implementation of the batched getri routine.
    template <bool run_on_device>
    auto getri_batch_impl(auto &&a, auto const &ipiv, [[maybe_unused]] auto &&work) {
      // check the dimensions of the input/output arrays/views
      auto const [m, n, n_b] = a.shape();
      EXPECTS(m == n);
      EXPECTS(ipiv.extent(0) == n);
      EXPECTS(ipiv.extent(1) == n_b);
 
      // arrays/views must be LAPACK compatible
      EXPECTS(a.indexmap().min_stride() == 1);
      EXPECTS(ipiv.indexmap().min_stride() == 1);
 
      // perform actual library call(s)
      auto info = array<int, 1>(n_b, 0);
      if constexpr (run_on_device) {
        using arr_t = std::remove_cvref_t<decltype(a)>;
 
        // resize/check work buffer
        resize_or_check_if_view(work, {a.size()});
        EXPECTS(work.indexmap().min_stride() == 1);
 
        // output buffer for inverted matrices
        auto c = cuarray_view<get_value_t<arr_t>, 3, F_layout>(a.shape(), work.data());
 
        auto a_ptrs = to_device(batch_ptrs(a));
        auto c_ptrs = to_device(batch_ptrs(c));
        auto info_d = vector<int, heap<mem::get_addr_space<arr_t>>>(n_b, 0);
        blas::device::getri_batch(n, a_ptrs.data(), get_ld(a(range::all, range::all, 0)), ipiv.data(), c_ptrs.data(),
                                  get_ld(c(range::all, range::all, 0)), info_d.data(), n_b);
        info = info_d;
 
        // copy result back to a
        a = c;
      } else {
        // for host, fall back to looping over batches
        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)     = getri(a_i, ipiv_i, work);
        }
      }
      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 getri_batch(A &&a, IPIV const &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 getri_batch(a, transpose(ipiv), work);
 
    // transpose A array/view if necessary and call the implementation
    if constexpr (has_C_layout<A>) {
      return detail::getri_batch_impl<run_on_device>(transpose(a), ipiv, std::forward<W>(work));
    } else {
      return detail::getri_batch_impl<run_on_device>(std::forward<A>(a), 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 getri(A &&a, IPIV const &ipiv, W &&work = vector_value_t<A>{}) { // NOLINT (temporary views are allowed here)
    return getri_batch(std::forward<A>(a), ipiv, std::forward<W>(work));
  }
 
} // namespace nda::lapack