scratch.hpp

Go to the documentation of this file.

#ifndef SCRATCH_HPP
#define SCRATCH_HPP
 
#include <mutex>
#include <iostream>
#include "space_settings.hpp"
 
class ScratchLock;
 
class Scratch {
public:
    Scratch(size_t n_double_host, size_t n_double_device);
    Scratch() : is_initialized(false) {}
 
    KOKKOS_INLINE_FUNCTION void pop_host(size_t n) {
        //std::lock_guard<std::mutex> lock(mutex_);
        host_offset_ -= n;
    }
 
    KOKKOS_INLINE_FUNCTION void pop_device(size_t n) {
        //std::lock_guard<std::mutex> lock(mutex_);
        device_offset_ -= n;
    }
 
    void report_peak_usage() const;
 
private:
    friend class ScratchLock;
 
    bool is_initialized;
    View<double*, CLayout, HostType> host_pool_;
    View<double*, CLayout, DeviceType> device_pool_;
 
    size_t host_offset_;
    size_t device_offset_;
 
    size_t host_unlimited_offset_;
    size_t device_unlimited_offset_;
 
    size_t host_peak_;
    size_t device_peak_;
 
    //std::mutex mutex_;
};
 
extern std::unique_ptr<Scratch> GLOBAL_SCRATCH;
 
class ScratchLock {
public:
    enum Layout{
        KokkosDefaultLayout = true
    };
 
    template<typename Device> View<double*, CLayout, Device> get_view(const std::string& label, int dim1);
    template<typename Device> View<double**, CLayout, Device> get_view(const std::string& label, int dim1, int dim2);
    template<typename Device> View<double***, CLayout, Device> get_view(const std::string& label, int dim1, int dim2, int dim3);
    template<typename Device> View<double****, CLayout, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4);
    template<typename Device> View<double*****, CLayout, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4, int dim5);
    template<typename Device> View<double******, CLayout, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4, int dim5, int dim6);
 
    template<typename Device> View<double*, Device> get_view(const std::string& label, int dim1, Layout kdl_true);
    template<typename Device> View<double**, Device> get_view(const std::string& label, int dim1, int dim2, Layout kdl_true);
    template<typename Device> View<double***, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, Layout kdl_true);
    template<typename Device> View<double****, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4, Layout kdl_true);
    template<typename Device> View<double*****, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4, int dim5, Layout kdl_true);
    template<typename Device> View<double******, Device> get_view(const std::string& label, int dim1, int dim2, int dim3, int dim4, int dim5, int dim6, Layout kdl_true);
 
    // Constructor
    ScratchLock(Scratch* scratch)
      : scratch_(scratch), is_used_(false) {}
 
    // Default Constructor
    KOKKOS_INLINE_FUNCTION ScratchLock()
      : scratch_(nullptr), is_used_(false) {}
 
    // Destructor
    KOKKOS_INLINE_FUNCTION ~ScratchLock() {
        if(is_used_){
            if (used_scratch_) {
                if (is_device_){
                    //printf("Releasing %s scratch range: %lu-%lu\n", is_device_ ? "device" : "host", allocation_offset_, scratch_->device_offset_);
                    scratch_->pop_device(allocation_size_);
                    assert(scratch_->device_offset_ == allocation_offset_ && "FILO order violated: device memory was not released in reverse allocation order.");
                }else{
                    //printf("Releasing %s scratch range: %lu-%lu\n", is_device_ ? "device" : "host", allocation_offset_, scratch_->host_offset_);
                    scratch_->pop_host(allocation_size_);
                    assert(scratch_->host_offset_ == allocation_offset_ && "FILO order violated: host memory was not released in reverse allocation order.");
                }
            }
            if (is_device_){
                //printf("Releasing %s tracked range: %lu-%lu\n", is_device_ ? "device" : "host", scratch_->device_unlimited_offset_-allocation_size_, scratch_->device_unlimited_offset_);
                scratch_->device_unlimited_offset_ -= allocation_size_;
            }else{
                //printf("Releasing %s tracked range: %lu-%lu\n", is_device_ ? "device" : "host", scratch_->host_unlimited_offset_-allocation_size_, scratch_->host_unlimited_offset_);
                scratch_->host_unlimited_offset_ -= allocation_size_;
            }
        }
    }
 
    // Copy constructor: don't transfer ownership. Only original has ownership!
    KOKKOS_INLINE_FUNCTION ScratchLock(const ScratchLock& other)
      : scratch_(other.scratch_),
        allocation_offset_(other.allocation_offset_),
        allocation_size_(other.allocation_size_),
        is_device_(other.is_device_),
        used_scratch_(other.used_scratch_),
        is_used_(false) {}
 
    // Copy assignment: don't transfer ownership. Only original has ownership!
    ScratchLock& operator=(const ScratchLock& other){
        scratch_ = other.scratch_;
        allocation_offset_ = other.allocation_offset_;
        allocation_size_ = other.allocation_size_;
        is_device_ = other.is_device_;
        used_scratch_ = other.used_scratch_;
        is_used_ = false;
        return *this;
    }
 
 
private:
    Scratch* scratch_;
    size_t allocation_size_;
    size_t allocation_offset_;
    bool is_device_;
    bool used_scratch_;
    bool is_used_;
 
    template<typename Device> double* get_ptr(size_t view_size);
};
 
#endif