/__w/Mila/Mila/Mila/Src/Dnn/Components/Activations/Gelu/Gelu.ixx

Construct GELU activation component with optional ExecutionContext ownership.

Construct GELU activation component with optional ExecutionContext ownership.Supports two construction modes:

Standalone mode (device_id provided):

• Creates and owns an ExecutionContext for the specified device.
• Registers the owned context with the base Component class via setExecutionContext().
• Backend operation is created immediately in onExecutionContextSet() hook.
• Use case: Unit tests, standalone component usage.

Shared mode (device_id not provided):

• Does not create ExecutionContext; expects parent to provide one.
• Parent (Network/CompositeComponent) calls setExecutionContext() after construction.
• Backend operation created when parent sets context.
• Use case: Components added to Network via addComponent<Gelu>(...).

Parameters

name	Component name identifier (mandatory).
config	GELU configuration (approximation method).
device_id	Optional device identifier. If provided, creates owned ExecutionContext for standalone mode. If nullopt, expects shared context from parent.

Exceptions

std::invalid_argument	if config is invalid (via config.validate()).
std::invalid_argument	if device_id.type does not match TDeviceType.
std::runtime_error	if ExecutionContext creation fails (standalone mode).
std::runtime_error	if backend operation creation fails in onExecutionContextSet().

Note

In standalone mode, setExecutionContext() is called to register the owned context with the base class, enabling getExecutionContext() and triggering the onExecutionContextSet() hook for operation creation.

// Standalone mode (owns context) GeluConfig config; Gelu<DeviceType::Cpu, TensorDataType::FP32> gelu("gelu", config, Device::Cpu());

 
module;
#include <memory>
#include <vector>
#include <string>
#include <sstream>
#include <type_traits>
#include <stdexcept>
#include <format>
#include <utility>
#include <optional>
#include <algorithm>
 
export module Dnn.Components.Gelu;
export import Dnn.Components.GeluConfig;
 
import Dnn.Component;
import Dnn.ComponentType;
import Dnn.Tensor;
import Dnn.ITensor;
import Dnn.TensorDataType;
import Dnn.TensorDataTypeTraits;
import Dnn.TensorTypes;
import Dnn.ApproximationMethod;
import Compute.Device;
import Compute.DeviceId;
import Compute.DeviceType;
import Compute.DeviceTypeTraits;
import Compute.IExecutionContext;
import Compute.ExecutionContextFactory;
import Compute.UnaryOperation;
import Compute.OperationTraits;
import Compute.CpuMemoryResource;
import Serialization.ModelArchive;
import Serialization.Tensor;
import Serialization.Mode;
import Serialization.Metadata;
import Logging.Logger;
 
namespace Mila::Dnn
{
    using namespace Mila::Dnn::Compute;
    using namespace Mila::Dnn::Serialization;

    export template<DeviceType TDeviceType, TensorDataType TPrecision>
        requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
    class Gelu : public Component<TDeviceType, TPrecision>
    {
    public:
        using MR = typename DeviceTypeTraits<TDeviceType>::memory_resource;
        using TensorType = Tensor<TPrecision, MR>;
        using ComponentBase = Component<TDeviceType, TPrecision>;

        explicit Gelu( const std::string& name, const GeluConfig& config, std::optional<DeviceId> device_id = std::nullopt )
            : ComponentBase( name ), config_( config )
        {
            config_.validate();
 
            if ( device_id.has_value() )
            {
                if ( device_id->type != TDeviceType )
                {
                    throw std::invalid_argument( "Gelu: device type mismatch" );
                }
 
                owned_exec_context_ = createExecutionContext( device_id.value() );
 
                this->setExecutionContext( owned_exec_context_.get() );
            }
        }
 
        ~Gelu() override = default;
 
        // ====================================================================
        // Computation Dispatch
        // ====================================================================

        TensorType& forward( const TensorType& input )
        {
            if ( !this->isBuilt() )
            {
                throw std::runtime_error( "Gelu::forward: component must be built before forward pass" );
            }
 
            const auto& input_shape = input.shape();
 
            if ( output_view_->shape() != input_shape )
            {
                output_view_.emplace( output_->view( input_shape ) );
            }
 
            operation_->forward( input, *output_view_ );
            
            return *output_view_;
        }

        TensorType& backward( const TensorType& input, const TensorType& output_grad )
        {
            if ( !this->isBuilt() )
            {
                throw std::runtime_error( "Gelu::backward: component must be built before backward pass" );
            }
 
            // REVIEW
            /*if ( !this->isTraining() )
            {
                throw std::runtime_error( "Gelu::backward: component must be in training mode to compute gradients" );
            }*/
 
            // Zero input gradient buffer before backward pass. No exeptions.
            // Backend ops use accumulation (atomicAdd/+=) which requires pre-zeroed buffers
            // to prevent gradient buildup across calls. Without this, gradients grow linearly
            // with each call -> explosion.
            zero( *input_grad_ );
 
            operation_->backward( input, output_grad, *input_grad_ );
 
            return *input_grad_;
        }

        void synchronize() override
        {
            this->getExecutionContext()->synchronize();
        }

        ApproximationMethod getApproximationMethod() const
        {
            return config_.getApproximationMethod();
        }
 
        // ====================================================================
        // Serialization
        // ====================================================================

        void save_( ModelArchive& archive, SerializationMode mode ) const override
        {
            (void)mode;
 
            SerializationMetadata meta;
            meta.set( "type", "Gelu" )
                .set( "version", int64_t( 1 ) )
                .set( "name", this->getName() )
                .set( "template_device", deviceTypeToString( TDeviceType ) )
                .set( "template_precision", static_cast<int64_t>(TPrecision) );
 
            archive.writeMetadata( "meta.json", meta );
 
            SerializationMetadata cfg;
            cfg.set( "approximation_method",
                static_cast<int64_t>(config_.getApproximationMethod()) );
 
            archive.writeMetadata( "config.json", cfg );
        }

        static std::unique_ptr<Gelu> fromArchive_(
            ModelArchive& archive,
            const std::string& component_name,
            IExecutionContext* exec_context )
        {
            try
            {
                SerializationMetadata meta = archive.readMetadata( "meta.json" );
                validateMetadata_( meta, component_name );
 
                SerializationMetadata cfg = archive.readMetadata( "config.json" );
 
                GeluConfig config;
                auto approx_method = static_cast<ApproximationMethod>(
                    cfg.getInt( "approximation_method" ));
                config.withApproximationMethod( approx_method );
                config.validate();
 
                return std::make_unique<Gelu>( component_name, config );
            }
            catch ( const std::exception& e )
            {
                throw std::runtime_error(
                    std::format( "Gelu::fromArchive: error for '{}': {}",
                        component_name, e.what() )
                );
            }
        }
 
        // ====================================================================
        // Parameters and Gradients
        // ====================================================================

        size_t parameterCount() const override
        {
            return 0;
        }

        std::vector<ITensor*> getParameters() const override
        {
            return {};
        }

        std::vector<ITensor*> getGradients() const override
        {
            return {};
        }
 
        // ====================================================================
        // Identification and Description
        // ====================================================================
 
        const ComponentType getType() const override
        {
            return ComponentType::Gelu;
        }
 
        // ====================================================================
        // State and Configuration
        // ====================================================================

        DeviceId getDeviceId() const override
        {
            return this->getExecutionContext()->getDeviceId();
        }

        MemoryStats getMemoryStats() const override
        {
            MemoryStats stats;
 
            if ( output_ != nullptr )
            {
                stats.device_state_bytes += output_->getStorageSize();
            }
            if ( input_grad_ != nullptr )
            {
                stats.device_gradient_bytes += input_grad_->getStorageSize();
            }
            
            return stats;
        }

        std::string toString() const override
        {
            std::ostringstream oss;
            oss << "--------------------" << std::endl;
            oss << "Gelu: " << this->getName() << std::endl;
            oss << "Device: " << deviceTypeToString( this->getDeviceType() ) << std::endl;
            // FIXME: oss << "Approximation Method: " << config_.toString( config_.getApproximationMethod() ) << std::endl;
 
            return oss.str();
        }
 
    protected:

        void onExecutionContextSet() override
        {
            createOperation();
        }

        void onBuilding( const BuildContext& build_context ) override
        {
            operation_->build( build_context );
            const auto& input_shape = build_context.inputShape();
 
            // Allocate owned output and input-gradient tensors with device binding.
            // Buffers are owned by this component and reused across calls.
            DeviceId dev_id = this->getExecutionContext()->getDeviceId();
 
            output_ = std::make_unique<TensorType>( dev_id, input_shape, this->getName() + ".output" );
            output_view_.emplace( output_->view( input_shape ) );
 
            if ( build_context.isTrainingMode() )
            {
                input_grad_ = std::make_unique<TensorType>( dev_id, input_shape, this->getName() + ".input_grad" );
                zero( *input_grad_ );
            }
        }

        void onTrainingModeChanging( TrainingMode training_mode ) override
        {
            operation_->setTrainingMode( training_mode );
        }
 
    private:
 
        using OpType = typename OperationTraits<OperationType::GeluOp, TDeviceType, TPrecision>::type;
 
        GeluConfig config_;
 
        std::unique_ptr<IExecutionContext> owned_exec_context_{ nullptr };
        std::shared_ptr<OpType> operation_{ nullptr };
 
        std::unique_ptr<TensorType> output_{ nullptr };
        std::optional<TensorType> output_view_;
        std::unique_ptr<TensorType> input_grad_{ nullptr };

        static void validateMetadata_( const SerializationMetadata& meta, const std::string& component_name )
        {
            int64_t version = meta.tryGetInt( "version" ).value_or( 0 );
            if ( version != 1 )
            {
                throw std::runtime_error(
                    std::format( "Gelu: unsupported version {} for '{}'",
                        version, component_name )
                );
            }
 
            std::string type = meta.tryGetString( "type" ).value_or( "" );
            if ( type != "Gelu" )
            {
                throw std::runtime_error(
                    std::format( "Gelu: type mismatch for '{}': expected 'Gelu', got '{}'",
                        component_name, type )
                );
            }
 
            std::string file_device = meta.tryGetString( "template_device" ).value_or( "" );
            int64_t file_precision = meta.tryGetInt( "template_precision" ).value_or( -1 );
 
            std::string expected_device = deviceTypeToString( TDeviceType );
            int64_t expected_precision = static_cast<int64_t>(TPrecision);
 
            if ( file_device != expected_device )
            {
                throw std::runtime_error(
                    std::format( "Gelu: device mismatch for '{}': archive='{}', expected='{}'",
                        component_name, file_device, expected_device )
                );
            }
 
            if ( file_precision != expected_precision )
            {
                throw std::runtime_error(
                    std::format( "Gelu: precision mismatch for '{}': archive={}, expected={}",
                        component_name, file_precision, expected_precision )
                );
            }
        }

        void createOperation()
        {
            operation_ = std::make_shared<OpType>( this->getExecutionContext(), config_ );
 
            if ( !operation_ )
            {
                throw std::runtime_error(
                    std::format( "Gelu: Failed to create compute backend operation for component '{}'",
                        this->getName() )
                );
            }
        }
    };
}