Tpetra_Details_StaticView.cpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#include "Tpetra_Details_StaticView.hpp"

namespace Tpetra {
namespace Details {
namespace Impl {

namespace {  // (anonymous)

// Motivating use cases for the initial size:
//
// 1. GMRES (need restart length (default 30) number of rows)
// 2. Single reduce CG (need 2 x 2)
constexpr size_t minimum_initial_size = sizeof(double) * 30 * 2;

// Intel 17 seems a bit buggy with respect to initialization of
// templated static classes, so let's make the compiler's job really
// easy by having nontemplated static raw pointers.

#ifdef KOKKOS_ENABLE_CUDA

void* cuda_memory_       = nullptr;
size_t cuda_memory_size_ = 0;

void finalize_cuda_memory() {
  if (cuda_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::CudaSpace>(cuda_memory_);
    cuda_memory_      = nullptr;
    cuda_memory_size_ = 0;
  }
}

void* cuda_uvm_memory_       = nullptr;
size_t cuda_uvm_memory_size_ = 0;

void finalize_cuda_uvm_memory() {
  if (cuda_uvm_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::CudaUVMSpace>(cuda_uvm_memory_);
    cuda_uvm_memory_      = nullptr;
    cuda_uvm_memory_size_ = 0;
  }
}

void* cuda_host_pinned_memory_       = nullptr;
size_t cuda_host_pinned_memory_size_ = 0;

void finalize_cuda_host_pinned_memory() {
  if (cuda_host_pinned_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::CudaHostPinnedSpace>(cuda_host_pinned_memory_);
    cuda_host_pinned_memory_      = nullptr;
    cuda_host_pinned_memory_size_ = 0;
  }
}
#endif  // KOKKOS_ENABLE_CUDA

#ifdef KOKKOS_ENABLE_HIP

void* hip_memory_       = nullptr;
size_t hip_memory_size_ = 0;

void finalize_hip_memory() {
  if (hip_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::HIPSpace>(hip_memory_);
    hip_memory_      = nullptr;
    hip_memory_size_ = 0;
  }
}

void* hip_host_pinned_memory_       = nullptr;
size_t hip_host_pinned_memory_size_ = 0;

void finalize_hip_host_pinned_memory() {
  if (hip_host_pinned_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::HIPHostPinnedSpace>(hip_host_pinned_memory_);
    hip_host_pinned_memory_      = nullptr;
    hip_host_pinned_memory_size_ = 0;
  }
}
#endif  // KOKKOS_ENABLE_HIP

#ifdef KOKKOS_ENABLE_SYCL

void* sycl_memory_       = nullptr;
size_t sycl_memory_size_ = 0;

void finalize_sycl_memory() {
  if (sycl_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::Experimental::SYCLDeviceUSMSpace>(sycl_memory_);
    sycl_memory_      = nullptr;
    sycl_memory_size_ = 0;
  }
}

void* sycl_shared_memory_       = nullptr;
size_t sycl_shared_memory_size_ = 0;

void finalize_sycl_shared_memory() {
  if (sycl_shared_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::Experimental::SYCLSharedUSMSpace>(sycl_shared_memory_);
    sycl_shared_memory_      = nullptr;
    sycl_shared_memory_size_ = 0;
  }
}
#endif  // KOKKOS_ENABLE_SYCL

void* host_memory_       = nullptr;
size_t host_memory_size_ = 0;

void finalize_host_memory() {
  if (host_memory_ != nullptr) {
    Kokkos::kokkos_free<Kokkos::HostSpace>(host_memory_);
    host_memory_      = nullptr;
    host_memory_size_ = 0;
  }
}

}  // namespace

#ifdef KOKKOS_ENABLE_CUDA

void* StaticKokkosAllocation<Kokkos::CudaSpace>::
    resize(Kokkos::CudaSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::CudaSpace;
  static bool created_finalize_hook = false;

  if (size > cuda_memory_size_) {
    if (cuda_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(cuda_memory_);
    }
    const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
    cuda_memory_          = Kokkos::kokkos_malloc<memory_space>(req_size);
    cuda_memory_size_     = size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_cuda_memory);
    created_finalize_hook = true;
  }

  return cuda_memory_;
}

void* StaticKokkosAllocation<Kokkos::CudaUVMSpace>::
    resize(Kokkos::CudaUVMSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::CudaUVMSpace;
  static bool created_finalize_hook = false;

  const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
  if (req_size > cuda_uvm_memory_size_) {
    if (cuda_uvm_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(cuda_uvm_memory_);
    }
    cuda_uvm_memory_      = Kokkos::kokkos_malloc<memory_space>(req_size);
    cuda_uvm_memory_size_ = req_size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_cuda_uvm_memory);
    created_finalize_hook = true;
  }

  return cuda_uvm_memory_;
}

void* StaticKokkosAllocation<Kokkos::CudaHostPinnedSpace>::
    resize(Kokkos::CudaHostPinnedSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::CudaHostPinnedSpace;
  static bool created_finalize_hook = false;

  const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
  if (req_size > cuda_host_pinned_memory_size_) {
    if (cuda_host_pinned_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(cuda_host_pinned_memory_);
    }
    cuda_host_pinned_memory_      = Kokkos::kokkos_malloc<memory_space>(req_size);
    cuda_host_pinned_memory_size_ = req_size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_cuda_host_pinned_memory);
    created_finalize_hook = true;
  }

  return cuda_host_pinned_memory_;
}

#endif  // KOKKOS_ENABLE_CUDA

#ifdef KOKKOS_ENABLE_HIP

void* StaticKokkosAllocation<Kokkos::HIPSpace>::
    resize(Kokkos::HIPSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::HIPSpace;
  static bool created_finalize_hook = false;

  if (size > hip_memory_size_) {
    if (hip_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(hip_memory_);
    }
    const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
    hip_memory_           = Kokkos::kokkos_malloc<memory_space>(req_size);
    hip_memory_size_      = size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_hip_memory);
    created_finalize_hook = true;
  }

  return hip_memory_;
}

void* StaticKokkosAllocation<Kokkos::HIPHostPinnedSpace>::
    resize(Kokkos::HIPHostPinnedSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::HIPHostPinnedSpace;
  static bool created_finalize_hook = false;

  const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
  if (req_size > hip_host_pinned_memory_size_) {
    if (hip_host_pinned_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(hip_host_pinned_memory_);
    }
    hip_host_pinned_memory_      = Kokkos::kokkos_malloc<memory_space>(req_size);
    hip_host_pinned_memory_size_ = req_size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_hip_host_pinned_memory);
    created_finalize_hook = true;
  }

  return hip_host_pinned_memory_;
}

#endif  // KOKKOS_ENABLE_HIP

#ifdef KOKKOS_ENABLE_SYCL

template <>
void* StaticKokkosAllocation<Kokkos::Experimental::SYCLDeviceUSMSpace>::
    resize(Kokkos::Experimental::SYCLDeviceUSMSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::Experimental::SYCLDeviceUSMSpace;
  static bool created_finalize_hook = false;

  if (size > sycl_memory_size_) {
    if (sycl_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(sycl_memory_);
    }
    const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
    sycl_memory_          = Kokkos::kokkos_malloc<memory_space>(req_size);
    sycl_memory_size_     = size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_sycl_memory);
    created_finalize_hook = true;
  }

  return sycl_memory_;
}

template <>
void* StaticKokkosAllocation<Kokkos::Experimental::SYCLSharedUSMSpace>::
    resize(Kokkos::Experimental::SYCLSharedUSMSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::Experimental::SYCLSharedUSMSpace;
  static bool created_finalize_hook = false;

  const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
  if (req_size > sycl_shared_memory_size_) {
    if (sycl_shared_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(sycl_shared_memory_);
    }
    sycl_shared_memory_      = Kokkos::kokkos_malloc<memory_space>(req_size);
    sycl_shared_memory_size_ = req_size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_sycl_shared_memory);
    created_finalize_hook = true;
  }

  return sycl_shared_memory_;
}

#endif  // KOKKOS_ENABLE_SYCL

void* StaticKokkosAllocation<Kokkos::HostSpace>::
    resize(Kokkos::HostSpace /* space */,
           const size_t size) {
  using memory_space                = Kokkos::HostSpace;
  static bool created_finalize_hook = false;

  const size_t req_size = size > minimum_initial_size ? size : minimum_initial_size;
  if (req_size > host_memory_size_) {
    if (host_memory_ != nullptr) {
      Kokkos::kokkos_free<memory_space>(host_memory_);
    }
    host_memory_      = Kokkos::kokkos_malloc<memory_space>(req_size);
    host_memory_size_ = req_size;
  }
  if (!created_finalize_hook) {
    Kokkos::push_finalize_hook(finalize_host_memory);
    created_finalize_hook = true;
  }

  return host_memory_;
}

}  // namespace Impl
}  // namespace Details
}  // namespace Tpetra