Method exec

dnn/src/cuda/pooling/algo.cpp:112–146

}

void PoolingForwardImpl::AlgoCUDNN::exec(const ExecArgs& args) const {
    TensorND src = *args.src_tensor;
    TensorND dst = *args.dst_tensor;
    auto wsb = get_workspace_bundle(args.workspace.raw_ptr, args);
    auto ctypecvt = CompTypeCvter<dtype::BFloat16, dtype::Float32>(
            concrete_handle(args.handle), &wsb);
    if (args.layout_src->dtype.enumv() == DTypeTrait<dtype::BFloat16>::enumv) {
        ctypecvt.src_to_comp_type(*args.src_tensor, src)
                .src_to_comp_type(*args.dst_tensor, dst);
    }
    {
        dt_float32 alpha = 1.0f, beta = 0.0f;
        TensorDesc src_desc, dst_desc;
        src_desc.set(src.layout, args.opr->param().format);
        dst_desc.set(dst.layout, args.opr->param().format);

        cudnnPoolingMode_t mode;
        init_mode(args, mode);

        cudnnPoolingDescriptor_t cudnn_desc;
        cudnn_check(cudnnCreatePoolingDescriptor(&cudnn_desc));
        cudnn_check(cudnnSetPooling2dDescriptor(
                cudnn_desc, mode, CUDNN_NOT_PROPAGATE_NAN, args.opr->param().window_h,
                args.opr->param().window_w, args.opr->param().pad_h,
                args.opr->param().pad_w, args.opr->param().stride_h,
                args.opr->param().stride_w));
        cudnn_check(cudnnPoolingForward(
                args.handle->cudnn_handle(), cudnn_desc, &alpha, src_desc.desc,
                src.raw_ptr(), &beta, dst_desc.desc, dst.raw_ptr()));
        cudnn_check(cudnnDestroyPoolingDescriptor(cudnn_desc));
    }
    if (args.layout_src->dtype.enumv() == DTypeTrait<dtype::BFloat16>::enumv) {
        ctypecvt.comp_to_dst_type(dst, *args.dst_tensor);
    }
}

#if CUDNN_VERSION >= 6000
bool PoolingForwardImpl::AlgoCUDNNMAXDETERMINISTIC::is_available(