algorithms.h

Go to the documentation of this file.

#pragma once
 
#include "backend/api.h"
#include "backend/types.h"
#include "cuda_check.h"
#include "mem/mempool.h"
 
#if defined(FZGMOD_BACKEND_SYCL)
 
#error "FZGMOD_BACKEND=SYCL is not implemented yet"
 
#endif
 
#include <thrust/device_ptr.h>
#include <thrust/execution_policy.h>
#include <thrust/scan.h>
 
#include <cstddef>
 
namespace fz {
namespace backend {
 
namespace detail {
 
inline auto parOn(fz::stream_t stream) {
#if defined(FZGMOD_BACKEND_HIP)
    return thrust::hip::par.on(stream);
#else
    return thrust::cuda::par.on(stream);
#endif
}
 
} // namespace detail
 
struct TempStorage {
    void* ptr       = nullptr;
    bool  from_pool = false;
};
 
template<typename Fn>
TempStorage withTempStorage(MemoryPool* pool, fz::stream_t stream, const char* tag, Fn&& fn) {
    size_t bytes = 0;
    fn(nullptr, bytes);
    void* d_temp    = pool ? pool->allocate(bytes, stream, tag) : nullptr;
    bool  from_pool = (d_temp != nullptr);
    if (!d_temp && bytes > 0) {
        FZ_CUDA_CHECK(cudaMalloc(&d_temp, bytes));
        from_pool = false;
    }
    fn(d_temp, bytes);
    return TempStorage{d_temp, from_pool};
}
 
inline void freeTempStorage(MemoryPool* pool, const TempStorage& temp, fz::stream_t stream) {
    if (!temp.ptr) return;
    if (temp.from_pool && pool) {
        pool->free(temp.ptr, stream);
    } else {
        FZ_CUDA_CHECK_WARN(cudaStreamSynchronize(stream));
        FZ_CUDA_CHECK_WARN(cudaFree(temp.ptr));
    }
}
 
template<typename T>
void exclusiveScan(fz::stream_t stream, T* d_in, T* d_out, size_t n) {
    thrust::device_ptr<T> in_ptr(d_in);
    thrust::device_ptr<T> out_ptr(d_out);
    thrust::exclusive_scan(detail::parOn(stream), in_ptr, in_ptr + n, out_ptr);
}
 
} // namespace backend
} // namespace fz