tsgAcceleratedDataStructures.hpp

/*
 * Copyright (c) 2017, Miroslav Stoyanov
 *
 * This file is part of
 * Toolkit for Adaptive Stochastic Modeling And Non-Intrusive ApproximatioN: TASMANIAN
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions
 *    and the following disclaimer in the documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse
 *    or promote products derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
 * OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * UT-BATTELLE, LLC AND THE UNITED STATES GOVERNMENT MAKE NO REPRESENTATIONS AND DISCLAIM ALL WARRANTIES, BOTH EXPRESSED AND IMPLIED.
 * THERE ARE NO EXPRESS OR IMPLIED WARRANTIES OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE, OR THAT THE USE OF THE SOFTWARE WILL NOT INFRINGE ANY PATENT,
 * COPYRIGHT, TRADEMARK, OR OTHER PROPRIETARY RIGHTS, OR THAT THE SOFTWARE WILL ACCOMPLISH THE INTENDED RESULTS OR THAT THE SOFTWARE OR ITS USE WILL NOT RESULT IN INJURY OR DAMAGE.
 * THE USER ASSUMES RESPONSIBILITY FOR ALL LIABILITIES, PENALTIES, FINES, CLAIMS, CAUSES OF ACTION, AND COSTS AND EXPENSES, CAUSED BY, RESULTING FROM OR ARISING OUT OF,
 * IN WHOLE OR IN PART THE USE, STORAGE OR DISPOSAL OF THE SOFTWARE.
 */
 
#ifndef __TASMANIAN_SPARSE_GRID_ACCELERATED_DATA_STRUCTURES_HPP
#define __TASMANIAN_SPARSE_GRID_ACCELERATED_DATA_STRUCTURES_HPP
 
#include "tsgAcceleratedHandles.hpp"


 
namespace TasGrid{
 
struct AccelerationContext; // forward declaration, CUDA and HIP GpuVector do not use the context, but DPC++ needs it

template<typename T>
class GpuVector{
public:
    GpuVector(GpuVector<T> const &) = delete;
    GpuVector<T>& operator =(GpuVector<T> const &) = delete;

    GpuVector(GpuVector<T> &&other) : num_entries(std::exchange(other.num_entries, 0)), gpu_data(std::exchange(other.gpu_data, nullptr))
    #ifdef Tasmanian_ENABLE_DPCPP
        , sycl_queue(other.sycl_queue)
    #endif
    {}
    GpuVector<T>& operator =(GpuVector<T> &&other){
        GpuVector<T> temp(std::move(other));
        std::swap(num_entries, temp.num_entries);
        std::swap(gpu_data, temp.gpu_data);
        #ifdef Tasmanian_ENABLE_DPCPP
        std::swap(sycl_queue, temp.sycl_queue);
        #endif
        return *this;
    }

    GpuVector() = default;
    GpuVector(AccelerationContext const *acc, size_t count){ resize(acc, count); }

    GpuVector(AccelerationContext const *acc, int dim1, int dim2){ resize(acc, Utils::size_mult(dim1, dim2)); }
    GpuVector(AccelerationContext const *acc, const std::vector<T> &cpu_data){ load(acc, cpu_data); }
    GpuVector(AccelerationContext const *acc, int dim1, int dim2, T const *cpu_data){ load(acc, Utils::size_mult(dim1, dim2), cpu_data); }
    template<typename IteratorLike>
    GpuVector(AccelerationContext const *acc, IteratorLike ibegin, IteratorLike iend) : GpuVector(){
        load(acc, std::distance(ibegin, iend), std::addressof(*ibegin));
    }
    ~GpuVector(){ clear(); }

    size_t size() const{ return num_entries; }
    T* data(){ return gpu_data; }
    const T* data() const{ return gpu_data; }

    void resize(AccelerationContext const *acc, size_t count);
    void clear();
    bool empty() const{ return (num_entries == 0); }

    template<typename U>
    void load(AccelerationContext const *acc, const std::vector<U> &cpu_data){
        load(acc, cpu_data.size(), cpu_data.data());
    }

    template<typename U>
    void load(AccelerationContext const *acc, size_t count, const U* cpu_data) {
        if constexpr (std::is_same_v<U, T>) {
            load_internal(acc, count, cpu_data);
        } else {
            // need to convert the array
            std::vector<T> converted(count);
            std::transform(cpu_data, cpu_data + count, converted.begin(), [](U const &x)->T{ return static_cast<T>(x); });
            load_internal(acc, count, converted.data());
        }
    }

    void unload(AccelerationContext const *acc, std::vector<T> &cpu_data) const{
        cpu_data.resize(num_entries);
        unload(acc, cpu_data.data());
    }
    std::vector<T> unload(AccelerationContext const *acc) const{
        std::vector<T> y;
        unload(acc, y);
        return y;
    }
    void unload(AccelerationContext const *acc, size_t num, T* cpu_data) const;
    void unload(AccelerationContext const *acc, T* cpu_data) const{ unload(acc, num_entries, cpu_data); }

    T* eject(){
        T* external = gpu_data;
        gpu_data = nullptr;
        num_entries = 0;
        return external;
    }

    using value_type = T;
 
private:
    void load_internal(AccelerationContext const *acc, size_t count, const T* cpu_data);
 
    size_t num_entries = 0; // keep track of the size, update on every call that changes the gpu_data
    T *gpu_data = nullptr; // the GPU array
    #ifdef Tasmanian_ENABLE_DPCPP
    void* sycl_queue = nullptr;
    #endif
};

struct GpuEngine{
    GpuEngine() = default;
    GpuEngine(GpuEngine const &) = delete;
    GpuEngine &operator = (GpuEngine const &) = delete;
    GpuEngine(GpuEngine &&) = default;
    GpuEngine &operator = (GpuEngine &&) = default;
 
    #ifdef Tasmanian_ENABLE_CUDA
    void setCuBlasHandle(void *handle);
    void setCuSparseHandle(void *handle);
    void setCuSolverDnHandle(void *handle);

    std::unique_ptr<int, HandleDeleter<AccHandle::Cublas>> cublas_handle;
    std::unique_ptr<int, HandleDeleter<AccHandle::Cusparse>> cusparse_handle;
    std::unique_ptr<int, HandleDeleter<AccHandle::Cusolver>> cusolver_handle;
    #endif
 
    #ifdef Tasmanian_ENABLE_HIP
    void setRocBlasHandle(void *handle);
    void setRocSparseHandle(void *handle);

    std::unique_ptr<int, HandleDeleter<AccHandle::Rocblas>> rblas_handle;
    std::unique_ptr<int, HandleDeleter<AccHandle::Rocsparse>> rsparse_handle;
    #endif
 
    #ifdef Tasmanian_ENABLE_DPCPP
    void setSyclQueue(void *queue);
    std::unique_ptr<int, HandleDeleter<AccHandle::Syclqueue>> internal_queue;
    #endif

    bool called_magma_init = false;
};


class AccelerationDomainTransform{
public:
    AccelerationDomainTransform() = default;
    AccelerationDomainTransform(AccelerationContext const *acc, std::vector<double> const &transform_a, std::vector<double> const &transform_b){
        set(acc, transform_a, transform_b);
    }

    template<typename T>
    void getCanonicalPoints(bool use01, T const gpu_transformed_x[], int num_x, GpuVector<T> &gpu_canonical_x);
    operator bool () const { return (acceleration != nullptr); }
    void set(AccelerationContext const *, std::vector<double> const &transform_a, std::vector<double> const &transform_b);
    void clear() {
        gpu_trans_a.clear();
        gpu_trans_b.clear();
        num_dimensions = 0;
        padded_size = 0;
        acceleration = nullptr;
    }
 
private:
    // these actually store the rate and shift and not the hard upper/lower limits
    GpuVector<double> gpu_trans_a, gpu_trans_b;
    int num_dimensions = 0, padded_size = 0;
    AccelerationContext const *acceleration = nullptr;
};

namespace TasGpu{

    template<typename T>
    void dtrans2can(AccelerationContext const *acc, bool use01, int dims, int num_x, int pad_size,
                    const double *gpu_trans_a, const double *gpu_trans_b,
                    const T *gpu_x_transformed, T *gpu_x_canonical);


    template<typename T>
    void devalpwpoly(AccelerationContext const *acc, int order, TypeOneDRule rule, int num_dimensions, int num_x, int num_basis, const T *gpu_x, const T *gpu_nodes, const T *gpu_support, T *gpu_y);


    template<typename T>
    void devalpwpoly_sparse(AccelerationContext const *acc, int order, TypeOneDRule rule, int dims, int num_x, const T *gpu_x,
                            const GpuVector<T> &gpu_nodes, const GpuVector<T> &gpu_support,
                            const GpuVector<int> &gpu_hpntr, const GpuVector<int> &gpu_hindx, const GpuVector<int> &gpu_hroots,
                            GpuVector<int> &gpu_spntr, GpuVector<int> &gpu_sindx, GpuVector<T> &gpu_svals);


    template<typename T>
    void devalseq(AccelerationContext const *acc, int dims, int num_x, const std::vector<int> &max_levels, const T *gpu_x,
                  const GpuVector<int> &num_nodes,
                  const GpuVector<int> &points, const GpuVector<T> &nodes, const GpuVector<T> &coeffs, T *gpu_result);


    template<typename T>
    void devalfor(AccelerationContext const *acc, int dims, int num_x, const std::vector<int> &max_levels, const T *gpu_x, const GpuVector<int> &num_nodes, const GpuVector<int> &points, T *gpu_wreal, typename GpuVector<T>::value_type *gpu_wimag);

    template<typename T>
    void devalglo(AccelerationContext const *acc, bool is_nested, bool is_clenshawcurtis0, int dims, int num_x, int num_p, int num_basis,
                  T const *gpu_x, GpuVector<T> const &nodes, GpuVector<T> const &coeff, GpuVector<T> const &tensor_weights,
                  GpuVector<int> const &nodes_per_level, GpuVector<int> const &offset_per_level, GpuVector<int> const &map_dimension, GpuVector<int> const &map_level,
                  GpuVector<int> const &active_tensors, GpuVector<int> const &active_num_points, GpuVector<int> const &dim_offsets,
                  GpuVector<int> const &map_tensor, GpuVector<int> const &map_index, GpuVector<int> const &map_reference, T *gpu_result);
 

    void fillDataGPU(AccelerationContext const *acc, double value, long long N, long long stride, double data[]);

    template<typename T> void load_n_internal(AccelerationContext const *acc, T const *cpu_data, size_t num_entries, T *gpu_data);

    template<typename T, typename U>
    void load_n(AccelerationContext const *acc, U const *cpu_data, size_t num_entries, T *gpu_data){
        if constexpr (std::is_same_v<U, T>) {
            load_n_internal(acc, cpu_data, num_entries, gpu_data);
        } else {
            std::vector<T> converted(num_entries);
            std::transform(cpu_data, cpu_data + num_entries, converted.begin(), [](U const &x)->T{ return static_cast<T>(x); });
            load_n_internal(acc, converted.data(), num_entries, gpu_data);
        }
    }
 
    // #define __TASMANIAN_COMPILE_FALLBACK_CUDA_KERNELS__ // uncomment to compile a bunch of custom CUDA kernels that provide some functionality similar to cuBlas
    #ifdef __TASMANIAN_COMPILE_FALLBACK_CUDA_KERNELS__
    // CUDA kernels that provide essentially the same functionality as cuBlas and MAGMA, but nowhere near as optimal
    // those functions should not be used in a Release or production builds
    // the kernels are useful because they are simple and do not depend on potentially poorly documented 3d party library
    // since the kernels are useful for testing and some debugging, the code should not be deleted (for now), but also don't waste time compiling in most cases
 
    void cudaDgemm(int M, int N, int K, const double *gpu_a, const double *gpu_b, double *gpu_c);
    // lazy cuda dgemm, nowhere near as powerful as cuBlas, but does not depend on cuBlas
    // gpu_a is M by K, gpu_b is K by N, gpu_c is M by N, all in column-major format
    // on exit gpu_c = gpu_a * gpu_b
 
    void cudaSparseMatmul(int M, int N, int num_nz, const int* gpu_spntr, const int* gpu_sindx, const double* gpu_svals, const double *gpu_B, double *gpu_C);
    // lazy cuda sparse dgemm, less efficient (especially for large N), but more memory conservative then cusparse as there is no need for a transpose
    // C is M x N, B is K x N (K is max(gpu_sindx)), both are given in row-major format, num_nz/spntr/sindx/svals describe row compressed A which is M by K
    // on exit C = A * B
 
    void cudaSparseVecDenseMat(int M, int N, int num_nz, const double *A, const int *indx, const double *vals, double *C);
    // dense matrix A (column major) times a sparse vector defiend by num_nz, indx, and vals
    // A is M by N, C is M by 1,
    // on exit C = A * (indx, vals)
 
    void convert_sparse_to_dense(int num_rows, int num_columns, const int *gpu_pntr, const int *gpu_indx, const double *gpu_vals, double *gpu_destination);
    // converts a sparse matrix to a dense representation (all data sits on the gpu and is pre-allocated)
    #endif
}

namespace AccelerationMeta{
    TypeAcceleration getIOAccelerationString(const char * name);
    std::map<std::string, TypeAcceleration> getStringToAccelerationMap();
    const char* getIOAccelerationString(TypeAcceleration accel);
    int getIOAccelerationInt(TypeAcceleration accel);
    TypeAcceleration getIOIntAcceleration(int accel);
    bool isAccTypeGPU(TypeAcceleration accel);

    inline bool isAvailable(TypeAcceleration accel){
        switch(accel){
            #ifdef Tasmanian_ENABLE_MAGMA
            case accel_gpu_magma: return true;
            #endif
            #ifdef Tasmanian_ENABLE_CUDA
            case accel_gpu_cuda: return true;
            case accel_gpu_cublas: return true;
            #endif
            #ifdef Tasmanian_ENABLE_HIP
            case accel_gpu_hip: return true;
            case accel_gpu_rocblas: return true;
            #endif
            #ifdef Tasmanian_ENABLE_DPCPP
            case accel_gpu_cuda: return true;
            case accel_gpu_cublas: return true;
            #endif
            #ifdef Tasmanian_ENABLE_BLAS
            case accel_cpu_blas: return true;
            #endif
            case accel_none:
                return true;
            default:
                return false;
        }
    }


    TypeAcceleration getAvailableFallback(TypeAcceleration accel);

    int getNumGpuDevices();


    void setDefaultGpuDevice(int deviceID);


    unsigned long long getTotalGPUMemory(int deviceID);


    std::string getGpuDeviceName(int deviceID);

    template<typename T> void recvGpuArray(AccelerationContext const*, size_t num_entries, const T *gpu_data, std::vector<T> &cpu_data);

    template<typename T> void delGpuArray(AccelerationContext const*, T *x);

    void *createCublasHandle();
    void deleteCublasHandle(void *);
}

struct AccelerationContext{
    AccelerationContext(AccelerationContext const &) = delete;
    AccelerationContext &operator = (AccelerationContext const &) = delete;
    AccelerationContext(AccelerationContext &&) = default;
    AccelerationContext &operator = (AccelerationContext &&) = default;

    enum AlgorithmPreference{
        algorithm_dense,
        algorithm_sparse,
        algorithm_autoselect
    };

    enum ChangeType{
        change_none,
        change_gpu_device,
        change_gpu_enabled,
        change_cpu_blas,
        change_sparse_dense
    };

    TypeAcceleration mode;
    AlgorithmPreference algorithm_select;
    int device;

    mutable std::optional<GpuEngine> engine;

    inline static constexpr TypeAcceleration getDefaultAccMode() {
        #ifdef Tasmanian_ENABLE_BLAS
        return accel_cpu_blas;
        #else
        return accel_none;
        #endif
    }
    inline static constexpr int getDefaultAccDevice() {
        #ifdef Tasmanian_ENABLE_DPCPP
        return -1;
        #else
        return 0;
        #endif
    }

    AccelerationContext() : mode(getDefaultAccMode()), algorithm_select(algorithm_autoselect), device(getDefaultAccDevice()){}

    ChangeType favorSparse(bool favor){
        AlgorithmPreference new_preference = [as=algorithm_select, favor]()->AlgorithmPreference{
            if (favor){
                return (as == algorithm_dense) ? algorithm_autoselect : algorithm_sparse;
            }else{
                return (as == algorithm_sparse) ? algorithm_autoselect : algorithm_dense;
            }
        }();
        if (new_preference != algorithm_select){
            algorithm_select = new_preference;
            return change_sparse_dense;
        }else{
            return change_none;
        }
    }

    bool blasCompatible() const{
        #ifdef Tasmanian_ENABLE_BLAS
        return (mode != accel_none);
        #else
        return false;
        #endif
    }

    bool useKernels() const{
        #if defined(Tasmanian_ENABLE_CUDA) || defined(Tasmanian_ENABLE_HIP)
        return ((mode == accel_gpu_cuda) or (mode == accel_gpu_magma));
        #else
        return false;
        #endif
    }

    ChangeType testEnable(TypeAcceleration acc, int new_gpu_id) const{
        TypeAcceleration effective_acc = AccelerationMeta::getAvailableFallback(acc);
        #ifdef Tasmanian_ENABLE_DPCPP
        if (AccelerationMeta::isAccTypeGPU(effective_acc) and ((new_gpu_id < -1 or new_gpu_id >= AccelerationMeta::getNumGpuDevices())))
            throw std::runtime_error("Invalid GPU device ID, see ./tasgrid -v for list of detected devices.");
        #else
        if (AccelerationMeta::isAccTypeGPU(effective_acc) and ((new_gpu_id < 0 or new_gpu_id >= AccelerationMeta::getNumGpuDevices())))
            throw std::runtime_error("Invalid GPU device ID, see ./tasgrid -v for list of detected devices.");
        #endif
        ChangeType mode_change = (effective_acc == mode) ? change_none : change_cpu_blas;
        ChangeType device_change = (device == new_gpu_id) ? change_none : change_gpu_device;
 
        if (on_gpu()){
            return (AccelerationMeta::isAccTypeGPU(effective_acc)) ? device_change : change_gpu_device;
        }else{
            return (AccelerationMeta::isAccTypeGPU(effective_acc)) ? change_gpu_enabled : mode_change;
        }
    }

    void enable(TypeAcceleration acc, int new_gpu_id){
        // get the effective new acceleration mode (use the fallback if acc is not enabled)
        TypeAcceleration effective_acc = AccelerationMeta::getAvailableFallback(acc);
        // if switching to a GPU mode, check if the device id is valid
        #ifdef Tasmanian_ENABLE_DPCPP
        if (AccelerationMeta::isAccTypeGPU(effective_acc) and ((new_gpu_id < -1 or new_gpu_id >= AccelerationMeta::getNumGpuDevices())))
            throw std::runtime_error("Invalid GPU device ID, see ./tasgrid -v for list of detected devices.");
        #else
        if (AccelerationMeta::isAccTypeGPU(effective_acc) and ((new_gpu_id < 0 or new_gpu_id >= AccelerationMeta::getNumGpuDevices())))
            throw std::runtime_error("Invalid GPU device ID, see ./tasgrid -v for list of detected devices.");
        #endif
        if (AccelerationMeta::isAccTypeGPU(effective_acc)){
            // if the new mode is GPU-based, make an engine or reset the engine if the device has changed
            // if the engine exists and the device is not changed, then keep the existing engine
            if (!engine or new_gpu_id != device)
                engine.emplace();
        }else{
            engine.reset();
        }
 
        // assign the new values for the mode and device
        mode = effective_acc;
        device = new_gpu_id;
    }
    void setDevice() const{ AccelerationMeta::setDefaultGpuDevice(device); }
    operator GpuEngine* () const{ return std::addressof(engine.value()); }
    bool on_gpu() const{ return !!engine; }
};
 
#ifdef Tasmanian_ENABLE_DPCPP
struct InternalSyclQueue{
    InternalSyclQueue() = default;
    void init_testing(int gpuid);
    operator std::unique_ptr<int, HandleDeleter<AccHandle::Syclqueue>> (){
        return std::unique_ptr<int, HandleDeleter<AccHandle::Syclqueue>>(test_queue.get(),
                                                                         HandleDeleter<AccHandle::Syclqueue>(false));
    }
    bool use_testing = false;
    std::unique_ptr<int, HandleDeleter<AccHandle::Syclqueue>> test_queue;
};
extern InternalSyclQueue test_queue;
#endif
 
}
 
#endif // __TASMANIAN_SPARSE_GRID_ACCELERATED_DATA_STRUCTURES_HPP