CPU implementation of the Encoder operation. More...

Inheritance diagram for Mila::Dnn::Compute::CpuEncoderOp:

[legend]

Collaboration diagram for Mila::Dnn::Compute::CpuEncoderOp:

[legend]

Public Types
using	ConfigType = LpeConfig
using	CpuExecutionContext = ExecutionContext<DeviceType::Cpu>
using	MR = CpuMemoryResource
using	OperationBase = UnaryOperation<DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32>
using	TensorType = Tensor<TensorDataType::FP32, MR>
Public Types inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32 >
using	MR
using	TensorInputType
using	TensorOutputType
Public Types inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
using	DataTypeTraits

Public Member Functions
	CpuEncoderOp (IExecutionContext *context, const LpeConfig &config)
	Construct with execution context and configuration.
	~CpuEncoderOp () override=default
void	backward (const ITensor &input, const ITensor &output_grad, ITensor &input_grad) const override
	Backward pass: accumulates gradients into embedding tables.
void	build (const BuildContext &config) override
	Prepare internal caches for a concrete input shape.
void	forward (const ITensor &input, ITensor &output) const override
	Forward pass: combines token and positional embeddings.
std::string	getName () const override
	Human-readable operation name.
OperationType	getOperationType () const override
	Operation type identifier.
void	setGradients (ITensor wte_grad, ITensor wpe_grad) override
	Bind gradient tensors for training.
void	setParameters (ITensor wte, ITensor wpe) override
	Bind parameter tensors for forward execution.
Public Member Functions inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32 >
virtual	~UnaryOperation ()=default
Public Member Functions inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
virtual	~Operation ()=default
virtual void	clearGradients () noexcept
	Clear any cached gradient pointers held by the operation.
virtual TensorDataType	getDataType () const
	Tensor data type for this operation.
virtual DeviceType	getDeviceType () const
	Device type for this operation.
virtual std::size_t	getStateMemorySize () const
	Returns the number of bytes of state memory allocated by this operation.
virtual bool	isBuilt () const
	Whether build() completed successfully for a concrete input shape.
virtual bool	isEvalMode () const
	Query whether operation is configured for training.
virtual void	setTrainingMode (TrainingMode training_mode)
	Configure operation training-mode behavior.

Private Member Functions
void	validateInputShape (const ITensor &input) const
void	validateInputShape (const shape_t &input_shape) const

Private Attributes
int64_t	batch_size_ { 0 }
LpeConfig	config_
IExecutionContext *	context_ { nullptr }
int64_t	embedding_dim_ { 0 }
bool	is_built_ { false }
int64_t	seq_length_ { 0 }
const float *	wpe_ { nullptr }
float *	wpe_grad_ { nullptr }
const float *	wte_ { nullptr }
float *	wte_grad_ { nullptr }

Additional Inherited Members
Static Public Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
static constexpr TensorDataType	data_type
static constexpr DeviceType	device_type
Static Protected Member Functions inherited from Mila::Dnn::Compute::UnaryOperation< DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32 >
static const TensorInputType &	asInputTensor (const ITensor &t)
static TensorOutputType &	asOutputTensor (ITensor &t)
Protected Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >
bool	is_built_
TrainingMode	training_mode_

Detailed Description

CPU implementation of the Encoder operation.

Contract and behavior:

Forward: for each batch b and sequence position t, computes output[b, t, c] = wte[input[b, t], c] + wpe[t, c] where input is a token-id tensor of shape [B, T] (INT32) and output is an embedding tensor of shape [B, T, C] (FP32).
Backward: accumulates gradients into two parameter tensors: dwte[input[b, t], :] += output_grad[b, t, :] dwpe[t, :] += output_grad[b, t, :] Gradients are accumulated (in-place) into wte_grad and wpe_grad.

Threading and safety:

Forward and backward use OpenMP parallelization across batch and sequence dimensions when available.
Backward requires atomic updates when multiple threads may update the same embedding row or positional row concurrently. The implementation uses OpenMP atomics for per-element accumulation.

Parameter binding and ownership:

setParameters(ITensor* wte, ITensor* wpe) binds raw parameter tensors. The operation does NOT take ownership; the caller (module) retains ownership and must ensure the lifetime exceeds operation usage.
Bound parameter shapes must match the EncoderConfig (vocabulary length, max sequence length and channels). If shapes mismatch, the method throws.

Edge-cases:

Out-of-range token indices are treated as fatal errors and will throw.
Token indices are discrete: no input gradient is produced. The input_grad parameter exists to satisfy the UnaryOperation interface but is unused.

Member Typedef Documentation

◆ ConfigType

using Mila::Dnn::Compute::CpuEncoderOp::ConfigType = LpeConfig

◆ CpuExecutionContext

using Mila::Dnn::Compute::CpuEncoderOp::CpuExecutionContext = ExecutionContext<DeviceType::Cpu>

◆ MR

using Mila::Dnn::Compute::CpuEncoderOp::MR = CpuMemoryResource

◆ OperationBase

using Mila::Dnn::Compute::CpuEncoderOp::OperationBase = UnaryOperation<DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32>

◆ TensorType

using Mila::Dnn::Compute::CpuEncoderOp::TensorType = Tensor<TensorDataType::FP32, MR>

Constructor & Destructor Documentation

◆ CpuEncoderOp()

Mila::Dnn::Compute::CpuEncoderOp::CpuEncoderOp	(	IExecutionContext *	context,
		const LpeConfig &	config )

inlineexplicit

Construct with execution context and configuration.

Preconditions:

context must be non-null and refer to a CPU execution context.
config must be valid (EncoderConfig::validate()).

Ownership:

The operation stores the provided execution context shared_ptr.

◆ ~CpuEncoderOp()

Mila::Dnn::Compute::CpuEncoderOp::~CpuEncoderOp ( )

overridedefault

Member Function Documentation

◆ backward()

void Mila::Dnn::Compute::CpuEncoderOp::backward	(	const ITensor &	input,
		const ITensor &	output_grad,
		ITensor &	input_grad ) const

inlineoverridevirtual

Backward pass: accumulates gradients into embedding tables.

Parameters:

input: INT32 token indices tensor with shape [B, T]
output_grad: FP32 gradients tensor with shape [B, T, C]
input_grad: unused (token indices are non-differentiable)

Semantics:

Accumulates gradients into wte_grad and wpe_grad in-place.
Uses atomic updates to ensure thread-safety when multiple threads update the same row.

Note:

Input token indices are validated and will throw on out-of-range indices.

Implements Mila::Dnn::Compute::UnaryOperation< DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32 >.

Here is the call graph for this function:

◆ build()

void Mila::Dnn::Compute::CpuEncoderOp::build ( const BuildContext & config )

inlineoverridevirtual

Prepare internal caches for a concrete input shape.

Validates the input shape and caches B, T and C for hot-path loops. Must be called (via Module::build) before forward/backward.

Reimplemented from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

Here is the call graph for this function:

◆ forward()

void Mila::Dnn::Compute::CpuEncoderOp::forward	(	const ITensor &	input,
		ITensor &	output ) const

inlineoverridevirtual

Forward pass: combines token and positional embeddings.

Parameters:

input: INT32 token indices tensor with shape [B, T]
output: FP32 tensor with shape [B, T, C] to receive embeddings

Preconditions:

build() and setParameters() must have been called.

Behavior:

Writes output in-place. Out-of-range token indices now throw.

Implements Mila::Dnn::Compute::UnaryOperation< DeviceType::Cpu, TensorDataType::INT32, TensorDataType::FP32 >.

Here is the call graph for this function:

◆ getName()

std::string Mila::Dnn::Compute::CpuEncoderOp::getName ( ) const

inlineoverridevirtual

Human-readable operation name.

Implements Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

◆ getOperationType()

OperationType Mila::Dnn::Compute::CpuEncoderOp::getOperationType ( ) const

inlineoverridevirtual

Operation type identifier.

Implements Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

◆ setGradients()

void Mila::Dnn::Compute::CpuEncoderOp::setGradients	(	ITensor *	wte_grad,
		ITensor *	wpe_grad )

inlineoverridevirtual

Bind gradient tensors for training.

Preconditions:

wte_grad and wpe_grad must be CPU tensors with shapes matching the corresponding parameter tensors.

Semantics:

Gradients are accumulated into these buffers (in-place).
Caller is responsible for zeroing gradients when appropriate.

Throws:

std::invalid_argument for null pointers, device mismatches or shape mismatches.

Reimplemented from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

Here is the call graph for this function:

◆ setParameters()

void Mila::Dnn::Compute::CpuEncoderOp::setParameters	(	ITensor *	wte,
		ITensor *	wpe )

inlineoverridevirtual

Bind parameter tensors for forward execution.

Preconditions:

wte and wpe must be CPU tensors with shapes: wte: [vocabulary_length, channels] wpe: [max_sequence_length, channels]

Ownership:

The operation stores raw data pointers to the tensor storage but does not take ownership of the ITensor objects.

Throws:

std::invalid_argument for null pointers, device mismatches or shape mismatches.

Reimplemented from Mila::Dnn::Compute::Operation< TDeviceType, TPrecision >.

Here is the call graph for this function:

◆ validateInputShape() [1/2]

void Mila::Dnn::Compute::CpuEncoderOp::validateInputShape ( const ITensor & input ) const

inlineprivate

Here is the call graph for this function:

Here is the caller graph for this function:

◆ validateInputShape() [2/2]

void Mila::Dnn::Compute::CpuEncoderOp::validateInputShape ( const shape_t & input_shape ) const

inlineprivate

Here is the call graph for this function:

Member Data Documentation

◆ batch_size_

int64_t Mila::Dnn::Compute::CpuEncoderOp::batch_size_ { 0 }

private

◆ config_

LpeConfig Mila::Dnn::Compute::CpuEncoderOp::config_

private

◆ context_

IExecutionContext* Mila::Dnn::Compute::CpuEncoderOp::context_ { nullptr }

private

◆ embedding_dim_

int64_t Mila::Dnn::Compute::CpuEncoderOp::embedding_dim_ { 0 }

private

◆ is_built_

bool Mila::Dnn::Compute::CpuEncoderOp::is_built_ { false }

private

◆ seq_length_

int64_t Mila::Dnn::Compute::CpuEncoderOp::seq_length_ { 0 }

private

◆ wpe_

const float* Mila::Dnn::Compute::CpuEncoderOp::wpe_ { nullptr }

private

◆ wpe_grad_

float* Mila::Dnn::Compute::CpuEncoderOp::wpe_grad_ { nullptr }

private

◆ wte_

const float* Mila::Dnn::Compute::CpuEncoderOp::wte_ { nullptr }

private

◆ wte_grad_

float* Mila::Dnn::Compute::CpuEncoderOp::wte_grad_ { nullptr }

private

The documentation for this class was generated from the following file:

/__w/Mila/Mila/Mila/Src/Dnn/Compute/Devices/Cpu/Operations/CpuEncoderOp.ixx

Public Types

Public Member Functions

Private Member Functions

Private Attributes

Additional Inherited Members

Detailed Description

Member Typedef Documentation

◆ ConfigType

◆ CpuExecutionContext

◆ MR

◆ OperationBase

◆ TensorType

Constructor & Destructor Documentation

◆ CpuEncoderOp()

◆ ~CpuEncoderOp()

Member Function Documentation

◆ backward()

◆ build()

◆ forward()

◆ getName()

◆ getOperationType()

◆ setGradients()

◆ setParameters()

◆ validateInputShape() [1/2]

◆ validateInputShape() [2/2]

Member Data Documentation

◆ batch_size_

◆ config_

◆ context_

◆ embedding_dim_

◆ is_built_

◆ seq_length_

◆ wpe_

◆ wpe_grad_

◆ wte_

◆ wte_grad_