feats: add ROCM related ⬆️

Joker2770 · Aug 16, 2024 · 43da8c5 · 43da8c5
1 parent e08f1df
commit 43da8c5
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -33,6 +33,12 @@ if(ENABLE_TENSORRT)
     add_definitions(-DUSE_TENSORRT)
 endif(ENABLE_TENSORRT)
 
+option(ENABLE_ROCM OFF)
+if(ENABLE_ROCM)
+    message("build project with ROCM")
+    add_definitions(-DUSE_ROCM)
+endif(ENABLE_ROCM)
+
 include_directories("${ONNXRUNTIME_ROOTDIR}/include"
                 "${CMAKE_SOURCE_DIR}/src"
                 "${CMAKE_SOURCE_DIR}/src/pbrain-Z2I/toml11"

diff --git a/src/common.h b/src/common.h
@@ -30,6 +30,7 @@ SOFTWARE.
 // #define USE_CUDA
 // #define USE_OPENVINO
 // #define USE_TENSORRT
+// #define USE_ROCM
 
 #define CHANNEL_SIZE 3
 

diff --git a/src/onnx.cpp b/src/onnx.cpp
@@ -105,6 +105,14 @@ NeuralNetwork::NeuralNetwork(const std::string &model_path, const unsigned int b
   CheckStatus(g_ort, g_ort->SessionOptionsAppendExecutionProvider_TensorRT(session_options, &TensorRT_Options));
 #endif
 
+#ifdef USE_ROCM
+  const OrtApiBase *ptr_api_base = OrtGetApiBase();
+  const OrtApi *g_ort = ptr_api_base->GetApi(ORT_API_VERSION);
+  OrtROCMProviderOptions ROCM_Options;
+  ROCM_Options.device_id = 0;
+  CheckStatus(g_ort, g_ort->SessionOptionsAppendExecutionProvider_ROCM(session_options, &ROCM_Options));
+#endif
+
 #ifdef _WIN32
   // std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
   std::wstring wstr(model_path.length(), L' ');

diff --git a/xmake.lua b/xmake.lua
@@ -9,13 +9,21 @@ target("Z2I")
     add_files("src/*.cpp")
     add_packages("onnxruntime")
     add_rpathdirs("@loader_path", "@loader_path/lib", "@executable_path", "@executable_path/lib")
+    -- add_defines("USE_CUDA")
+    -- add_defines("USE_OPENVINO")
+    -- add_defines("USE_TENSORRT")
+    -- add_defines("USE_ROCM")
 
 target("pbrain-Z2I")
     set_kind("binary")
     add_files("src/pbrain-Z2I/*.cpp")
     add_deps("Z2I")
     add_packages("onnxruntime")
     add_rpathdirs("@loader_path", "@loader_path/lib", "@executable_path", "@executable_path/lib")
+    -- add_defines("USE_CUDA")
+    -- add_defines("USE_OPENVINO")
+    -- add_defines("USE_TENSORRT")
+    -- add_defines("USE_ROCM")
     after_build(function (target)
         os.cp("$(scriptdir)/config/*.toml", target:targetdir())
     end)
@@ -26,3 +34,7 @@ target("train_and_eval")
     add_deps("Z2I")
     add_packages("onnxruntime")
     add_rpathdirs("@loader_path", "@loader_path/lib", "@executable_path", "@executable_path/lib")
+    -- add_defines("USE_CUDA")
+    -- add_defines("USE_OPENVINO")
+    -- add_defines("USE_TENSORRT")
+    -- add_defines("USE_ROCM")