device_backend.h

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Artificial Intelligence Techniques, S.L. and Roberto López González
 
#pragma once
 
#include <array>
#include <utility>
 
#include "opennn/core/opennn_types.h"
#include "opennn/core/configuration.h"
 
namespace opennn
{
 
template<typename Handle>
struct CudnnDescriptor
{
    Handle handle = nullptr;
#ifdef OPENNN_HAS_CUDA
    DnnStatus (*deleter)(Handle) = nullptr;
#else
    void (*deleter)(Handle) = nullptr;
#endif
 
    CudnnDescriptor() = default;
 
    CudnnDescriptor(CudnnDescriptor&& other) noexcept
        : handle(std::exchange(other.handle, nullptr)),
          deleter(std::exchange(other.deleter, nullptr)) {}
 
    CudnnDescriptor& operator=(CudnnDescriptor&& other) noexcept
    {
        if (this != &other)
        {
            reset();
            handle = std::exchange(other.handle, nullptr);
            deleter = std::exchange(other.deleter, nullptr);
        }
        return *this;
    }
 
    CudnnDescriptor(const CudnnDescriptor&) = delete;
    CudnnDescriptor& operator=(const CudnnDescriptor&) = delete;
 
    ~CudnnDescriptor() { reset(); }
 
    void reset()
    {
        if (handle && deleter) deleter(handle);
        handle = nullptr;
        deleter = nullptr;
    }
 
    Handle get() const noexcept { return handle; }
    operator Handle() const noexcept { return handle; }
    explicit operator bool() const noexcept { return handle != nullptr; }
};
 
}
 
namespace opennn::device
{
 
enum class CopyKind
{
    HostToHost,
    HostToDevice,
    DeviceToHost,
    DeviceToDevice
};
 
constexpr bool is_cuda_build() noexcept
{
#ifdef OPENNN_HAS_CUDA
    return true;
#else
    return false;
#endif
}
bool has_cuda_device() noexcept;
 
// Why has_cuda_device() is false, in words a user can act on: no NVIDIA
// driver, a driver older than the CUDA this build needs (with both versions),
// no card, or the error CUDA gave. Empty when a device is available.
string cuda_unavailable_reason() noexcept;
int cuda_compute_capability() noexcept;
size_t available_memory();
string gpu_info_string() noexcept;
 
// Where tuned cuBLASLt plans persist for this card and library, or empty when
// the cache is disabled, there is no CUDA device, or the build has no CUDA.
string lt_plan_cache_directory() noexcept;
bool cuda_allocation_growth_forbidden() noexcept;
void set_cuda_allocation_growth_forbidden(bool forbidden) noexcept;
void begin_cuda_shutdown() noexcept;
 
enum class GraphWorkspaceKind
{
    SharedScratch,
    Bf16Input,
    Bf16Gradient,
    Bf16ToFp32,
    Int8Dequant,
    PoolingMask,
    NormPartials,
    GradientPartials,
    FlashAttention,
    Count
};
 
inline constexpr std::array<const char*, static_cast<size_t>(GraphWorkspaceKind::Count)>
graph_workspace_labels = {"shared_scratch", "bf16_input", "bf16_gradient",
                          "bf16_to_fp32", "int8_dequant", "pooling_mask", "norm_partials",
                          "gradient_partials", "flash_attention"};
 
struct GraphWorkspaceView
{
    void* data = nullptr;
    Index bytes = 0;
};
 
using GraphWorkspaceRequirements = std::array<Index, static_cast<size_t>(GraphWorkspaceKind::Count)>;
using GraphWorkspaceViews = std::array<GraphWorkspaceView, static_cast<size_t>(GraphWorkspaceKind::Count)>;
 
#ifdef OPENNN_HAS_CUDA
bool cuda_matmul_plan_creation_forbidden() noexcept;
optional<void*> graph_workspace_override(GraphWorkspaceKind, Index minimum_bytes);
#endif
 
class CudaGraphWorkspaceScope
{
public:
    explicit CudaGraphWorkspaceScope(GraphWorkspaceRequirements&,
                                     const GraphWorkspaceViews* = nullptr);
 
    CudaGraphWorkspaceScope(const CudaGraphWorkspaceScope&) = delete;
    CudaGraphWorkspaceScope& operator=(const CudaGraphWorkspaceScope&) = delete;
 
    ~CudaGraphWorkspaceScope() noexcept;
 
private:
    GraphWorkspaceRequirements* previous_requirements = nullptr;
    const GraphWorkspaceViews* previous_views = nullptr;
    GraphWorkspaceViews owned_views;
};
 
int64_t conv_workspace_limit_bytes() noexcept;
void    set_conv_workspace_cap(int64_t mode) noexcept;
void    set_conv_workspace_auto_limit_bytes(int64_t) noexcept;
 
bool conv_autotune_enabled() noexcept;
void set_conv_autotune(bool enabled) noexcept;
 
// TF32 for the fp32 paths on tensor cores: cuBLAS/cuBLASLt fp32 GEMMs and
// cuDNN fp32 convolution engines. On by default (OPENNN_ALLOW_TF32=0 turns it
// off), which is also PyTorch's setting in the benchmark. Off means IEEE fp32
// arithmetic at several times the cost, and it exists for exactly one reason:
// validating a GPU result against a CPU reference to a tolerance TF32's
// 10-bit mantissa cannot meet. Takes effect on cuBLAS handles already created
// and on plans built after the call; cached plans carry the setting in their
// keys.
bool allow_tf32() noexcept;
void set_allow_tf32(bool enabled) noexcept;
void refresh_blas_math_mode();
 
enum class BatchNormBackwardRung { Auto, StagedFp32, PlainNative, OwnKernel };
enum class BatchNormForwardRung { Auto, CudnnGraph, OwnKernel };
enum class MaxPoolingRung { Auto, Cudnn, OwnKernel };
enum class AttentionRung { Auto, CudnnGraph, FlashAttention };
 
template<typename Rung> Rung rung() noexcept;
template<typename Rung> void set_rung(Rung) noexcept;
 
class CudaAllocationGrowthGuard
{
public:
 
    explicit CudaAllocationGrowthGuard(bool enabled,
                                       bool forbid_matmul_plan_creation = true);
 
    CudaAllocationGrowthGuard(const CudaAllocationGrowthGuard&) = delete;
    CudaAllocationGrowthGuard& operator=(const CudaAllocationGrowthGuard&) = delete;
 
    ~CudaAllocationGrowthGuard() noexcept;
 
private:
    bool active = false;
    bool previous = false;
    bool guard_matmul_plans = false;
    bool previous_matmul_plan_guard = false;
};
 
void* allocate(Device, Index);
void deallocate(Device, void*, Index) noexcept;
 
// Explicit session trimming returns only the caller's released blocks to CUDA.
class CudaBlockCacheBypass
{
public:
    CudaBlockCacheBypass() noexcept;
    ~CudaBlockCacheBypass() noexcept;
    CudaBlockCacheBypass(const CudaBlockCacheBypass&) = delete;
    CudaBlockCacheBypass& operator=(const CudaBlockCacheBypass&) = delete;
private:
    bool previous;
};
 
void set_zero(void*, Index, Device);
void set_zero_async(void*, Index, DeviceStream = nullptr);
 
void copy_async(void*, const void*, Index, CopyKind, DeviceStream = nullptr);
void copy_async(void*, const void*, Index, Device, Device, DeviceStream = nullptr);
void synchronize(DeviceStream = nullptr);
void check_last_error();
void reset_last_error() noexcept;
 
#ifdef OPENNN_HAS_CUDA
class CublasPointerModeGuard
{
public:
 
    CublasPointerModeGuard(BlasHandle, BlasPointerMode);
 
    CublasPointerModeGuard(const CublasPointerModeGuard&) = delete;
    CublasPointerModeGuard& operator=(const CublasPointerModeGuard&) = delete;
 
    ~CublasPointerModeGuard() noexcept;
 
private:
 
    BlasHandle handle = nullptr;
    BlasPointerMode previous_mode = CUBLAS_POINTER_MODE_HOST;
};
 
class CublasMathModeGuard
{
public:
 
    CublasMathModeGuard(BlasHandle, BlasMathMode);
 
    CublasMathModeGuard(const CublasMathModeGuard&) = delete;
    CublasMathModeGuard& operator=(const CublasMathModeGuard&) = delete;
 
    ~CublasMathModeGuard() noexcept;
 
private:
 
    BlasHandle handle = nullptr;
    BlasMathMode previous_mode = CUBLAS_DEFAULT_MATH;
};
#endif
 
class PinnedBuffer
{
public:
 
    PinnedBuffer() = default;
    explicit PinnedBuffer(Index byte_count);
 
    PinnedBuffer(const PinnedBuffer&) = delete;
    PinnedBuffer& operator=(const PinnedBuffer&) = delete;
 
    PinnedBuffer(PinnedBuffer&&) noexcept;
    PinnedBuffer& operator=(PinnedBuffer&&) noexcept;
 
    ~PinnedBuffer() noexcept;
 
    void resize_bytes(Index);
    void grow_to(Index);
 
    void* data() noexcept { return pointer; }
    const void* data() const noexcept { return pointer; }
 
    template<typename T>
    T* as() noexcept { return static_cast<T*>(pointer); }
 
    template<typename T>
    const T* as() const noexcept { return static_cast<const T*>(pointer); }
 
    Index byte_size() const noexcept { return allocated_bytes; }
    bool empty() const noexcept { return allocated_bytes == 0; }
    explicit operator bool() const noexcept { return pointer != nullptr; }
 
private:
 
    void reset() noexcept;
    void swap(PinnedBuffer&) noexcept;
 
    void* pointer = nullptr;
    Index allocated_bytes = 0;
};
 
void record_event(DeviceEvent, DeviceStream);
void synchronize_event(DeviceEvent);
void stream_wait_event(DeviceStream, DeviceEvent);
 
class CudaEvent
{
public:
 
    CudaEvent() = default;
    explicit CudaEvent(unsigned flags);
 
    CudaEvent(const CudaEvent&) = delete;
    CudaEvent& operator=(const CudaEvent&) = delete;
 
    CudaEvent(CudaEvent&&) noexcept;
    CudaEvent& operator=(CudaEvent&&) noexcept;
 
    ~CudaEvent() noexcept;
 
    void create();
 
    DeviceEvent get() const noexcept { return handle; }
    explicit operator bool() const noexcept { return handle != nullptr; }
 
private:
 
    void reset() noexcept;
 
    DeviceEvent handle = nullptr;
};
 
struct GraphExecDeleter { void operator()(DeviceGraphExec) const noexcept; };
 
using GraphExecHandle = unique_ptr<remove_pointer_t<DeviceGraphExec>, GraphExecDeleter>;
 
class StreamCapture
{
public:
    explicit StreamCapture(DeviceStream);
 
    StreamCapture(const StreamCapture&) = delete;
    StreamCapture& operator=(const StreamCapture&) = delete;
 
    ~StreamCapture() noexcept;
 
    void end(GraphExecHandle&);
 
private:
#ifdef OPENNN_HAS_CUDA
    DeviceStream stream = nullptr;
    bool finished = false;
#endif
};
 
void launch_graph(const GraphExecHandle&, DeviceStream);
 
constexpr int MAX_LANES = 4;
int lanes_available() noexcept;
int active_lane() noexcept;
void set_active_lane(int lane);
DeviceStream lane_stream(int lane);
 
DeviceStream get_compute_stream();
DeviceStream get_transfer_stream();
 
BlasHandle get_cublas_handle();
BlasLtHandle get_cublas_lt_handle();
DnnHandle get_cudnn_handle();
DnnOpTensorDescriptor get_op_tensor_add_descriptor();
 
}
 
namespace opennn
{
 
ThreadPoolDevice& get_device();
void set_threads_number(int);
 
struct TensorView;
 
void* ensure_workspace_bytes(device::GraphWorkspaceKind, Index bytes);
template<typename T>
T* ensure_workspace(device::GraphWorkspaceKind kind, Index count)
{
    return static_cast<T*>(ensure_workspace_bytes(kind, count * Index(sizeof(T))));
}
inline bfloat16* ensure_bf16_gradient_workspace(Index n) { return ensure_workspace<bfloat16>(device::GraphWorkspaceKind::Bf16Gradient, n); }
inline bfloat16* ensure_int8_dequant_workspace(Index n)  { return ensure_workspace<bfloat16>(device::GraphWorkspaceKind::Int8Dequant, n); }
inline float*    ensure_bf16_to_fp32_workspace(Index n)  { return ensure_workspace<float>(device::GraphWorkspaceKind::Bf16ToFp32, n); }
inline void*     ensure_shared_scratch(size_t bytes)     { return ensure_workspace_bytes(device::GraphWorkspaceKind::SharedScratch, Index(bytes)); }
inline float*    ensure_flash_attention_workspace(Index n) { return ensure_workspace<float>(device::GraphWorkspaceKind::FlashAttention, n); }
 
void release_thread_workspaces();
 
const void* data_for_gemm_dtype(const TensorView&, Type);
 
const void* bias_for_gemm_bf16(const TensorView&);
 
// D = alpha * op(A) * op(B) + beta * C, with C the addend when one is given
// and D itself otherwise -- the two scalars the wrapper used to hardcode as
// &one and &zero. lda, ldb and ldd are 0 for "derive from m, n and k", which
// is what it did before they were parameters. The operands carry a type each
// because gemm_strided_batched_cuda below already takes an Atype and a Btype:
// nothing can move off that call and onto this one until this one can say the
// same thing.
void run_lt_matmul_cached(
    int, int, int,
    BlasOperation transA,
    BlasOperation transB,
    LinearEpilogue epilogue,
    const void*, const void*, void*,
    const void*,
    DeviceDataType dtype_a   = CUDA_R_32F,
    DeviceDataType dtype_b   = CUDA_R_32F,
    DeviceDataType out_dtype = CUDA_R_32F,
    const void* aux_pointer  = nullptr,
    const void* addend       = nullptr,
    float alpha              = 1.0f,
    float beta               = 0.0f,
    int lda = 0, int ldb = 0, int ldd = 0);
 
void gemm_strided_batched_cuda(BlasOperation transa, BlasOperation transb,
                               int, int, int,
                               const void*, DeviceDataType Atype, int, long long stride_a,
                               const void*, DeviceDataType Btype, int, long long stride_b,
                               void*, DeviceDataType Ctype, int, long long stride_c,
                               int,
                               float alpha = 1.0f, float beta = 0.0f);
 
}