[WASI-NN] Add AUTO device support in ggml backend

Signed-off-by: dm4 <dm4@secondstate.io>
WasmEdge · Oct 12, 2023 · 7a36709 · 7a36709
1 parent db51a07
commit 7a36709
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/plugins/wasi_nn/types.h b/plugins/wasi_nn/types.h
@@ -22,7 +22,7 @@ enum class ErrNo : uint32_t {
 
 enum class TensorType : uint8_t { F16 = 0, F32 = 1, U8 = 2, I32 = 3 };
 
-enum class Device : uint32_t { CPU = 0, GPU = 1, TPU = 2 };
+enum class Device : uint32_t { CPU = 0, GPU = 1, TPU = 2, AUTO = 3 };
 
 enum class Backend : uint8_t {
   OpenVINO = 0,

diff --git a/plugins/wasi_nn/wasinnenv.cpp b/plugins/wasi_nn/wasinnenv.cpp
@@ -25,8 +25,10 @@ std::map<std::string_view, Backend> BackendMap = {
     {"autodetect"sv, Backend::Autodetect},
     {"ggml"sv, Backend::GGML}};
 
-std::map<std::string_view, Device> DeviceMap = {
-    {"cpu"sv, Device::CPU}, {"gpu"sv, Device::GPU}, {"tpu"sv, Device::TPU}};
+std::map<std::string_view, Device> DeviceMap = {{"cpu"sv, Device::CPU},
+                                                {"gpu"sv, Device::GPU},
+                                                {"tpu"sv, Device::TPU},
+                                                {"auto"sv, Device::AUTO}};
 
 bool load(const std::filesystem::path &Path, std::vector<uint8_t> &Data) {
   std::ifstream File(Path, std::ios::binary);
@@ -62,9 +64,9 @@ WasiNNEnvironment::WasiNNEnvironment() noexcept {
     std::vector<std::vector<uint8_t>> Models;
     Models.reserve(Paths.size());
     std::transform(Encode.begin(), Encode.end(), Encode.begin(),
-                   [](unsigned char c) { return std::tolower(c); });
+                   [](unsigned char C) { return std::tolower(C); });
     std::transform(Target.begin(), Target.end(), Target.begin(),
-                   [](unsigned char c) { return std::tolower(c); });
+                   [](unsigned char C) { return std::tolower(C); });
     auto Backend = BackendMap.find(Encode);
     auto Device = DeviceMap.find(Target);
     if (Backend != BackendMap.end() && Device != DeviceMap.end()) {