A registry for operations that can be created based on operation names, type information, and device type. More...

Classes
struct	TypeID
	Type ID structure to uniquely identify operations based on input types and device type. More...

struct	TypeIDHash
	Hash function for TypeID to use in unordered_map. More...

Public Member Functions
template<DeviceType TDeviceType, typename TInput1 , typename TInput2 , typename TOutput > requires ValidTensorTypes<TInput1, TInput2> && ValidFloatTensorType<TOutput>
std::shared_ptr< BinaryOperation< TDeviceType, TInput1, TInput2, TOutput > >	createBinaryOperation (const std::string &operation_name, std::shared_ptr< DeviceContext > context, const ComponentConfig &config) const
	Create a binary operation based on the type information, device type, and operation name.

template<DeviceType TDeviceType, typename TInput , typename TOutput >
std::shared_ptr< UnaryOperation< TDeviceType, TInput, TOutput > >	createUnaryOperation (const std::string &operation_name, std::shared_ptr< DeviceContext > context, const ComponentConfig &config) const
	Create a unary operation based on the type information, device type, and operation name.

std::optional< FusedOpMeta >	findFusedMatch (const std::vector< std::string > &child_types, DeviceType device_type, std::type_index precision_type, const std::string &variant="Default")
	Find a fused operation match for a sequence of module types.

template<DeviceType TDeviceType, typename TInput1 , typename TInput2 , typename TOutput > requires ValidTensorTypes<TInput1, TInput2>&& ValidFloatTensorType<TOutput>
void	registerBinaryOperation (const std::string &operation_name, std::function< std::shared_ptr< BinaryOperation< TDeviceType, TInput1, TInput2, TOutput > >(std::shared_ptr< DeviceContext >, const ComponentConfig &)> creator)
	Register a binary operation creator for a specific device type.

template<typename TOutput >
void	registerFusedOperation (const std::vector< OperationType > &operation_types, const std::string &fused_op_name, const std::string &variant="Default")
	Register a fused operation.

template<DeviceType TDeviceType = DeviceType::Cuda, typename TInput = float, typename TOutput = TInput> requires ValidTensorType<TInput>&& ValidFloatTensorType<TOutput>
void	registerUnaryOperation (const std::string &operation_name, std::function< std::shared_ptr< UnaryOperation< TDeviceType, TInput, TOutput > >(std::shared_ptr< DeviceContext >, const ComponentConfig &)> creator)
	Register a unary operation creator for a specific device type.

Static Public Member Functions
static OperationRegistry &	instance ()
	Get the singleton instance of the OperationRegistry.

Private Types
using	GenericCreator = std::function< std::shared_ptr< void >(std::shared_ptr< DeviceContext >, const ComponentConfig &)>
	Type alias for a generic operation creator function.

Private Member Functions
	OperationRegistry ()=default
	Default constructor.

Private Attributes
std::vector< FusedOpMeta >	fused_ops_
	List of registered fused operations.

std::unordered_map< TypeID, std::unordered_map< std::string, GenericCreator >, TypeIDHash >	registry_
	Registry mapping type info to operations.

Detailed Description

A registry for operations that can be created based on operation names, type information, and device type.

This singleton class manages the registration and creation of neural network operations in the Mila framework. It provides a unified interface to access different operation implementations across various device types (CPU, CUDA) and with different data types.

The registry supports:

Type-safe registration and lookup of operations
Operation variants for specialized implementations
Fused operations for performance optimization
Automatic device-specific operation selection

Member Typedef Documentation

◆ GenericCreator

using Mila::Dnn::Compute::OperationRegistry::GenericCreator = std::function<std::shared_ptr<void>( std::shared_ptr<DeviceContext>, const ComponentConfig& )>

private

Type alias for a generic operation creator function.

Constructor & Destructor Documentation

◆ OperationRegistry()

Mila::Dnn::Compute::OperationRegistry::OperationRegistry ( )

privatedefault

Default constructor.

Member Function Documentation

◆ createBinaryOperation()

template<DeviceType TDeviceType, typename TInput1 , typename TInput2 , typename TOutput >
requires ValidTensorTypes<TInput1, TInput2> && ValidFloatTensorType<TOutput>

std::shared_ptr< BinaryOperation< TDeviceType, TInput1, TInput2, TOutput > > Mila::Dnn::Compute::OperationRegistry::createBinaryOperation	(	const std::string &	operation_name,
		std::shared_ptr< DeviceContext >	context,
		const ComponentConfig &	config
	)		const

inline

Create a binary operation based on the type information, device type, and operation name.

Template Parameters

TDeviceType	The device type for the operation (defaults to CUDA).
TInput1	The first input tensor element type.
TInput2	The second input tensor element type (defaults to TInput1).
TOutput	The output tensor element type (defaults to TInput2).

Parameters

operation_name	The name of the operation.
context	The device context to use for the operation.
precision_policy	The compute precision policy to use.

Returns: std::shared_ptr<BinaryOperation<TDeviceType, TInput1, TInput2, TOutput>> The created binary operation.

Exceptions

std::runtime_error	If the type combination, device type, or operation name is invalid.
std::invalid_argument	If the context is null.

Here is the caller graph for this function:

◆ createUnaryOperation()

template<DeviceType TDeviceType, typename TInput , typename TOutput >

std::shared_ptr< UnaryOperation< TDeviceType, TInput, TOutput > > Mila::Dnn::Compute::OperationRegistry::createUnaryOperation	(	const std::string &	operation_name,
		std::shared_ptr< DeviceContext >	context,
		const ComponentConfig &	config
	)		const

inline

Create a unary operation based on the type information, device type, and operation name.

Template Parameters

TDeviceType	The device type for the operation (defaults to CUDA).
TInput	The input tensor element type.
TOutput	The output tensor element type (defaults to TInput).

Parameters

operation_name	The name of the operation.
context	The device context to use for the operation.
precision_policy	The compute precision policy to use.

Returns: std::shared_ptr<UnaryOperation<TDeviceType, TInput, TOutput>> The created unary operation.

Exceptions

std::runtime_error	If the type combination, device type, or operation name is invalid.
std::invalid_argument	If the context is null.

Here is the caller graph for this function:

◆ findFusedMatch()

std::optional< FusedOpMeta > Mila::Dnn::Compute::OperationRegistry::findFusedMatch	(	const std::vector< std::string > &	child_types,
		DeviceType	device_type,
		std::type_index	precision_type,
		const std::string &	variant = `"Default"`
	)

inline

Find a fused operation match for a sequence of module types.

Parameters

child_types	The sequence of module types to match.
device_type	The device type.
precision_type	The precision type.
variant	The variant of the operation (defaults to "Default").

Returns: std::optional<FusedOpMeta> The matched fused operation metadata if found, or nullopt if no match.

◆ instance()

static OperationRegistry & Mila::Dnn::Compute::OperationRegistry::instance ( )

inlinestatic

Get the singleton instance of the OperationRegistry.

Returns: OperationRegistry& The singleton instance.

Here is the caller graph for this function:

◆ registerBinaryOperation()

template<DeviceType TDeviceType, typename TInput1 , typename TInput2 , typename TOutput >
requires ValidTensorTypes<TInput1, TInput2>&& ValidFloatTensorType<TOutput>

void Mila::Dnn::Compute::OperationRegistry::registerBinaryOperation	(	const std::string &	operation_name,
		std::function< std::shared_ptr< BinaryOperation< TDeviceType, TInput1, TInput2, TOutput > >(std::shared_ptr< DeviceContext >, const ComponentConfig &)>	creator
	)

inline

Register a binary operation creator for a specific device type.

Template Parameters

TDeviceType	The device type for the operation (defaults to CUDA).
TInput1	The first input tensor element type.
TInput2	The second input tensor element type (defaults to TInput1).
TOutput	The output tensor element type (defaults to TInput2).

Parameters

operation_name	The name of the operation.
creator	The function that creates the binary operation.
precision_policy	The compute precision policy to use (defaults to Auto).

Here is the caller graph for this function:

◆ registerFusedOperation()

template<typename TOutput >

void Mila::Dnn::Compute::OperationRegistry::registerFusedOperation	(	const std::vector< OperationType > &	operation_types,
		const std::string &	fused_op_name,
		const std::string &	variant = `"Default"`
	)

inline

Register a fused operation.

Template Parameters

TOutput The precision type of the operation.

Parameters

operation_types	The sequence of operation types to fuse.
fused_op_name	The name of the fused operation.
variant	The variant of the operation (defaults to "Default").

Here is the call graph for this function:

◆ registerUnaryOperation()

template<DeviceType TDeviceType = DeviceType::Cuda, typename TInput = float, typename TOutput = TInput>
requires ValidTensorType<TInput>&& ValidFloatTensorType<TOutput>

void Mila::Dnn::Compute::OperationRegistry::registerUnaryOperation	(	const std::string &	operation_name,
		std::function< std::shared_ptr< UnaryOperation< TDeviceType, TInput, TOutput > >(std::shared_ptr< DeviceContext >, const ComponentConfig &)>	creator
	)

inline

Register a unary operation creator for a specific device type.

Template Parameters

TDeviceType	The device type for the operation (defaults to CUDA).
TInput	The input tensor element type.
TOutput	The output tensor element type (defaults to TInput).

Parameters

operation_name	The name of the operation.
creator	The function that creates the unary operation.
precision_policy	The compute precision policy to use (defaults to Auto).

Here is the caller graph for this function:

Member Data Documentation

◆ fused_ops_

std::vector<FusedOpMeta> Mila::Dnn::Compute::OperationRegistry::fused_ops_

private

List of registered fused operations.

◆ registry_

std::unordered_map<TypeID, std::unordered_map<std::string, GenericCreator>, TypeIDHash> Mila::Dnn::Compute::OperationRegistry::registry_

private

Registry mapping type info to operations.

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

/home/runner/work/Mila/Mila/Mila/Src/Dnn/Compute/Operations/OperationRegistry.ixx

Classes

Public Member Functions

Static Public Member Functions

Private Types

Private Member Functions

Private Attributes

Detailed Description

Member Typedef Documentation

◆ GenericCreator

Constructor & Destructor Documentation

◆ OperationRegistry()

Member Function Documentation

◆ createBinaryOperation()

◆ createUnaryOperation()

◆ findFusedMatch()

◆ instance()

◆ registerBinaryOperation()

◆ registerFusedOperation()

◆ registerUnaryOperation()

Member Data Documentation

◆ fused_ops_

◆ registry_