target.h

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright (c) NumeriCor GmbH
// SPDX-License-Identifier: LicenseRef-APL-1.1
 
#ifndef TARGETS_H
#define TARGETS_H
 
#include <stdexcept>
#include <iostream>
#include <string>
#include <cstring>
#ifdef HAS_ROCM_MODEL
#include <hip/hip_runtime.h>
#endif
#ifdef HAS_CUDA_MODEL
#include <cuda_runtime.h>
#endif
 
namespace limpet
{
 
enum Target {
  AUTO = -2,       // !< special value for chosing the target automatically
  UNKNOWN = -1,    
  CPU,        
  MLIR_CPU,   
  MLIR_ROCM,  
  MLIR_CUDA,  
  N_TARGETS,  
};
 
Target get_target_from_string(std::string const str);
 
std::string get_string_from_target(Target const target);
 
std::string get_target_list_string();
 
bool is_gpu(Target const target);
 
bool is_concrete(Target const target);
 
template <typename T>
struct TargetAllocator {
  typedef T value_type; 
 
  TargetAllocator(Target target, bool always_managed = false) : _target(target), _always_managed(always_managed) {
    if (!is_concrete(target)) {
      throw std::invalid_argument("attempting to construct a TargetAllocator with an invalid target");
    }
  }
 
  Target get_target() const {
    return this->_target;
  }
 
  void set_target(Target new_target) {
    if (!is_concrete(new_target)) {
      throw std::invalid_argument("attempting to set TargetAllocator to an invalid target");
    }
    else {
      this->_target = new_target;
    }
  }
 
  T* allocate(std::size_t n, bool do_zero = true) {
    using type = typename std::conditional<std::is_void<T>::value, char, T>::type;
    T* ptr = nullptr;
    std::size_t n_bytes = n * sizeof(type);
    switch (this->_target) {
      case Target::MLIR_CPU:
      case Target::CPU: {
        std::allocator<type> std_alloc;
        if(n)
        {
          ptr = std_alloc.allocate(n);
          // Explicitely cast the pointer to silence warning on setting a
          // dynamic class object to 0.
          // The memory allocated by this function should be initialized after,
          // for example by using a placement new
          if(do_zero)
            std::memset((void *) ptr, 0, n_bytes);
        }
 
 
        break;
                        }
      case Target::MLIR_ROCM:
#ifdef HAS_ROCM_MODEL
                        {
        T* hip_ptr;
        hipError_t error;
        // Force usage of HIP managed memory even if it's not supported.
        // This should still preserve the semantics of managed memory, but
        // is slower. It can be useful to allocate small shared data between CPU
        // and GPU.
        if (this->_always_managed) {
          error = hipMallocManaged(&hip_ptr, n_bytes);
        }
        else {
          int device;
          int managed_memory = 0;
          hipGetDevice(&device);
          hipDeviceGetAttribute(&managed_memory, hipDeviceAttributeManagedMemory, device);
          // Check for HIP managed memory and fallback to a normal hipMalloc if
          // it's not supported. It should still allow the CPU to write to GPU memory.
          if (managed_memory) {
            error = hipMallocManaged(&hip_ptr, n_bytes);
          }
          else {
            error = hipMalloc(&hip_ptr, n_bytes);
          }
        }
        if (error != hipSuccess) {
          throw std::runtime_error(hipGetErrorString(error));
        }
        error = hipMemset(hip_ptr, 0, n_bytes);
        if (error != hipSuccess) {
          throw std::runtime_error(hipGetErrorString(error));
        }
        ptr = hip_ptr;
                        }
#else
        throw std::invalid_argument("target MLIR_ROCM is unavailable");
#endif
        break;
      case Target::MLIR_CUDA:
#ifdef HAS_CUDA_MODEL
        {
          T* cuda_ptr;
          cudaError_t error;
          // Always use unified memory for CUDA
          error = cudaMallocManaged(&cuda_ptr, n_bytes);
          if (error != cudaSuccess) {
            throw std::runtime_error(cudaGetErrorString(error));
          }
          error = cudaMemset(cuda_ptr, 0, n_bytes);
          if (error != cudaSuccess) {
            throw std::runtime_error(cudaGetErrorString(error));
          }
          ptr = cuda_ptr;
        }
#else
        throw std::invalid_argument("target MLIR_CUDA is unavailable");
#endif
        break;
      default:
        throw std::invalid_argument("unknown allocation target");
    }
    return ptr;
  }
 
  // Size is not needed
  void deallocate(T* ptr, std::size_t n = 0) {
    using type = typename std::conditional<std::is_void<T>::value, char, T>::type;
    switch (this->_target) {
      case Target::MLIR_CPU:
      case Target::CPU: {
        std::allocator<type> std_alloc;
        std_alloc.deallocate((type*) ptr, n);
      }
        break;
      case Target::MLIR_ROCM:
#ifdef HAS_ROCM_MODEL
        {
          hipError_t error;
          error = hipFree(ptr);
          if (error != hipSuccess) {
            throw std::runtime_error(hipGetErrorString(error));
          }
        }
#else
        throw std::invalid_argument("target MLIR_ROCM is unavailable");
#endif
        break;
      case Target::MLIR_CUDA:
#ifdef HAS_CUDA_MODEL
        {
          cudaError_t error;
          error = cudaFree(ptr);
          if (error != cudaSuccess) {
            throw std::runtime_error(cudaGetErrorString(error));
          }
        }
#else
        throw std::invalid_argument("target MLIR_CUDA is unavailable");
#endif
        break;
      default:
        throw std::invalid_argument("unknown allocation target");
    }
  }
 
 private:
  Target _target; 
  bool _always_managed; 
};
 
template<typename T>
T* allocate_on_target(Target target, std::size_t n, bool always_managed = false, bool do_zero = true) {
  TargetAllocator<T> alloc(target, always_managed);
  return alloc.allocate(n, do_zero);
}
template<typename T>
void deallocate_on_target(Target target, T* ptr) {
  TargetAllocator<T> alloc(target);
  return alloc.deallocate(ptr);
}
 
}  // namespace limpet
 
#endif  // TARGETS_H