CudaOps.h

Go to the documentation of this file.

#pragma once
 
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
 
namespace Mila::Dnn::Compute
{
    // REVIEW: These declaration should no longer be needed.
 
    // Attention functions
    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 );
 
    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 );
 
    // SoftmaxCrossEntropy functions
    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 );
 
    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 );
}