cudnn_frontend_utilities.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
 
#ifdef OPENNN_HAS_CUDA
 
#include <cudnn_frontend.h>
#include "opennn/core/log.h"
 
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <string_view>
 
#include "opennn/core/device_backend.h"
#include "opennn/core/cuda/energy_meter.h"
#include "opennn/core/profiler.h"
#include "opennn/core/string_utilities.h"
#include "opennn/core/tensor_types.h"
 
namespace opennn::cudnn_frontend
{
using namespace ::cudnn_frontend;
 
inline constexpr size_t graph_cache_capacity = 8;
 
inline const auto check_status = [](auto status, const string& what) {
    throw_if(status.is_bad(),
             "cudnn-frontend {}: {}", what, status.get_message());
};
 
inline int device_sm_version()
{
    // Encoded as 100 * major + 10 * minor (8.6 -> 860, 12.0 -> 1200); the
    // plan-cache directory name embeds it, so the encoding must stay fixed.
    // The device current on the first caller's thread is queried once and the
    // static serves every later call, since frontend_enabled() sits on hot
    // paths. A failed query leaves a zero component and disables the frontend.
    struct QueriedCapability
    {
        int encoded = 0;
 
        static int component_of(const int device_ordinal, const cudaDeviceAttr attribute)
        {
            int component = 0;
            cudaDeviceGetAttribute(&component, attribute, device_ordinal);
            return component;
        }
 
        QueriedCapability()
        {
            int device_ordinal = 0;
            cudaGetDevice(&device_ordinal);
 
            const int major = component_of(device_ordinal, cudaDevAttrComputeCapabilityMajor);
            const int minor = component_of(device_ordinal, cudaDevAttrComputeCapabilityMinor);
 
            encoded = major * 100 + minor * 10;
        }
    };
 
    static const QueriedCapability capability;
    return capability.encoded;
}
 
inline bool frontend_enabled()
{
    return device_sm_version() >= 700;
}
 
inline bool frontend_verbose()
{
    static const bool verbose = env_flag_enabled("OPENNN_CUDNN_FRONTEND_VERBOSE");
    return verbose;
}
 
inline bool bn_frontend_enabled()
{
    return frontend_enabled() && device_sm_version() >= 800;
}
 
// Reproducibility runs must reject cuDNN engines whose reductions can change
// between identical launches. Keep this opt-in because deterministic plans
// can be slower than the normal heuristic choice.
inline bool deterministic_algorithms_enabled()
{
    static const bool enabled = env_flag_enabled("OPENNN_DETERMINISTIC", false);
    return enabled;
}
 
// The convolution autotune's second stage. cuDNN's engines for one shape
// differ in power the way cuBLASLt's tiles do -- occupancy, split-K, tile
// shape -- and the frontend's autotune ranks them by time alone. With the
// meter available, every candidate within OPENNN_CONV_ENERGY_TOLERANCE
// percent of the fastest is run for OPENNN_CONV_ENERGY_WINDOW_MS against the
// board's energy counter, and the one that costs the least energy per run is
// taken -- the rule the GEMM tile chooser applies with a modelled power, here
// with a measured one, because a convolution engine has no tile model to
// stand on. A candidate must beat the fastest by more than 2% of energy to
// displace it, so that two readings inside the meter's noise do not flip the
// choice between runs. 0 percent turns the stage off; so does a machine
// without NVML.
inline bool conv_energy_autotune_enabled()
{
    static const bool enabled = env_flag_enabled("OPENNN_CONV_ENERGY_AUTOTUNE", true);
    return enabled;
}
 
inline float conv_energy_tolerance()
{
    static const float tolerance =
        float(std::clamp(env_int_or("OPENNN_CONV_ENERGY_TOLERANCE", 10), 0LL, 100LL)) / 100.0f;
    return tolerance;
}
 
inline int64_t conv_energy_window_ms()
{
    static const int64_t window = std::clamp(env_int_or("OPENNN_CONV_ENERGY_WINDOW_MS", 300), 100LL, 5000LL);
    return window;
}
 
inline bool conv_energy_verbose()
{
    static const bool verbose = env_flag_enabled("OPENNN_CONV_ENERGY_VERBOSE");
    return verbose;
}
 
inline bool conv_energy_stage_active()
{
    return conv_energy_autotune_enabled() && conv_energy_tolerance() > 0.0f
        && device::energy_meter_available();
}
 
inline bool graph_timing_enabled()
{
    static const bool enabled = env_flag_enabled("OPENNN_GRAPH_TIMING");
    return enabled;
}
 
inline profiler::Stats& graph_timing_stats()
{
    static profiler::Stats times;
    static const bool registered = [] {
        atexit(+[] {
            profiler::Stats& stats = graph_timing_stats();
            const double total_ms = stats.total_ms();
            stats.log(
                        format("total_gpu_ms={:.1f}", total_ms),
                        total_ms,
                        "GRAPH_TIMING");
        });
        return true;
    }();
    (void)registered;
    return times;
}
 
template<typename TensorMap>
inline void execute_graph(graph::Graph& graph, 
                          TensorMap& tensors,
                          void* workspace, 
                          const string& what, 
                          const string& timing_label)
{
    if (timing_label.empty())
        return check_status(graph.execute(device::get_cudnn_handle(), tensors, workspace), what);
 
    thread_local device::CudaEvent begin(cudaEventDefault);
    thread_local device::CudaEvent end(cudaEventDefault);
    device::record_event(begin.get(), device::get_compute_stream());
 
    check_status(graph.execute(device::get_cudnn_handle(), tensors, workspace), what);
 
    device::record_event(end.get(), device::get_compute_stream());
    device::synchronize_event(end.get());
 
    float milliseconds = 0;
    CHECK_CUDA(cudaEventElapsedTime(&milliseconds, begin.get(), end.get()));
 
    graph_timing_stats().add(timing_label, milliseconds);
}
 
inline void* shared_workspace(int64_t bytes)
{
    return bytes > 0 ? ensure_shared_scratch(size_t(bytes)) : nullptr;
}
 
template<typename GraphCache, typename Body>
bool run_frontend(GraphCache& cache, const char* label, Body&& body)
{
    const lock_guard lock(cache.access_mutex);
    if (cache.disabled) return false;
 
    try
    {
        body(cache);
        return true;
    }
    catch (const exception& e)
    {
        cache.disabled = true;
        if(frontend_verbose())
            logging::warning() << label << ": cudnn-frontend path unavailable ("
                               << e.what() << ").\n";
        return false;
    }
}
 
[[noreturn]] inline void throw_frontend_unavailable(const string& what)
{
    if (!frontend_enabled())
        throw runtime_error(what + " requires a GPU of compute capability 7.0 or higher.");
 
    throw runtime_error(
        what + ": no usable cuDNN plan for this shape. The cudnn-frontend message above "
        "names the cause; under memory pressure it is normally a failed workspace or "
        "autotune allocation. Reduce the batch size, or cap the convolution workspace "
        "with device::set_conv_workspace_cap().");
}
 
// Not a switch: the enumeration gains a member with every cuDNN release and
// -Wswitch-enum would name each one; every type this library does not use is
// sized as four bytes, which is the largest it allocates for.
inline int64_t element_bytes(DataType_t dtype)
{
    if (dtype == DataType_t::BFLOAT16 || dtype == DataType_t::HALF) return 2;
    if (dtype == DataType_t::INT8 || dtype == DataType_t::UINT8
        || dtype == DataType_t::FP8_E4M3 || dtype == DataType_t::FP8_E5M2) return 1;
    return 4;
}
 
inline DataType_t to_dtype(Type t)
{
    switch (t)
    {
        case Type::Auto:
        case Type::FP32: return DataType_t::FLOAT;
        case Type::BF16: return DataType_t::BFLOAT16;
        case Type::INT8: return DataType_t::INT8;
    }
 
    throw invalid_argument("Unsupported data type for a cuDNN graph.");
}
 
inline vector<int64_t> nhwc_strides(int64_t c, int64_t h, int64_t w)
{
    return {h * w * c, 1, w * c, c};
}
 
inline shared_ptr<graph::Graph> new_graph(Type dtype = Type::FP32)
{
    auto g = make_shared<graph::Graph>();
    g->set_io_data_type(to_dtype(dtype))
      .set_intermediate_data_type(DataType_t::FLOAT)
      .set_compute_data_type(DataType_t::FLOAT);
    return g;
}
 
inline shared_ptr<graph::Tensor_attributes>
nhwc_tensor(graph::Graph& graph, const char* name,
            int64_t n, int64_t c, int64_t h, int64_t w)
{
    return graph.tensor(graph::Tensor_attributes()
                        .set_name(name)
                        .set_dim({n, c, h, w})
                        .set_stride(nhwc_strides(c, h, w)));
}
 
inline void set_nhwc_output(shared_ptr<graph::Tensor_attributes>& tensor,
                     int64_t n, int64_t c, int64_t h, int64_t w)
{
    tensor->set_output(true)
           .set_dim({n, c, h, w})
           .set_stride(nhwc_strides(c, h, w));
}
 
using VariantPack = unordered_map<shared_ptr<graph::Tensor_attributes>, void*>;
 
inline vector<int64_t> bhsd_strides(int64_t h, int64_t s, int64_t d, bool interleaved)
{
    return interleaved ? vector<int64_t>{s * h * d, d, h * d, 1}
                       : vector<int64_t>{h * s * d, s * d, d, 1};
}
 
inline shared_ptr<graph::Tensor_attributes>
bhsd_tensor(graph::Graph& graph, const char* name,
            int64_t b, int64_t h, int64_t s, int64_t d, bool interleaved)
{
    return graph.tensor(graph::Tensor_attributes()
                        .set_name(name)
                        .set_dim   ({b, h, s, d})
                        .set_stride(bhsd_strides(h, s, d, interleaved)));
}
 
inline void set_bhsd_output(shared_ptr<graph::Tensor_attributes>& tensor,
                            int64_t b, int64_t h, int64_t s, int64_t d, bool interleaved)
{
    tensor->set_output(true)
           .set_dim   ({b, h, s, d})
           .set_stride(bhsd_strides(h, s, d, interleaved));
}
 
inline shared_ptr<graph::Tensor_attributes>
per_channel_tensor(graph::Graph& graph, const char* name, int64_t channels)
{
    return graph.tensor(graph::Tensor_attributes()
                        .set_name(name)
                        .set_data_type(DataType_t::FLOAT)
                        .set_dim({1, channels, 1, 1})
                        .set_stride(nhwc_strides(channels, 1, 1)));
}
 
inline void set_per_channel_output(shared_ptr<graph::Tensor_attributes>& tensor, int64_t channels)
{
    tensor->set_output(true)
           .set_data_type(DataType_t::FLOAT)
           .set_dim({1, channels, 1, 1})
           .set_stride(nhwc_strides(channels, 1, 1));
}
 
template <typename... Args>
inline string timing_label(format_string<Args...> fmt, Args&&... args)
{
    if (!graph_timing_enabled()) return {};
    return format(fmt, std::forward<Args>(args)...);
}
 
inline shared_ptr<graph::Tensor_attributes>
scalar_tensor(graph::Graph& graph, const char* name, DataType_t dtype,
              bool pass_by_value = false, int64_t batch = 1)
{
    return graph.tensor(graph::Tensor_attributes()
                        .set_name(name)
                        .set_dim   ({batch, 1, 1, 1})
                        .set_stride({1, 1, 1, 1})
                        .set_data_type(dtype)
                        .set_is_pass_by_value(pass_by_value));
}
 
inline int64_t candidate_limit_from(long long value)
{
    return value <= 0 ? numeric_limits<int64_t>::max() : int64_t(value);
}
 
inline int64_t autotune_candidate_limit()
{
    static const int64_t limit = candidate_limit_from(env_int_or("OPENNN_AUTOTUNE_CANDIDATES", 8));
    return limit;
}
 
inline int64_t autotune_candidate_limit(const string& tag)
{
    static const auto per_kind = []
    {
        map<string, int64_t> limits;
        for (const char* kind : {"forward", "wgrad", "dgrad"})
        {
            string name = "OPENNN_AUTOTUNE_CANDIDATES_" + string(kind);
            for (char& c : name) c = char(toupper(c));
            if (const char* setting = getenv(name.c_str()); setting && *setting)
                limits[kind] = candidate_limit_from(env_int_or(name.c_str(), 0));
        }
        return limits;
    }();
 
    const auto found = per_kind.find(tag);
    return found == per_kind.end() ? autotune_candidate_limit() : found->second;
}
 
inline vector<HeurMode_t> heuristic_modes()
{
    static const vector<HeurMode_t> modes = []
    {
        const char* setting = getenv("OPENNN_CUDNN_HEURISTICS");
        const string value = setting ? setting : "A";
        if (value == "B")  return vector<HeurMode_t>{HeurMode_t::B, HeurMode_t::FALLBACK};
        if (value == "AB") return vector<HeurMode_t>{HeurMode_t::A, HeurMode_t::B, HeurMode_t::FALLBACK};
        return vector<HeurMode_t>{HeurMode_t::A, HeurMode_t::FALLBACK};
    }();
    return modes;
}
 
inline const vector<NumericalNote_t>& conv_engine_notes()
{
    static const vector<NumericalNote_t> notes = []
    {
        vector<NumericalNote_t> selected;
        const char* setting = getenv("OPENNN_CONV_ENGINE_NOTES");
        const string value = setting ? setting : "";
        if (value.find("WINOGRAD") != string::npos) selected.push_back(NumericalNote_t::WINOGRAD);
        if (value.find("FFT") != string::npos)      selected.push_back(NumericalNote_t::FFT);
        return selected;
    }();
    return notes;
}
 
inline bool build_top_candidates(graph::Graph& graph, int64_t limit)
{
    const int64_t count = graph.get_execution_plan_count();
    int64_t built = 0;
 
    for (int64_t index = 0; index < count && built < limit; ++index)
        if (graph.build_plan_at_index(index).is_good()) ++built;
 
    return built > 0;
}
 
inline bool plan_cache_enabled()
{
    static const bool enabled = env_flag_enabled("OPENNN_CUDNN_PLAN_CACHE", true);
    return enabled;
}
 
inline bool sdpa_autotune_enabled();
inline int64_t sdpa_workspace_cap_override_bytes();
 
inline const std::filesystem::path& plan_cache_directory()
{
    static const std::filesystem::path directory = []
    {
        const char* override_path = getenv("OPENNN_CUDNN_PLAN_CACHE_DIR");
 
        std::filesystem::path root;
 
        if (override_path && *override_path)
            root = std::filesystem::path(override_path);
        else
        {
            std::error_code error;
            const std::filesystem::path temporary = std::filesystem::temp_directory_path(error);
            if (error) return std::filesystem::path{};
            root = temporary / "opennn-cudnn-plans";
        }
 
        return root / format("sm{}-cudnn{}", device_sm_version(), CUDNN_VERSION);
    }();
 
    return directory;
}
 
inline std::filesystem::path plan_cache_file(const graph::Graph& graph)
{
    // The cache key is the serialised graph, and serialising asks cuDNN for
    // CUDNN_ATTR_EXECUTION_PLAN_JSON_REPRESENTATION -- an attribute the
    // Windows builds of cuDNN 9 do not implement. store_cached_plan() already
    // treats a failed serialise as "do not cache"; this path did not, so the
    // throw escaped through load_cached_plan(), the frontend disabled itself,
    // and every convolution became fatal. A machine that cannot serialise a
    // plan should lose the cache, not the GPU.
    json structure;
    try
    {
        graph.serialize(structure);
    }
    catch (const std::exception&)
    {
        static std::once_flag reported;
        std::call_once(reported, []{
            logging::warning() << "cudnn plan cache: this cuDNN build cannot serialise "
                         "execution plans; continuing without the cache.\n"; });
        return {};
    }
    structure.erase("gid");
 
    size_t selection = size_t(autotune_candidate_limit());
    for (const char* kind : {"forward", "wgrad", "dgrad"})
        selection = selection * 31 + size_t(autotune_candidate_limit(kind));
    for (const HeurMode_t mode : heuristic_modes()) selection = selection * 31 + size_t(mode) + 1;
    for (const NumericalNote_t note : conv_engine_notes()) selection = selection * 31 + size_t(note) + 7;
    selection = selection * 31 + (deterministic_algorithms_enabled() ? 11u : 0u);
 
    selection = selection * 31 + (sdpa_autotune_enabled() ? 3u : 0u);
 
    // The cached entry is the plan that won an autotune, and the SDPA workspace
    // cap changes which plan that is, so a capped build must not load an
    // uncapped winner from the same directory. Only the override belongs here:
    // the shape half of the cap is a function of the extents, which the
    // serialised graph already carries.
    selection = selection * 31 + size_t(sdpa_workspace_cap_override_bytes() + 2);
 
    // Whether the attention plan was autotuned belongs in the key for the
    // same reason the workspace cap does: a heuristic plan stored by a run
    // with the autotune off must not be handed to a run that asked for the
    // measured one, or the knob measures nothing once the cache is warm.
    const size_t key = std::hash<json>{}(structure)
        ^ (std::hash<int64_t>{}(device::conv_workspace_limit_bytes()) << 1)
        ^ (std::hash<bool>{}(device::conv_autotune_enabled()) << 2)
        ^ (std::hash<size_t>{}(selection) << 3)
        ^ (std::hash<bool>{}(sdpa_autotune_enabled()) << 4)
        ^ (std::hash<bool>{}(device::allow_tf32()) << 6)
        ^ (std::hash<size_t>{}(conv_energy_stage_active()
                                   ? size_t(1000 + int(conv_energy_tolerance() * 100.0f + 0.5f))
                                   : size_t(0)) << 5);
 
    return plan_cache_directory() / format("{:016x}.plan", key);
}
 
inline bool load_cached_plan(graph::Graph& graph, const cudnnHandle_t handle,
                             int64_t& workspace_bytes)
{
    if (!plan_cache_enabled() || plan_cache_directory().empty()) return false;
 
    std::error_code failed;
    const std::filesystem::path file = plan_cache_file(graph);
    if (file.empty()) return false;
    if (!std::filesystem::exists(file, failed) || failed) return false;
 
    std::ifstream stream(file, std::ios::binary);
    if (!stream) return false;
 
    const vector<uint8_t> blob((std::istreambuf_iterator<char>(stream)),
                                std::istreambuf_iterator<char>());
    if (blob.empty()) return false;
 
    constexpr bool enforce_precompiled = false;
    constexpr bool run_warmup = false;
 
    if (graph.deserialize(handle, blob, enforce_precompiled, run_warmup).is_bad()) return false;
 
    return graph.get_workspace_size(workspace_bytes).is_good();
}
 
inline void store_cached_plan(graph::Graph& graph)
{
    if (!plan_cache_enabled() || plan_cache_directory().empty()) return;
 
    vector<uint8_t> blob;
 
    // Status-checking this is not enough: cudnn-frontend throws out of
    // serialize() rather than returning when the backend refuses the JSON
    // attribute, which is what Windows cuDNN does. The status check stays for
    // the builds that report it politely.
    try
    {
        if (graph.serialize(blob, false).is_bad() || blob.empty()) return;
    }
    catch (const std::exception&)
    {
        return;
    }
 
    std::error_code failed;
    std::filesystem::create_directories(plan_cache_directory(), failed);
    if (failed) return;
 
    static atomic<uint64_t> sequence{0};
 
    const std::filesystem::path file = plan_cache_file(graph);
    if (file.empty()) return;
    const std::filesystem::path pending = file.string()
        + format(".{:x}-{}.tmp", std::hash<thread::id>{}(this_thread::get_id()), sequence++);
 
    {
        std::ofstream stream(pending, std::ios::binary | std::ios::trunc);
        if (!stream) return;
        stream.write(reinterpret_cast<const char*>(blob.data()), std::streamsize(blob.size()));
        if (!stream) { std::filesystem::remove(pending, failed); return; }
    }
 
    std::filesystem::rename(pending, file, failed);
    if (failed) std::filesystem::remove(pending, failed);
}
 
inline int64_t autotune_workspace_bytes(const graph::Graph& graph)
{
    int64_t maximum = 0;
    const int64_t count = graph.get_execution_plan_count();
 
    for (int64_t index = 0; index < count; ++index)
    {
        int64_t bytes = 0;
        if (graph.get_workspace_size_plan_at_index(index, bytes).is_good())
            maximum = max(maximum, bytes);
    }
 
    return maximum;
}
 
// On by default since the transformer cells were re-measured with it
// (2026-09-06, RTX 5070 Ti, cuDNN 9.25.1): timing cuDNN's attention engines
// under the workspace cap below instead of taking heuristic A's first pick
// read 5,413 against 5,335 sequences/s on inference and 1,353 against 1,330
// on training, repeatable to the unit on launches that otherwise do not move.
// OPENNN_SDPA_AUTOTUNE=0 restores the heuristic plan.
inline bool sdpa_autotune_enabled()
{
    static const bool enabled = env_flag_enabled("OPENNN_SDPA_AUTOTUNE", true);
    return enabled;
}
 
// Megabytes, mirroring OPENNN_CONV_WORKSPACE_MB: 0 removes the cap and restores
// pick-by-time over every candidate, a positive value pins it, and unset leaves
// the shape-derived bound below. It exists so the cap stays measurable rather
// than baked in -- it is the knob a sweep varies now that the autotune is on
// by default (see sdpa_autotune_enabled()).
inline int64_t sdpa_workspace_cap_override_bytes()
{
    static const int64_t bytes = []
    {
        const long long megabytes = env_int_or("OPENNN_SDPA_WORKSPACE_MB", -1);
        return megabytes < 0 ? int64_t(-1) : int64_t(megabytes) * 1024 * 1024;
    }();
    return bytes;
}
 
// A plan may not ask for more scratch than an fp32 copy of the node's own
// Q, K and V. That is the same rule the convolution path applies -- bound the
// workspace by the tensors the layer already pays for -- but it has to be
// re-derived here rather than reused, because
// device::conv_workspace_limit_bytes() is the largest *convolution* activation
// slot (forward_propagation.cpp) and the cnn benchmark pins it flat at 16 MiB.
// Neither number is a statement about an attention node, and 16 MiB would bar
// every plan the transformer cell needs.
//
// The bound is exactly the size of the fp32 dQ/dK/dV accumulators a flash
// engine legitimately keeps, so it excludes no flash plan, while an engine that
// materialises the B*H*Sq*Skv score matrix in fp32 is admitted for at most one
// such buffer and excluded outright once Sq passes 3*head_dim, where that
// matrix overtakes the accumulators and then keeps growing quadratically. At
// the transformer cell's 32 x 8 x 130 x 64 that is a 24.5 MiB cap against a
// 16.5 MiB score matrix and 12.2 MiB of bf16 Q/K/V.
inline int64_t sdpa_workspace_cap_bytes(int64_t batch, int64_t heads,
                                        int64_t q_seq, int64_t src_seq, int64_t head_dim)
{
    if (const int64_t override_bytes = sdpa_workspace_cap_override_bytes(); override_bytes >= 0)
        return override_bytes;
 
    const int64_t operands   = 4 * batch * heads * head_dim * (q_seq + 2 * src_seq);
    const int64_t statistics = 4 * batch * heads * q_seq;
 
    // The floor keeps a small shape -- a unit test, a short prompt -- from being
    // capped below a plan's fixed scratch (tile counters, semaphores, the
    // frontend's own node workspace), where every candidate is small anyway and
    // there is nothing to exclude.
    return max(operands + statistics, int64_t(4) * 1024 * 1024);
}
 
inline bool finalize_attention(graph::Graph& graph, const string& tag, int64_t& workspace_bytes,
                               bool allow_autotune = false, int64_t workspace_cap = 0)
{
    const cudnnHandle_t handle = device::get_cudnn_handle();
 
    check_status(graph.validate(), tag + " validate");
 
    if (load_cached_plan(graph, handle, workspace_bytes)) return false;
 
    check_status(graph.build_operation_graph(handle), tag + " build_operation_graph");
 
    if (allow_autotune && sdpa_autotune_enabled())
    {
        check_status(graph.create_execution_plans({HeurMode_t::A, HeurMode_t::B, HeurMode_t::FALLBACK}),
                     tag + " create_execution_plans");
 
        if (workspace_cap > 0)
            graph.deselect_workspace_greater_than(workspace_cap);
 
        if (build_top_candidates(graph, autotune_candidate_limit()))
        {
            // What the slot carries is the workspace the graph would execute
            // with, which is the candidate plan's; autotune_now() overwrites it
            // with the winner's and sizes its own tuning scratch from
            // autotune_workspace_bytes(). Parking that max-over-candidates here
            // instead left the slot's persistent size at the largest built
            // candidate's, and that is the size shared_workspace() grows to on
            // any path where the tuning run does not happen.
            check_status(graph.get_workspace_size(workspace_bytes), tag + " workspace");
            return true;
        }
 
        // Nothing survived the cap, so the heuristic plan below has to be chosen
        // without it. The cap narrows an autotune; it must never be the reason
        // attention is left with no plan at all.
        graph.deselect_workspace_greater_than(numeric_limits<int64_t>::max());
    }
 
    check_status(graph.create_execution_plans({HeurMode_t::A}), tag + " create_execution_plans");
    check_status(graph.build_plans(handle, BuildPlanPolicy_t::HEURISTICS_CHOICE), tag + " build_plans");
    check_status(graph.get_workspace_size(workspace_bytes), tag + " workspace");
 
    store_cached_plan(graph);
    return false;
}
 
inline shared_ptr<graph::Tensor_attributes>
seq_len_scalar(graph::Graph& graph, const char* name, int64_t batch = 1)
{
    return scalar_tensor(graph, name, DataType_t::INT32, false, batch);
}
 
// fp32_graph: the graph's operands are fp32, so its tensor-core engines would
// run TF32; when TF32 is disallowed those engines are deselected and the plan
// is IEEE fp32 (see device::allow_tf32()).
inline bool finalize(graph::Graph& graph, int64_t& workspace_bytes, const string& tag,
                     const bool fp32_graph = false)
{
    const cudnnHandle_t handle = device::get_cudnn_handle();
    const bool request_autotune = device::conv_autotune_enabled();
    const bool ieee_fp32 = fp32_graph && !device::allow_tf32();
 
    workspace_bytes = 0;
 
    check_status(graph.validate(), tag + " validate");
 
    if (load_cached_plan(graph, handle, workspace_bytes)) return false;
 
    check_status(graph.build_operation_graph(handle), tag + " build_operation_graph");
    const int64_t conv_workspace_cap = device::conv_workspace_limit_bytes();
 
    const bool convolution_graph = tag == "forward" || tag == "wgrad" || tag == "dgrad";
    const auto prepare_candidates = [&](bool restrict_notes)
    {
        check_status(graph.create_execution_plans(heuristic_modes()), tag + " create_execution_plans");
        if (conv_workspace_cap > 0)
            graph.deselect_workspace_greater_than(conv_workspace_cap);
        // IEEE fp32 means no tensor-core (TF32) engine, no engine that
        // down-converts its inputs, and no engine that reduces in a narrower
        // type -- the three notes cuDNN uses to say "not the arithmetic the
        // operands are declared in".
        if (ieee_fp32)
            graph.deselect_numeric_notes({NumericalNote_t::TENSOR_CORE,
                                          NumericalNote_t::DOWN_CONVERT_INPUTS,
                                          NumericalNote_t::REDUCED_PRECISION_REDUCTION});
        if (deterministic_algorithms_enabled())
            graph.deselect_numeric_notes({NumericalNote_t::NONDETERMINISTIC});
        if (restrict_notes)
            graph.select_numeric_notes(conv_engine_notes());
    };
 
    const auto build_candidates = [&]() -> bool
    {
        if (request_autotune)
            return build_top_candidates(graph, autotune_candidate_limit(tag));
        return graph.build_plans(handle, BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good();
    };
 
    bool restricted = convolution_graph && !conv_engine_notes().empty();
    prepare_candidates(restricted);
    bool built = build_candidates();
    if (!built && restricted)
    {
        restricted = false;
        prepare_candidates(false);
        built = build_candidates();
    }
    if (!built)
        check_status(graph.build_plans(handle, BuildPlanPolicy_t::HEURISTICS_CHOICE), tag + " build_plans");
 
    if (request_autotune && built) return true;
 
    check_status(graph.get_workspace_size(workspace_bytes), tag + " get_workspace_size");
 
    store_cached_plan(graph);
 
    return false;
}
 
inline void report_autotune_skipped(const char* tag, const char* reason)
{
    if(frontend_verbose() || conv_energy_verbose())
        logging::warning() << (tag ? tag : "autotune")
             << ": autotune skipped, keeping the heuristic plan (" << reason << ").\n";
}
 
// run_slot() labels its calls "ConvolutionOperator fwd", "... wgrad",
// "... dgrad", "... bgrad" and "... fwd folded"; attention's are "sdpa fwd"
// and "sdpa bwd". The stage applies to the convolution family only.
inline bool is_convolution_tag(const char* tag)
{
    return tag && std::string_view(tag).find("ConvolutionOperator") != std::string_view::npos;
}
 
// Runs after graph.autotune() has ranked the built plans by time, winner at
// index 0. Re-times each plan briefly (the frontend keeps its timings
// private), meters the energy of every plan inside the tolerance, and makes
// the cheapest of them the candidate. build_plan_at_index() on an already
// built plan is the frontend's way of selecting it: it skips the build and
// sets the candidate, which is what execute() and serialize() then use.
template<typename TensorMap>
inline void select_plan_by_energy(graph::Graph& graph, TensorMap& tensors, void* workspace, const char* tag)
{
    const int64_t count = graph.get_execution_plan_count();
    if (count < 2) return;
 
    const cudnnHandle_t handle = device::get_cudnn_handle();
    cudaStream_t stream = nullptr;
    cudnnGetStream(handle, &stream);
    cudaEvent_t start = nullptr, stop = nullptr;
    if (cudaEventCreate(&start) != cudaSuccess || cudaEventCreate(&stop) != cudaSuccess) return;
 
    struct Reading { int64_t index; float microseconds; double millijoules; };
    vector<Reading> timed;
 
    // Time: the median of five back-to-back runs, after one warm-up, on the
    // same stream the plan will execute on.
    for (int64_t index = 0; index < count; ++index)
    {
        if (graph.execute_plan_at_index(handle, tensors, workspace, index).is_bad()) continue;
        float samples[5] = {};  // microseconds
        bool ok = true;
        for (float& sample : samples)
        {
            cudaEventRecord(start, stream);
            ok = graph.execute_plan_at_index(handle, tensors, workspace, index).is_good();
            cudaEventRecord(stop, stream);
            cudaEventSynchronize(stop);
            if (!ok) break;
            float milliseconds = 0.0f;
            cudaEventElapsedTime(&milliseconds, start, stop);
            sample = milliseconds * 1000.0f;
        }
        if (!ok) continue;
        std::sort(samples, samples + 5);
        timed.push_back({index, samples[2], 0.0});
    }
    if (timed.size() < 2) { cudaEventDestroy(start); cudaEventDestroy(stop); device::reset_last_error(); return; }
 
    const auto fastest = std::min_element(timed.begin(), timed.end(),
        [](const Reading& a, const Reading& b) { return a.microseconds < b.microseconds; });
    const float budget = fastest->microseconds * (1.0f + conv_energy_tolerance());
 
    // Energy: each admissible plan runs back to back for the window while
    // the driver's power ring samples the board every ~20 ms; the mean of
    // the samples inside the window times the window is the energy, divided
    // by the runs. A 300 ms window holds about fifteen samples, enough to
    // put the mean well inside the 2% margin below; a window that turns out
    // to hold fewer than five is not trusted and leaves the time-ranked
    // choice alone.
    vector<Reading*> admissible;
    for (Reading& reading : timed)
        if (reading.microseconds <= budget) admissible.push_back(&reading);
    if (admissible.size() < 2) { cudaEventDestroy(start); cudaEventDestroy(stop); device::reset_last_error(); return; }
 
    for (Reading* reading : admissible)
    {
        const int64_t runs = std::max<int64_t>(10, int64_t(double(conv_energy_window_ms()) * 1000.0
                                                           / double(std::max(reading->microseconds, 1.0f))));
        (void)graph.execute_plan_at_index(handle, tensors, workspace, reading->index);
        cudaStreamSynchronize(stream);
        const unsigned long long window_begin = device::energy_meter_begin();
        cudaEventRecord(start, stream);
        for (int64_t run = 0; run < runs; ++run)
            (void)graph.execute_plan_at_index(handle, tensors, workspace, reading->index);
        cudaEventRecord(stop, stream);
        cudaEventSynchronize(stop);
        float milliseconds = 0.0f;
        cudaEventElapsedTime(&milliseconds, start, stop);
        int samples = 0;
        const double watts = device::energy_meter_window_watts(window_begin, samples);
        reading->millijoules = samples >= 5 && watts > 0.0
            ? watts * double(milliseconds) / double(runs)      // W * ms = mJ
            : 0.0;
        // The window's own timing is the better time figure -- many runs, not
        // five -- so it replaces the short one for the decision below.
        reading->microseconds = milliseconds * 1000.0f / float(runs);
    }
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    device::reset_last_error();
 
    Reading* baseline = nullptr;
    for (Reading* reading : admissible)
        if (reading->millijoules > 0.0 && (!baseline || reading->microseconds < baseline->microseconds))
            baseline = reading;
    if (!baseline) return;
 
    Reading* chosen = baseline;
    for (Reading* reading : admissible)
        if (reading->millijoules > 0.0 && reading->millijoules < chosen->millijoules * 0.98
            && reading->microseconds <= baseline->microseconds * (1.0f + conv_energy_tolerance()))
            chosen = reading;
 
    if (conv_energy_verbose())
    {
        string chosen_name, fastest_name;
        (void)graph.get_plan_name_at_index(chosen->index, chosen_name);
        (void)graph.get_plan_name_at_index(baseline->index, fastest_name);
        auto message = logging::info();
        message << (tag ? tag : "conv") << ": energy autotune over " << admissible.size() << " of "
             << timed.size() << " plans: "
             << (chosen == baseline ? "kept the fastest " : "chose ") << chosen_name
             << " (" << chosen->microseconds << " us, " << chosen->millijoules << " mJ)";
        if (chosen != baseline)
            message << " over " << fastest_name << " (" << baseline->microseconds << " us, "
                 << baseline->millijoules << " mJ)";
        message << "\n";
    }
 
    if (chosen != baseline)
    {
        const auto selected = graph.build_plan_at_index(handle, chosen->index);
        (void)selected;
    }
}
 
template<typename TensorMap>
inline void autotune_now(bool& pending, graph::Graph& graph,
                         TensorMap& tensors, int64_t& workspace_bytes,
                         const char* tag = nullptr)
{
    if (!pending) return;
    pending = false;
    if (device::lanes_available() > 1) device::synchronize();
 
    Buffer tune_workspace{Device::CUDA};
    try
    {
        const int64_t tune_bytes = autotune_workspace_bytes(graph);
        if (tune_bytes > 0) tune_workspace.resize_bytes(Index(tune_bytes), Device::CUDA);
        check_status(graph.autotune(device::get_cudnn_handle(), tensors, tune_workspace.data()), "autotune");
 
        if (is_convolution_tag(tag) && conv_energy_stage_active())
            select_plan_by_energy(graph, tensors, tune_workspace.data(), tag);
 
        store_cached_plan(graph);
    }
    catch (const exception& e)
    {
        report_autotune_skipped(tag, e.what());
    }
    catch (...)
    {
        report_autotune_skipped(tag, "unknown error");
    }
 
    device::reset_last_error();
 
    workspace_bytes = graph.get_workspace_size();
}
 
template<typename TensorMap>
inline void autotune_with_scratch(bool& pending, graph::Graph& graph,
                                  const TensorMap& tensors, int64_t& workspace_bytes,
                                  const char* tag = nullptr)
{
    if (!pending) return;
    if (device::lanes_available() > 1) device::synchronize();
 
    TensorMap scratch = tensors;
    vector<Buffer> buffers;
    buffers.reserve(scratch.size());
 
    try
    {
        for (auto& [tensor, pointer] : scratch)
        {
            if (tensor->get_is_pass_by_value()) continue;
 
            int64_t elements = 1;
            for (const int64_t dimension : tensor->get_dim()) elements *= dimension;
 
            Buffer& buffer = buffers.emplace_back(Device::CUDA);
            buffer.resize_bytes(Index(elements * element_bytes(tensor->get_data_type())),
                                Device::CUDA);
            pointer = buffer.data();
        }
    }
    catch (const exception& e)
    {
        pending = false;
        buffers.clear();
        device::reset_last_error();
        report_autotune_skipped(tag, e.what());
        workspace_bytes = graph.get_workspace_size();
        return;
    }
 
    autotune_now(pending, graph, scratch, workspace_bytes, tag);
}
 
struct GraphSlot
{
    shared_ptr<graph::Graph> graph;
    int64_t workspace_bytes = 0;
    bool autotune_pending = false;
 
    explicit operator bool() const noexcept { return graph != nullptr; }
    graph::Graph& operator*() const noexcept { return *graph; }
 
    void build(shared_ptr<graph::Graph> built, const string& tag, const bool fp32_graph = false)
    {
        graph.reset();
        autotune_pending = finalize(*built, workspace_bytes, tag, fp32_graph);
        graph = std::move(built);
    }
 
    // Attention finalizes under its own policy -- its own autotune switch, its
    // own heuristic modes and its own workspace cap -- so it cannot share
    // build()'s convolution one. Keeping it a named second entry point makes
    // that difference visible instead of hiding it behind a slot that looks
    // interchangeable. The cap comes from the caller because it is a function
    // of the node's extents, and the graph exposes those only through a full
    // serialise, not an accessor; sdpa_workspace_cap_bytes() derives it from
    // the cache key instead, and 0 leaves the candidates uncapped.
    void build_attention(shared_ptr<graph::Graph> built, const string& tag,
                         bool allow_autotune = true, int64_t workspace_cap = 0)
    {
        PROFILE_SCOPE_HOST("cudnn:build " + tag);
 
        graph.reset();
        autotune_pending = finalize_attention(*built, tag, workspace_bytes,
                                              allow_autotune, workspace_cap);
        graph = std::move(built);
    }
};
 
template<typename TensorMap>
inline void run_slot(GraphSlot& slot, TensorMap& tensors, const char* what,
                     const string& timing_label, bool scratch_autotune)
{
    if (scratch_autotune)
        autotune_with_scratch(slot.autotune_pending, *slot.graph, tensors, slot.workspace_bytes, what);
    else
        autotune_now(slot.autotune_pending, *slot.graph, tensors, slot.workspace_bytes, what);
    execute_graph(*slot.graph, tensors, shared_workspace(slot.workspace_bytes), what, timing_label);
}
 
}
 
#endif