network.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 <cstdint>
#include <functional>
#include <initializer_list>
 
#include "opennn/core/configuration.h"
#include "opennn/core/tensor_types.h"
#include "opennn/core/variable.h"
#include "opennn/network/layers/layer.h"
 
namespace opennn
{
 
class Network;
class TokenizerOperator;
 
namespace detail
{
class ParameterStore
{
public:
    enum class State { Host, DeviceMaster, DeviceMasterWithMirror, DeviceCompact };
 
    State state() const noexcept
    {
        if(master.get_device() != Device::CUDA) return State::Host;
        if(!master.owns_memory()) return State::DeviceCompact;
        return bf16.empty() && int8.empty()
            ? State::DeviceMaster : State::DeviceMasterWithMirror;
    }
 
    bool low_precision_ready(Type type) const noexcept
    {
        if(type == Type::BF16) return !bf16.empty();
        if(type == Type::INT8) return !int8.empty();
        return true;
    }
 
    void clear_derived()
    {
        bf16.resize_bytes(0, Device::CUDA);
        fp32_inference.resize_bytes(0, Device::CUDA);
        int8.resize_bytes(0, Device::CUDA);
        bf16_compact = false;
    }
 
    void validate(Type type) const
    {
        throw_if(type == Type::BF16 && bf16_compact && bf16.empty(),
                 "ParameterStore: compact BF16 layout has no BF16 storage.");
        throw_if(state() == State::DeviceCompact
                 && bf16.empty() && fp32_inference.empty() && int8.empty(),
                 "ParameterStore: compact master view has no owning storage.");
        throw_if(type == Type::BF16 && !int8.empty(),
                 "ParameterStore: BF16 precision cannot own INT8 storage.");
    }
 
private:
    Buffer master;
    Buffer bf16{Device::CUDA};
    Buffer fp32_inference{Device::CUDA};
    Buffer int8{Device::CUDA};
    bool bf16_compact = false;
 
    friend class ::opennn::Network;
};
}
 
enum class NetworkTask
{
    Generic,
    Approximation,
    Classification,
    Forecasting,
    AnomalyDetection,
    ImageClassification,
    ObjectDetection,
    TextClassification,
    LanguageModeling
};
 
class Network
{
 
public:
 
    struct HostParametersGuard
    {
        explicit HostParametersGuard(Network& n)
            : network(n), was_on_device(n.get_parameters_device() == Device::CUDA)
        {
            if (was_on_device) network.copy_parameters_host();
        }
 
        ~HostParametersGuard()
        {
            if (!was_on_device) return;
            try { network.copy_parameters_device(); } catch (...) {}
        }
 
        HostParametersGuard(const HostParametersGuard&) = delete;
        HostParametersGuard& operator=(const HostParametersGuard&) = delete;
 
        Network& network;
        const bool was_on_device;
    };
 
    Network();
 
    virtual ~Network() = default;
 
    Network(const filesystem::path&);
 
    NetworkTask get_task() const noexcept { return task; }
    void set_task(NetworkTask new_task) noexcept { task = new_task; }
 
    Index add_layer(unique_ptr<Layer>,
                    const vector<Index>& = {});
 
    const EffectiveConfig& get_config() const noexcept { return config; }
    Device get_device() const noexcept { return config.device; }
    bool is_gpu() const noexcept { return config.device == Device::CUDA; }
    Type get_training_type()  const noexcept { return config.training_type; }
    const PrecisionPlan& get_precision() const noexcept { return precision; }
    Type get_activation_type() const noexcept { return precision.activations; }
 
    void set_training_activation_recomputation(bool enabled) noexcept
    {
        training_activation_recomputation = enabled;
    }
    bool get_training_activation_recomputation() const noexcept
    {
        return training_activation_recomputation;
    }
 
    void warn_if_stale_configuration() const;
 
    vector<vector<TensorSpec>> get_parameter_specs() const
    {
        return collect_layer_specs([](const Layer& layer) { return layer.get_parameter_specs(); });
    }
 
    vector<vector<TensorSpec>> get_state_specs() const
    {
        return collect_layer_specs([](const Layer& layer) { return layer.get_state_specs(); });
    }
 
    vector<vector<TensorSpec>> get_forward_specs(Index batch_size) const
    {
        auto specs = collect_layer_specs([batch_size](const Layer& layer) { return layer.get_forward_specs(batch_size); });
        if (!is_gpu()) force_specs_to_fp32(specs);
        return specs;
    }
 
    vector<vector<TensorSpec>> get_backward_specs(Index batch_size) const
    {
        auto specs = collect_layer_specs([batch_size](const Layer& layer) { return layer.get_backward_specs(batch_size); });
        if (!is_gpu()) force_specs_to_fp32(specs);
        return specs;
    }
 
    Index get_states_size() const     { return get_aligned_size(get_state_specs()); }
 
    void compile();
    void compile(Device device);
 
    void set_attention_sdpa_auto(bool);
    void set_attention_sdpa_min_sequence_length(Index);
    bool has(const string&) const;
    bool has(LayerType) const;
    bool has_recurrent_layers() const;
 
    bool is_empty() const noexcept { return layers.empty(); }
 
    // Anything derived from the parameters -- a packed cuDNN weight space, a
    // batch-norm fold, a quantized copy -- has to know when they change. There
    // is no way to observe that passively, because callers write straight
    // through the mutable handles below, so handing one out counts as a change.
    // That over-counts the few readers that take a mutable handle and only read,
    // which costs those callers a recompute and cannot make a cache stale. The
    // reverse policy, trusting writers to announce themselves, is one missed
    // call site away from silently wrong weights.
    // A handle retained across an execution has already consumed that one
    // invalidation; call mark_parameters_changed() before executing after any
    // later writes through the retained handle.
    //
    // Cache holders should store the version they were built at and compare, not
    // assume any particular increment.
    uint64_t get_parameters_version() const noexcept { return parameters_version; }
 
    void mark_parameters_changed() noexcept { ++parameters_version; }
 
    float* get_parameters_data() { mark_parameters_changed(); return parameter_store.master.as<float>(); }
    const float* get_parameters_data() const { return parameter_store.master.as<float>(); }
    VectorMap get_parameters_map() &
    {
        mark_parameters_changed();
        return parameter_store.master.as_vector();
    }
    ConstVectorMap get_parameters_map() const &
    {
        return parameter_store.master.as_vector();
    }
    Index get_parameters_buffer_size() const noexcept { return parameter_store.master.size_in_floats(); }
    Device get_parameters_device() const noexcept { return parameter_store.master.get_device(); }
    float* get_states_data() { return states.as<float>(); }
    const float* get_states_data() const { return states.as<float>(); }
    Index get_states_buffer_size() const noexcept { return states.size_in_floats(); }
 
    const vector<Variable>& get_input_variables() const noexcept { return input_variables; }
    vector<string> get_input_feature_names() const { return get_variable_feature_names(input_variables); }
 
    const vector<Variable>& get_output_variables() const noexcept { return output_variables; }
    vector<string> get_output_feature_names() const { return get_variable_feature_names(output_variables); }
 
    const vector<unique_ptr<Layer>>& get_layers() const noexcept { return layers; }
    const unique_ptr<Layer>& get_layer(const Index layer_index) const { return layers[layer_index]; }
    const unique_ptr<Layer>& get_layer(const string&) const;
 
    Index get_layer_index(const string&) const;
 
    const vector<vector<Index>>& get_source_layers() const noexcept { return source_layers; }
    vector<vector<pair<size_t, size_t>>> get_consumer_edges() const;
 
    Layer* get_first(const string&);
    Layer* get_first(LayerType);
    const Layer* get_first(const string&) const;
    const Layer* get_first(LayerType) const;
 
    void set_input_variables(const vector<Variable>& new_input_variables) { input_variables = new_input_variables; }
    void set_output_variables(const vector<Variable>& new_output_variables) { output_variables = new_output_variables; }
 
    void set_input_names(const vector<string>&);
    void set_output_names(const vector<string>&);
 
    void set_input_shape(const Shape&);
 
    void clear();
    void steal_from(Network& src);
 
    Index get_layers_number() const noexcept { return ssize(layers); }
    Index get_layers_number(const string&) const;
    Index get_layers_number(LayerType) const;
 
    Index get_first_trainable_layer_index() const;
    Index get_last_trainable_layer_index() const;
 
 
    Index get_inputs_number() const { return get_input_shape().size(); }
    Index get_outputs_number() const { return get_output_shape().size(); }
 
    Shape get_input_shape() const;
    Shape get_output_shape() const;
 
    ActivationFunction get_output_activation() const;
    Index get_parameters_number() const;
 
    void set_parameters(const VectorR&);
    void set_states(const VectorR&);
    void set_parameters_random() { initialize_parameters(&Operator::set_parameters_random); }
    void set_parameters_glorot() { initialize_parameters(&Operator::set_parameters_glorot); }
    void set_parameters_pytorch() { initialize_parameters(&Operator::set_parameters_pytorch); }
    void link_parameters();
 
    enum class ParameterStorage
    {
        Host,
        DeviceMaster,
        DeviceMasterWithMirror,
        DeviceCompact
    };
 
    ParameterStorage get_parameter_storage() const noexcept
    {
        static_assert(int(ParameterStorage::DeviceCompact)
                      == int(detail::ParameterStore::State::DeviceCompact));
        return static_cast<ParameterStorage>(parameter_store.state());
    }
 
    bool fp32_master_released() const noexcept
    {
        return get_parameter_storage() == ParameterStorage::DeviceCompact;
    }
 
    // Whether the derived storage this precision needs has been built yet. A
    // separate question from where the master lives: FP32 needs nothing derived,
    // so it is always ready.
    bool low_precision_storage_ready() const noexcept
    {
        return parameter_store.low_precision_ready(config.training_type);
    }
 
    void link_gradients(const Buffer&) const;
    void link_gradients(span<const TensorView> layer_gradients) const;
 
    void link_states();
    void link_states(Device);
    MatrixR calculate_outputs(const vector<TensorView>&);
 
    void calculate_outputs(const vector<TensorView>&, MatrixR& outputs);
 
    TensorView calculate_outputs_resident(const vector<TensorView>&,
                                          ForwardPropagation&,
                                          bool upload_parameters = true);
 
    MatrixR calculate_outputs(const MatrixR&);
 
    MatrixR calculate_outputs(const Tensor3&);
 
    MatrixR calculate_outputs(const Tensor4&);
 
    // Text input for a single-input network with a layer named "tokenizer".
    MatrixR calculate_text_outputs(const Tensor<string, 1>&);
 
    void set_tokenizer(unique_ptr<TokenizerOperator>, const string& label = "tokenizer");
    const TokenizerOperator* get_tokenizer(const string& label = "tokenizer") const;
    void set_vocabulary(const vector<string>&, const string& label = "tokenizer");
    const vector<string>& get_vocabulary(const string& label = "tokenizer") const;
 
    void calculate_outputs(const MatrixR&, MatrixR& outputs);
 
    void calculate_outputs(const Tensor3&, MatrixR& outputs);
 
    void calculate_outputs(const Tensor4&, MatrixR& outputs);
 
    Tensor3 calculate_outputs(const Tensor3&, const Tensor3&);
 
    void calculate_outputs(const Tensor3&, const Tensor3&, Tensor3& outputs);
 
    void from_JSON(const JsonDocument&);
 
    void to_JSON(JsonWriter&) const;
 
    void save(const filesystem::path&) const;
    void save_parameters_binary(const filesystem::path&) const;
    void save_states_binary(const filesystem::path&) const;
 
    void load(const filesystem::path&);
    void load_parameters_binary(const filesystem::path&);
 
    void load_parameters_bf16_inference_binary(const filesystem::path&);
    void load_states_binary(const filesystem::path&);
 
    void forward_propagate(const vector<TensorView>&,
                          ForwardPropagation&,
                          ForwardPropagationMode = ForwardPropagationMode::Inference) const;
 
    void forward_propagate(const vector<TensorView>&,
                          ForwardPropagation&,
                          ForwardPropagationMode,
                          Index,
                          Index) const;
 
    void forward_propagate(const vector<TensorView>&,
                          const VectorR&,
                          ForwardPropagation&);
 
    void cast_parameters_to_bf16();
 
    void release_bf16_fp32_parameter_master_for_inference();
 
    void upload_parameters_bf16_inference();
 
    void upload_parameters_int8_inference();
 
    bfloat16* get_parameters_bf16_mirror_data()
    {
        return config.training_type == Type::BF16 && parameter_store.master.owns_memory()
            ? parameter_store.bf16.as<bfloat16>()
            : nullptr;
    }
 
    void copy_parameters_device();
    void copy_parameters_host();
 
    void copy_states_device();
    void copy_states_host();
 
    vector<string> get_layer_labels() const;
 
protected:
 
    explicit Network(NetworkTask);
    Network(const filesystem::path&, NetworkTask);
 
    // Compile a newly built model and initialize its parameters with Glorot.
    void finalize_build();
    // Update execution settings without changing parameter/state layouts.
    // Preserve learned values and the compiled device/precision configuration.
    void configure_layers(const function<void(Layer&)>&);
    void set_attention_and_dense_dropout(float, initializer_list<string_view> dense_prefixes);
 
    // For a model factory that builds the layers and loads a BF16 inference
    // binary in one step: compiles without the fp32 master the loader would
    // discard. The network has no parameter storage of any kind until the
    // load completes, so nothing may observe it in between; a factory keeps
    // it private until then and drops it if the load throws.
    void compile_and_load_parameters_bf16_inference_binary(const filesystem::path&);
 
    NetworkTask task = NetworkTask::Generic;
 
    vector<Variable> input_variables;
    vector<Variable> output_variables;
 
    vector<unique_ptr<Layer>> layers;
 
    vector<vector<Index>> source_layers;
 
    struct DeviceResidency
    {
        const void* parameter_master = nullptr;
        const void* bf16_mirror = nullptr;
        const void* fp32_inference = nullptr;
        const void* int8_storage = nullptr;
        const void* states = nullptr;
 
        friend bool operator==(const DeviceResidency&, const DeviceResidency&) = default;
    };
 
    DeviceResidency get_device_residency() const noexcept;
 
    // Bumped whenever a mutable handle to the parameters is handed out, so a
    // derived cache can tell that what it was built from has moved on.
    uint64_t parameters_version = 1;
 
    detail::ParameterStore parameter_store;
 
    Buffer states;
 
    EffectiveConfig config;
    PrecisionPlan precision;
 
    bool training_activation_recomputation = false;
 
    mutable bool stale_configuration_warned = false;
 
    mutable const void* linked_gradient_base = nullptr;
 
private:
 
    void compile(EffectiveConfig, bool allocate_parameter_master = true);
 
    void read_parameters_bf16_inference_binary(const filesystem::path&, Index parameters_number);
 
    MatrixR calculate_outputs_device(const vector<TensorView>&, ForwardPropagation&);
 
    void calculate_outputs_device(const vector<TensorView>&, ForwardPropagation&,
                                  MatrixR& outputs);
 
    struct HostStatesGuard
    {
        explicit HostStatesGuard(Network& n)
            : HostStatesGuard(n, n.states.get_device() == Device::CUDA) {}
 
        HostStatesGuard(Network& n, bool stage)
            : network(n), was_on_device(stage)
        {
            if (was_on_device) network.copy_states_host();
        }
 
        ~HostStatesGuard()
        {
            if (!was_on_device) return;
            try { network.copy_states_device(); } catch (...) {}
        }
 
        HostStatesGuard(const HostStatesGuard&) = delete;
        HostStatesGuard& operator=(const HostStatesGuard&) = delete;
 
        Network& network;
        const bool was_on_device;
    };
 
    void initialize_parameters(void (Operator::*)());
 
    uint64_t parameter_layout_fingerprint() const;
    uint64_t state_layout_fingerprint() const;
 
    void clear_low_precision_parameter_storage();
    void activate_transposed_inference_weights();
 
    struct ParameterSlot
    {
        Layer* layer = nullptr;
        Shape shape;
        Type dtype = Type::FP32;
        bool tied = false;
        Index scale_channels = 0;
        int   scale_axis = 0;
        Index master_offset = 0;
        Index bf16_offset = 0;
        Index int8_offset = 0;
        Index fp32_offset = 0;
    };
 
    struct ParameterSlotTotals
    {
        Index bf16_elements = 0;
        Index int8_elements = 0;
        Index fp32_elements = 0;
    };
 
    ParameterSlotTotals for_each_parameter_slot(
        const function<void(const ParameterSlot&)>& visit,
        const function<void(Layer&)>& begin_layer = {}) const;
 
    void allocate_compact_parameter_storage(const ParameterSlotTotals&);
    void use_compact_parameter_storage();
 
    static void force_specs_to_fp32(vector<vector<TensorSpec>>& specs)
    {
        for (auto& layer_specs : specs)
            for (auto& spec : layer_specs)
                if (spec.dtype == Type::BF16)
                    spec.dtype = Type::FP32;
    }
 
    template<typename Fn>
    vector<vector<TensorSpec>> collect_layer_specs(Fn fn) const
    {
        vector<vector<TensorSpec>> out(layers.size());
        ranges::transform(layers, out.begin(),
                          [&](const unique_ptr<Layer>& layer) { return fn(*layer); });
        return out;
    }
};
 
}