Structural.h

Go to the documentation of this file.

 
#pragma once
 
#include <cuda_runtime.h>
 
namespace Mila::Dnn::Compute::Cuda
{
    void cuda_split3_fp32(
        const float* __restrict__ src,
        float* __restrict__ out_a,
        float* __restrict__ out_b,
        float* __restrict__ out_c,
        int src_rows,
        int dim_a, int dim_b, int dim_c,
        cudaStream_t stream );
 
    void cuda_split3_bf16(
        const __nv_bfloat16* __restrict__ src,
        __nv_bfloat16* __restrict__ out0,
        __nv_bfloat16* __restrict__ out1,
        __nv_bfloat16* __restrict__ out2,
        int rows,
        int D0, int D1, int D2,
        cudaStream_t stream );
}