Add support for CUDAExecutionProvider to Onnx::Session (#137)

Author: SimBe195

src/Onnx/Session.cc

@@ -5,6 +5,10 @@
 
 #include <chrono>
 
+#ifdef MODULE_CUDA
+#include <cuda_runtime.h>
+#endif
+
 #include "Util.hh"
 
 namespace Onnx {
@@ -21,6 +25,14 @@ const Core::ParameterInt Session::paramInterOpNumThreads("inter-op-num-threads",
                                                          "number of threads to use between ops",
                                                          1);
 
+const Core::Choice Session::executionProviderChoice(
+        "cpu", ExecutionProviderType::cpu,
+        "cuda", ExecutionProviderType::cuda,
+        Core::Choice::endMark());
+
+const Core::ParameterChoice Session::paramExecutionProviderType(
+        "execution-provider-type", &Session::executionProviderChoice, "type of execution provider", ExecutionProviderType::cpu);
+
 Session::Session(Core::Configuration const& config)
         : Precursor(config),
           file_(paramFile(config)),
@@ -34,6 +46,37 @@ Session::Session(Core::Configuration const& config)
     Ort::SessionOptions session_opts;
     session_opts.SetIntraOpNumThreads(intraOpNumThreads_);
     session_opts.SetInterOpNumThreads(interOpNumThreads_);
+
+    auto providers = Ort::GetAvailableProviders();
+    switch (paramExecutionProviderType(config)) {
+        case ExecutionProviderType::cpu: {
+            if (std::find(providers.begin(), providers.end(), "CPUExecutionProvider") == providers.end()) {
+                error() << "Requested CPU execution provider for ONNX session but it is not available.";
+            }
+            break;
+        }
+        case ExecutionProviderType::cuda: {
+            if (std::find(providers.begin(), providers.end(), "CUDAExecutionProvider") == providers.end()) {
+                error() << "Requested CUDA execution provider for ONNX session but it is not available.";
+            }
+#ifdef MODULE_CUDA
+            int deviceCount = 0;
+            if (cudaGetDeviceCount(&deviceCount) != cudaSuccess or deviceCount == 0) {
+                error() << "Requested CUDA execution provider but no CUDA device was found.";
+            }
+            OrtCUDAProviderOptionsV2* cuda_opts = nullptr;
+            Ort::ThrowOnError(Ort::GetApi().CreateCUDAProviderOptions(&cuda_opts));
+            session_opts.AppendExecutionProvider_CUDA_V2(*cuda_opts);
+            Ort::GetApi().ReleaseCUDAProviderOptions(cuda_opts);
+            break;
+#else
+            error() << "Requested CUDA execution provider but RASR was not compiled with MODULE_CUDA which is required for it.";
+#endif
+        }
+        default:
+            error() << "Execution provider for ONNX session not known.";
+    }
+
     session_ = Ort::Session(env_, file_.c_str(), session_opts);
 
     size_t num_inputs  = session_.GetInputCount();

src/Onnx/Session.hh

@@ -14,6 +14,11 @@
 
 namespace Onnx {
 
+enum ExecutionProviderType {
+    cpu,
+    cuda
+};
+
 class Session : public Core::Component {
 public:
     using Precursor = Core::Component;
@@ -22,6 +27,9 @@ public:
     static const Core::ParameterInt    paramIntraOpNumThreads;
     static const Core::ParameterInt    paramInterOpNumThreads;
 
+    static const Core::Choice          executionProviderChoice;
+    static const Core::ParameterChoice paramExecutionProviderType;
+
     Session(Core::Configuration const& config);
     virtual ~Session() = default;