Function test

examples/cpp/swin/swin_example.cc:30–321

template<typename T>
void test(int version, int model_type, int window_size, int img_size, int batch)
{
    cudnnHandle_t    cudnn_handle;
    cublasHandle_t   cublas_handle;
    cublasLtHandle_t cublaslt_handle;
    cudaStream_t     stream = 0;
    checkCUDNN(cudnnCreate(&cudnn_handle));
    checkCUDNN(cudnnSetStream(cudnn_handle, stream));
    check_cuda_error(cublasCreate(&cublas_handle));
    check_cuda_error(cublasSetStream(cublas_handle, stream));
    check_cuda_error(cublasLtCreate(&cublaslt_handle));

    cublasAlgoMap* cublas_algo_map = new cublasAlgoMap(GEMM_CONFIG);

    std::mutex* cublas_wrapper_mutex = new std::mutex();

    cublasMMWrapper* cublas_wrapper =
        new cublasMMWrapper(cublas_handle, cublaslt_handle, stream, cublas_algo_map, cublas_wrapper_mutex, nullptr);

    if (std::is_same<T, half>::value) {
        cublas_wrapper->setFP16GemmConfig();
    }
#ifdef ENABLE_BF16
    else if (std::is_same<T, __nv_bfloat16>::value) {
        cublas_wrapper->setBF16GemmConfig();
    }
#endif
    else if (std::is_same<T, float>::value) {
        cublas_wrapper->setFP32GemmConfig();
    }

    int embed_dim;
    int shift_size;
    int depths[4], num_heads[4];

    if (model_type == 0) {  // tiny
        embed_dim    = 96;
        shift_size   = window_size / 2;
        depths[0]    = 2;
        depths[1]    = 2;
        depths[2]    = 6;
        depths[3]    = 2;
        num_heads[0] = 3;
        num_heads[1] = 6;
        num_heads[2] = 12;
        num_heads[3] = 24;
    }
    else if (model_type == 1) {  // small
        embed_dim    = 96;
        shift_size   = window_size / 2;
        depths[0]    = 2;
        depths[1]    = 2;
        depths[2]    = 18;
        depths[3]    = 2;
        num_heads[0] = 3;
        num_heads[1] = 6;
        num_heads[2] = 12;
        num_heads[3] = 24;