Use --device-id option

zakki · May 4, 2018 · 14cbff2 · 14cbff2
1 parent 133bfcb
commit 14cbff2
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/src/UctSearch.cpp b/src/UctSearch.cpp
@@ -416,6 +416,12 @@ SetDeviceId(const int id)
   device_id = id;
 }
 
+int
+GetDeviceId()
+{
+  return device_id;
+}
+
 void
 SetNoExpand(bool flag)
 {

diff --git a/src/UctSearch.h b/src/UctSearch.h
@@ -239,6 +239,8 @@ void SetUseNN(bool flag);
 
 void SetDeviceId( const int id );
 
+int GetDeviceId();
+
 // Policy networkの手を打つ
 int PolicyNetworkGenmove(game_info_t *game, int color);
 

diff --git a/src/lz/Leela.cpp b/src/lz/Leela.cpp
@@ -340,6 +340,7 @@ void benchmark(GameState& game) {
 #endif
 
 extern char uct_params_path[1024];
+extern int GetDeviceId();
 
 int InitializeLeela() {
   auto input = std::string{};
@@ -366,6 +367,10 @@ int InitializeLeela() {
   cfg_weightsfile = uct_params_path;
   cfg_weightsfile += "/lz.bin";
 
+  int device_id = GetDeviceId();
+  if (device_id >= 0)
+    cfg_gpus.push_back(device_id);
+
   init_global_objects();
 
 #if 0