HIP/Morpheus_Workspace.hpp

 
#ifndef MORPHEUS_DENSEVECTOR_HIP_WORKSPACE_IMPL_HPP
#define MORPHEUS_DENSEVECTOR_HIP_WORKSPACE_IMPL_HPP
 
#include <Morpheus_Macros.hpp>
#if defined(MORPHEUS_ENABLE_HIP)
 
#include <impl/Morpheus_HIPUtils.hpp>
 
#ifdef MORPHEUS_ENABLE_TPL_HIPBLAS
#include <hipblas.h>
#endif  // MORPHEUS_ENABLE_TPL_HIPBLAS
 
namespace Morpheus {
namespace Impl {
 
class HIPWorkspace {
 public:
  HIPWorkspace() : _workspace(nullptr), _nbytes(0) {}
 
  ~HIPWorkspace() { _free(); }
 
  template <typename ValueType>
  void allocate(size_t N) {
    size_t bytes = N * sizeof(ValueType);
    if ((bytes > _nbytes) && (_nbytes > 0)) {
      _free();
    }
    _malloc(bytes);
  }
 
  template <typename ValueType>
  ValueType* data() {
    return reinterpret_cast<ValueType*>(_workspace);
  }
 
 private:
  void* _workspace;
  size_t _nbytes;
 
  void _malloc(size_t bytes) {
    if (_workspace == nullptr) {
      _nbytes = bytes;
      if (hipSuccess != hipMalloc(&_workspace, bytes)) {
        std::cout << "Malloc Failed on Device" << std::endl;
        exit(EXIT_FAILURE);
      }
    }
  }
 
  void _free() {
    if (_workspace != nullptr) {
      if (hipSuccess == hipFree(_workspace)) {
        _workspace = nullptr;
        _nbytes    = 0;
      }
    }
  }
};
 
#ifdef MORPHEUS_ENABLE_TPL_HIPBLAS
 
class HIPblasWorkspace : public HIPWorkspace {
 public:
  HIPblasWorkspace() : _handle(nullptr) {}
 
  ~HIPblasWorkspace() {
    if (_handle != nullptr) {
      if (hipblasDestroy(_handle) != HIPBLAS_STATUS_SUCCESS) {
        _handle = nullptr;
      }
    }
  }
 
  void init() {
    if (_handle == nullptr) {
      if (hipblasCreate(&_handle) != HIPBLAS_STATUS_SUCCESS) {
        printf("HIPBLAS initialization failed\n");
        exit(EXIT_FAILURE);
      }
      hipblasSetPointerMode(_handle, HIPBLAS_POINTER_MODE_DEVICE);
    }
  }
 
  hipblasHandle_t handle() { return _handle; }
 
 private:
  hipblasHandle_t _handle;
};
 
extern HIPblasWorkspace hipblasdotspace;
 
#endif  // MORPHEUS_ENABLE_TPL_HIPBLAS
 
extern HIPWorkspace hipdotspace;
 
}  // namespace Impl
}  // namespace Morpheus
 
#endif  // MORPHEUS_ENABLE_HIP
#endif  // MORPHEUS_DENSEVECTOR_HIP_WORKSPACE_IMPL_HPP