Add onnxruntime gpu for cmake (#153)

* add onnxruntime gpu for cmake

* fix clang

* fix typo

* cpplint
This commit is contained in:
Yuekai Zhang
2023-05-12 22:30:47 +08:00
committed by GitHub
parent cea718e3d8
commit b8fbf8e5ce
4 changed files with 90 additions and 18 deletions

View File

@@ -78,6 +78,12 @@ target_link_libraries(sherpa-onnx-core
kaldi-native-fbank-core
)
if(SHERPA_ONNX_ENABLE_GPU)
target_link_libraries(sherpa-onnx-core
onnxruntime_providers_cuda
)
endif()
if(SHERPA_ONNX_ENABLE_CHECK)
target_compile_definitions(sherpa-onnx-core PUBLIC SHERPA_ONNX_ENABLE_CHECK=1)

View File

@@ -4,8 +4,10 @@
#include "sherpa-onnx/csrc/session.h"
#include <algorithm>
#include <string>
#include <utility>
#include <vector>
#include "sherpa-onnx/csrc/macros.h"
#include "sherpa-onnx/csrc/provider.h"
@@ -27,10 +29,20 @@ static Ort::SessionOptions GetSessionOptionsImpl(int32_t num_threads,
case Provider::kCPU:
break; // nothing to do for the CPU provider
case Provider::kCUDA: {
OrtCUDAProviderOptions options;
options.device_id = 0;
// set more options on need
sess_opts.AppendExecutionProvider_CUDA(options);
std::vector<std::string> available_providers =
Ort::GetAvailableProviders();
if (std::find(available_providers.begin(), available_providers.end(),
"CUDAExecutionProvider") != available_providers.end()) {
// The CUDA provider is available, proceed with setting the options
OrtCUDAProviderOptions options;
options.device_id = 0;
// set more options on need
sess_opts.AppendExecutionProvider_CUDA(options);
} else {
SHERPA_ONNX_LOGE(
"Please compile with -DSHERPA_ONNX_ENABLE_GPU=ON. Fallback to "
"cpu!");
}
break;
}
case Provider::kCoreML: {