Loading...
Searching...
No Matches
Mila::Dnn::Compute Namespace Reference
Namespaces | |
| namespace | Cpu |
| namespace | Cuda |
| namespace | OperationNames |
Concepts | |
| concept | GqaOpConcept |
| Contract for GroupedQueryAttentionOp: positional forward and backward. | |
| concept | IsAcceleratorMemoryResource |
| Concept for identifying compute accelerator memory resources. | |
| concept | IsCacheCoherent |
| Concept for cache-coherent memory resources. | |
| concept | IsDeviceMemoryResource |
| Concept for identifying device-based memory resources. | |
| concept | IsHighBandwidth |
| Concept for memory resources with high bandwidth characteristics. | |
| concept | IsHostMemoryResource |
| Concept for identifying host-accessible memory resources. | |
| concept | LinearOpConcept |
| Contract for LinearOp: typed forward matmul and backward weight/input gradients. | |
| concept | OptimizedForCoalescing |
| Concept for memory resources optimized for coalesced access patterns. | |
| concept | RequiresContextBinding |
| Concept for memory resources requiring CUDA context binding. | |
| concept | SamplingOpConcept |
| Contract for SamplingOp: in-place token sampling from a logits tensor. | |
| concept | SupportsConcurrentKernels |
| Concept for memory resources supporting concurrent kernel execution. | |
| concept | SupportsPeerAccess |
| Concept for memory resources optimized for device-to-device transfers. | |
| concept | SupportsSIMD |
| Concept for CPU memory resources with SIMD support. | |
| concept | SupportsTextureMemory |
| Concept for memory resources supporting texture memory access. | |
| concept | SupportsThreading |
| Concept for memory resources with threading support. | |
| concept | SupportsUnifiedMemory |
| Concept for CUDA memory resources with unified memory support. | |
| concept | UnaryOpConcept |
| Contract for policy-free unary ops (Softmax, RmsNorm, LayerNorm, Residual, ...). | |
Typedefs | |
| using | AdamWConfig = Mila::Dnn::Optimizers::AdamWConfig |
| using | Mila::Dnn::Compute::CpuExecutionContext = ExecutionContext<DeviceType::Cpu> |
| using | Mila::Dnn::Compute::CudaExecutionContext = ExecutionContext<DeviceType::Cuda> |
| using | Mila::Dnn::Compute::HostMemoryResource = CpuMemoryResource |
| Alias for CpuMemoryResource that represents host-accessible memory. | |
| using | OptimizerConfig = Mila::Dnn::Optimizers::AdamWConfig |
Enumerations | |
| enum class | Mila::Dnn::Compute::DeviceType { Cpu , Cuda , Metal , Rocm } |
| Enumeration of supported compute device types. More... | |
| enum class | Mila::Dnn::Compute::OperationType { CrossEntropyOp , TokenEmbeddingOp , LpeOp , RopeOp , FusedOp , LinearOp , GeluOp , SwigluOp , LayerNormOp , RmsNormOp , MultiHeadAttentionOp , GroupedQueryAttentionOp , ResidualOp , SoftmaxOp , DropoutOp , SamplingOp , SoftmaxCrossEntropyOp } |
| Enumeration of all supported neural network operation types. More... | |
Functions | |
| template<typename Tp, typename Tg> | |
| void | adamw_update (Tp *params_memory, float *master_params_memory, Tg *grads_memory, float *m_memory, float *v_memory, size_t num_parameters, ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride, int num_slices, float learning_rate, float beta1, float beta2, int t, float eps, float weight_decay, float grad_scale, unsigned int seed, cudaStream_t stream) |
| template<DeviceType TDeviceType> | |
| const ExecutionContext< TDeviceType > * | Mila::Dnn::Compute::cast_context_ (const IExecutionContext *ctx) noexcept |
| template<DeviceType TDeviceType> | |
| ExecutionContext< TDeviceType > * | Mila::Dnn::Compute::cast_context_ (IExecutionContext *ctx) noexcept |
| std::unique_ptr< IExecutionContext > | Mila::Dnn::Compute::createExecutionContext (DeviceId device_id) |
| Create execution context for specified device. | |
| void | Mila::Dnn::Compute::cublasLtCheckStatus (cublasStatus_t status, const std::source_location &location=std::source_location::current()) |
| Checks the status of a cuBLASLt operation and throws if an error occurred. | |
| void | cuda_mha_forward_fp16 (half *Y, half *qkvr, half *att, const half *X, int B, int T, int C, int NH, cudaStream_t stream) |
| void | cuda_mha_forward_fp32 (float *Y, float *qkvr, float *att, const float *X, int B, int T, int C, int NH, cudaStream_t stream) |
| template<typename TPrecision> | |
| void | cuda_softmax_crossentropy_backward (TPrecision *dX, const TPrecision *dY_loss, const TPrecision *Y, const int *targets, int batch_size, int seq_len, int vocab_size, cudaStream_t stream) |
| template<typename TPrecision> | |
| void | cuda_softmax_crossentropy_forward (TPrecision *Y_loss, TPrecision *Y, const TPrecision *X, const int *targets, int batch_size, int seq_len, int vocab_size, cudaStream_t stream) |
| void | Mila::Dnn::Compute::cudaCheckLastError (const std::source_location &location=std::source_location::current()) |
| Checks the last CUDA error and throws if an error occurred. | |
| void | Mila::Dnn::Compute::cudaCheckStatus (cudaError_t status, const std::source_location &location=std::source_location::current()) |
| Checks the status of a CUDA operation and throws if an error occurred. | |
| std::string | Mila::Dnn::Compute::deviceTypeToString (DeviceType device_type) |
| Converts a DeviceType to its string representation. | |
| DeviceId | Mila::Dnn::Compute::getBestDevice (DeviceType type, bool preferMemory=false) |
| Gets the best DeviceId of a specific type based on performance characteristics. | |
| std::size_t | Mila::Dnn::Compute::getDeviceCount (DeviceType type) noexcept |
| Count instantiated devices of the given DeviceType. | |
| template<DeviceType TDevice, TensorDataType TDataType> | |
| std::vector< std::string > | Mila::Dnn::Compute::getRegisteredOperations () |
| Templated helper returning registered operation names for a compile-time device and tensor data type. | |
| template<DeviceType TDevice, TensorDataType TDataType> | |
| bool | Mila::Dnn::Compute::isOperationRegistered (const std::string &operation_name) |
| Templated helper that checks whether a named operation is registered for a compile-time device and tensor data type. | |
| std::vector< std::string > | Mila::Dnn::Compute::listDevicesByName () |
| Lists all available compute devices by name. | |
| std::vector< std::string > | Mila::Dnn::Compute::listDevicesByType (DeviceType type) |
| Lists compute devices of a specific type. | |
| template<typename OpType, DeviceType DT, typename ConfigT> | |
| std::shared_ptr< OpType > | makeOpInstance (IExecutionContext *ctx, const ConfigT &cfg) |
| Attempt to construct an OpType instance from a raw IExecutionContext*. | |
| std::string_view | Mila::Dnn::Compute::operationTypeToString (OperationType op) |
| template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TA, Dnn::TensorDataType TB, Dnn::TensorDataType TP = TA> | |
| void | Mila::Dnn::Compute::registerBinaryOpType (const std::string &opName) |
| Register a binary operation type with OperationRegistry using a common factory pattern. | |
| template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TA, Dnn::TensorDataType TB = TA, Dnn::TensorDataType TP = TA> | |
| void | Mila::Dnn::Compute::registerPairedOpType (const std::string &opName) |
| Register a paired operation type with OperationRegistry using a common factory pattern. | |
| template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TInput, Dnn::TensorDataType TPrecision = TInput> | |
| void | Mila::Dnn::Compute::registerUnaryOpType (std::string_view op_name) |
| Register a unary operation type with OperationRegistry using a common factory pattern. | |
| DeviceType | Mila::Dnn::Compute::toDeviceType (std::string_view device_type) |
| Converts a string to the corresponding DeviceType. | |
| template<DeviceType TDeviceType> | |
| ExecutionContext< TDeviceType > * | Mila::Dnn::Compute::validateExecutionContext_ (IExecutionContext *context, const std::string &op_name) |
Variables | |
| constexpr float | GELU_SCALING_FACTOR = 0.7978845608f |
| Scaling factor for GELU tanh approximation: sqrt(2/pi). | |
Typedef Documentation
◆ AdamWConfig
◆ CpuExecutionContext
|
export |
◆ CudaExecutionContext
|
export |
◆ HostMemoryResource
|
export |
Alias for CpuMemoryResource that represents host-accessible memory.
This alias provides a semantic name that describes the memory's accessibility characteristics rather than its implementation details. Use HostMemoryResource when you need memory that can be directly accessed from host (CPU) code.
- See also
- CpuMemoryResource
◆ OptimizerConfig
Enumeration Type Documentation
◆ DeviceType
|
exportstrong |
◆ OperationType
|
exportstrong |
Enumeration of all supported neural network operation types.
This enumeration defines the different types of operations that can be executed by the compute framework. Each operation type corresponds to a specific neural network function or layer.
| Enumerator | |
|---|---|
| CrossEntropyOp | Cross entropy loss operation (host-based; used by GPT reference implementation). |
| TokenEmbeddingOp | Token embedding operation. |
| LpeOp | Learned Positional Embedding operation for transformer architecture. |
| RopeOp | Rotary Position Embedding operation for transformer architecture. |
| FusedOp | Fused operation combining multiple operations for performance optimization. |
| LinearOp | Linear (fully connected/dense) layer operation. |
| GeluOp | Gaussian Error Linear Unit activation function. |
| SwigluOp | SwiGLU activation function. |
| LayerNormOp | Layer normalization operation. |
| RmsNormOp | RMS normalization operation. |
| MultiHeadAttentionOp | Multi-head attention operation (MHA) for transformers. |
| GroupedQueryAttentionOp | Grouped Query Attention (GQA). |
| ResidualOp | Residual connection operation. |
| SoftmaxOp | Softmax activation function. |
| DropoutOp | Dropout regularization operation. |
| SamplingOp | Device-side token sampling from logits. |
| SoftmaxCrossEntropyOp | WIP: Fused softmax + cross-entropy loss — targeted for Llama training. |
Function Documentation
◆ adamw_update()
template<typename Tp, typename Tg>
| void Mila::Dnn::Compute::adamw_update | ( | Tp * | params_memory, |
| float * | master_params_memory, | ||
| Tg * | grads_memory, | ||
| float * | m_memory, | ||
| float * | v_memory, | ||
| size_t | num_parameters, | ||
| ptrdiff_t | w_stride, | ||
| ptrdiff_t | g_stride, | ||
| ptrdiff_t | s_stride, | ||
| int | num_slices, | ||
| float | learning_rate, | ||
| float | beta1, | ||
| float | beta2, | ||
| int | t, | ||
| float | eps, | ||
| float | weight_decay, | ||
| float | grad_scale, | ||
| unsigned int | seed, | ||
| cudaStream_t | stream ) |
Here is the caller graph for this function:
◆ cast_context_() [1/2]
template<DeviceType TDeviceType>
|
nodiscardexportnoexcept |
◆ cast_context_() [2/2]
template<DeviceType TDeviceType>
|
nodiscardexportnoexcept |
Here is the caller graph for this function:
◆ createExecutionContext()
|
export |
Create execution context for specified device.
Factory function returning type-erased IExecutionContext. Hides device-specific implementation details from users.
- Examples
- /__w/Mila/Mila/Mila/Src/Dnn/Components/Activations/Gelu/Gelu.ixx, and /__w/Mila/Mila/Mila/Src/Dnn/Components/Normalization/Softmax.ixx.
Here is the call graph for this function:
Here is the caller graph for this function:
◆ cublasLtCheckStatus()
|
inlineexport |
Checks the status of a cuBLASLt operation and throws if an error occurred.
- Parameters
-
status The cuBLASLt error status code to check. location Source location information (automatically populated by default).
- Exceptions
-
CublasLtError if the status is not CUBLAS_STATUS_SUCCESS.
Here is the caller graph for this function:
◆ cuda_mha_forward_fp16()
| void Mila::Dnn::Compute::cuda_mha_forward_fp16 | ( | half * | Y, |
| half * | qkvr, | ||
| half * | att, | ||
| const half * | X, | ||
| int | B, | ||
| int | T, | ||
| int | C, | ||
| int | NH, | ||
| cudaStream_t | stream ) |
◆ cuda_mha_forward_fp32()
| void Mila::Dnn::Compute::cuda_mha_forward_fp32 | ( | float * | Y, |
| float * | qkvr, | ||
| float * | att, | ||
| const float * | X, | ||
| int | B, | ||
| int | T, | ||
| int | C, | ||
| int | NH, | ||
| cudaStream_t | stream ) |
◆ cuda_softmax_crossentropy_backward()
template<typename TPrecision>
| void Mila::Dnn::Compute::cuda_softmax_crossentropy_backward | ( | TPrecision * | dX, |
| const TPrecision * | dY_loss, | ||
| const TPrecision * | Y, | ||
| const int * | targets, | ||
| int | batch_size, | ||
| int | seq_len, | ||
| int | vocab_size, | ||
| cudaStream_t | stream ) |
◆ cuda_softmax_crossentropy_forward()
template<typename TPrecision>
| void Mila::Dnn::Compute::cuda_softmax_crossentropy_forward | ( | TPrecision * | Y_loss, |
| TPrecision * | Y, | ||
| const TPrecision * | X, | ||
| const int * | targets, | ||
| int | batch_size, | ||
| int | seq_len, | ||
| int | vocab_size, | ||
| cudaStream_t | stream ) |
◆ cudaCheckLastError()
|
export |
Checks the last CUDA error and throws if an error occurred.
- Parameters
-
location Source location information (automatically populated by default).
- Exceptions
-
CudaError if the last error is not cudaSuccess.
Here is the caller graph for this function:
◆ cudaCheckStatus()
|
export |
Checks the status of a CUDA operation and throws if an error occurred.
- Parameters
-
status The CUDA error status code to check. location Source location information (automatically populated by default).
- Exceptions
-
CudaError if the status is not cudaSuccess.
Here is the caller graph for this function:
◆ deviceTypeToString()
|
export |
Converts a DeviceType to its string representation.
- Parameters
-
device_type The device type to convert.
- Returns
- std::string The string representation of the device type (Device::Cpu() or "CUDA").
- Exceptions
-
std::invalid_argument If the device type is invalid.
- Examples
- /__w/Mila/Mila/Mila/Src/Dnn/Components/Activations/Gelu/Gelu.ixx, and /__w/Mila/Mila/Mila/Src/Dnn/Components/Normalization/Softmax.ixx.
Here is the caller graph for this function:
◆ getBestDevice()
|
export |
Gets the best DeviceId of a specific type based on performance characteristics.
For CUDA the function consults Cuda::getBestDeviceId(preferMemory) and returns the matching DeviceId when available. If no device of the requested type is found the function returns a DeviceId with the requested type and index == -1 to indicate "none".
- Parameters
-
type The device type to filter by (e.g., DeviceType::Cuda) preferMemory When true, prioritizes memory bandwidth over compute capability
- Returns
- DeviceId Best available device id for the requested type, or index == -1 if none
Here is the call graph for this function:
◆ getDeviceCount()
|
exportnoexcept |
Count instantiated devices of the given DeviceType.
Lightweight, thread-safe convenience helper intended for tests and diagnostics. Non-throwing; returns 0 on error.
Here is the call graph for this function:
◆ getRegisteredOperations()
template<DeviceType TDevice, TensorDataType TDataType>
|
export |
Templated helper returning registered operation names for a compile-time device and tensor data type.
Here is the call graph for this function:
Here is the caller graph for this function:
◆ isOperationRegistered()
template<DeviceType TDevice, TensorDataType TDataType>
|
export |
Templated helper that checks whether a named operation is registered for a compile-time device and tensor data type.
Here is the call graph for this function:
Here is the caller graph for this function:
◆ listDevicesByName()
|
export |
Lists all available compute devices by name.
- Returns
- std::vector<std::string> A list of device identifiers (e.g., "CUDA:0", "CPU:0").
Here is the call graph for this function:
◆ listDevicesByType()
|
export |
Lists compute devices of a specific type.
- Parameters
-
type The device type to filter by
- Returns
- std::vector<std::string> List of matching device identifiers
Here is the call graph for this function:
◆ makeOpInstance()
template<typename OpType, DeviceType DT, typename ConfigT>
| std::shared_ptr< OpType > Mila::Dnn::Compute::makeOpInstance | ( | IExecutionContext * | ctx, |
| const ConfigT & | cfg ) |
Attempt to construct an OpType instance from a raw IExecutionContext*.
Supports three constructor shapes (in priority order):
- OpType(IExecutionContext*, const ConfigT&)
- OpType(ExecutionContext<DT>*, const ConfigT&)
- OpType(std::shared_ptr<ExecutionContext<DT>>, const ConfigT&)
Throws std::invalid_argument when a required dynamic cast fails.
Here is the caller graph for this function:
◆ operationTypeToString()
|
export |
◆ registerBinaryOpType()
template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TA, Dnn::TensorDataType TB, Dnn::TensorDataType TP = TA>
|
export |
Register a binary operation type with OperationRegistry using a common factory pattern.
Template parameter ordering:
- TDataType : DeviceType
- OpType : Concrete operation class (must define using ConfigType = ...)
- TA : Input A precision
- TB : Input B precision
- TP : Compute/output precision (defaults to TA)
Here is the call graph for this function:
Here is the caller graph for this function:
◆ registerPairedOpType()
template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TA, Dnn::TensorDataType TB = TA, Dnn::TensorDataType TP = TA>
|
export |
Register a paired operation type with OperationRegistry using a common factory pattern.
Template parameter ordering:
- TDataType : DeviceType
- OpType : Concrete operation class (must define using ConfigType = ...)
- TA : Input A precision
- TB : Input B precision (defaults to TA)
- TP : Compute/output precision (defaults to TA)
Here is the call graph for this function:
Here is the caller graph for this function:
◆ registerUnaryOpType()
template<DeviceType TDataType, typename OpType, Dnn::TensorDataType TInput, Dnn::TensorDataType TPrecision = TInput>
|
export |
Register a unary operation type with OperationRegistry using a common factory pattern.
Template parameter ordering and names:
- TDataType : DeviceType
- OpType : Concrete operation class (must define using ConfigType = ...)
- TInput : Abstract input tensor precision
- TPrecision: Compute/output precision (defaults to TInput)
The factory lambda casts ComponentConfig -> OpType::ConfigType and forwards the IExecutionContext* + concrete config to the operation constructor via makeOpInstance.
Here is the call graph for this function:
Here is the caller graph for this function:
◆ toDeviceType()
|
export |
Converts a string to the corresponding DeviceType.
Performs case-insensitive matching to convert device type strings to the corresponding enum value.
- Parameters
-
device_type The string representation of the device type.
- Returns
- DeviceType The corresponding device type enum value.
- Exceptions
-
std::invalid_argument If the string does not represent a valid device type. Valid options are: Device::Cpu(), "CUDA", "AUTO".
Here is the caller graph for this function:
◆ validateExecutionContext_()
template<DeviceType TDeviceType>
|
export |
Here is the call graph for this function:
Here is the caller graph for this function:
Variable Documentation
◆ GELU_SCALING_FACTOR
|
constexpr |
Scaling factor for GELU tanh approximation: sqrt(2/pi).
Used in the tanh approximation formula: GELU(x) ~= 0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3)))
