From 7a36709fbfc0dcbbf3ea7d55e677568c891ec620 Mon Sep 17 00:00:00 2001
From: dm4 <dm4@secondstate.io>
Date: Thu, 5 Oct 2023 10:43:51 +0800
Subject: [PATCH] [WASI-NN] Add AUTO device support in ggml backend

Signed-off-by: dm4 <dm4@secondstate.io>
---
 plugins/wasi_nn/types.h       |  2 +-
 plugins/wasi_nn/wasinnenv.cpp | 10 ++++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/plugins/wasi_nn/types.h b/plugins/wasi_nn/types.h
index 6e2a812861ef..98eca0416a87 100644
--- 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
index af567393263c..6c90ae54b125 100644
--- 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()) {