Cleanup utils.h (OpenNMT#756)

* Move some tokenization utilities out of TranslatorPool implementation

* Move get_gpu_count to devices.h

* Move random utilities to separate file

* Move get_preferred_size_multiple to types utils

* Move mayiuse_* functions to types utils

* Move env reading utilities to separate file

* Move thread affinity function to thread pool implementation

* Cleanup include

* Fix include

CMakeLists.txt

@@ -86,6 +86,7 @@ set(SOURCES
   src/cpu/primitives.cc
   src/decoding.cc
   src/devices.cc
+  src/env.cc
   src/layers/attention.cc
   src/layers/common.cc
   src/layers/decoder.cc
@@ -131,6 +132,7 @@ set(SOURCES
   src/ops/transpose.cc
   src/padder.cc
   src/profiler.cc
+  src/random.cc
   src/sampling.cc
   src/scoring.cc
   src/storage_view.cc

cli/translate.cc

@@ -5,6 +5,7 @@
 
 #include <ctranslate2/translator_pool.h>
 #include <ctranslate2/utils.h>
+#include <ctranslate2/random.h>
 #include <ctranslate2/devices.h>
 #include <ctranslate2/profiler.h>
 

include/ctranslate2/batch_reader.h

@@ -39,6 +39,27 @@ namespace ctranslate2 {
     }
   };
 
+  // Helper class to read tokens from a text stream.
+  template <typename Tokenizer>
+  class TextLineReader {
+  public:
+    TextLineReader(Tokenizer& tokenizer)
+      : _tokenizer(tokenizer)
+    {
+    }
+
+    bool operator()(std::istream& in, std::vector<std::string>& tokens) {
+      std::string line;
+      if (!std::getline(in, line))
+        return false;
+      tokens = _tokenizer(line);
+      return true;
+    }
+
+  private:
+    Tokenizer& _tokenizer;
+  };
+
 
   // Base class to produce batches.
   class BatchReader {

include/ctranslate2/devices.h

@@ -16,6 +16,8 @@ namespace ctranslate2 {
   int get_device_index(Device device);
   void set_device_index(Device device, int index);
 
+  int get_gpu_count();
+
   void synchronize_device(Device device, int index);
 
   class ScopedDeviceSetter {

include/ctranslate2/random.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <random>
+
+namespace ctranslate2 {
+
+  void set_random_seed(const unsigned int seed);
+  unsigned int get_random_seed();
+  std::mt19937& get_random_generator();
+
+}

include/ctranslate2/translator_pool.h

@@ -253,8 +253,8 @@ namespace ctranslate2 {
                                            const bool with_scores = false) {
       TranslationStats stats;
 
-      TokensReader<SourceTokenizer> source_reader(source_tokenizer);
-      TokensReader<TargetTokenizer> target_reader(target_tokenizer);
+      TextLineReader<SourceTokenizer> source_reader(source_tokenizer);
+      TextLineReader<TargetTokenizer> target_reader(target_tokenizer);
 
       auto writer = [&detokenizer, &stats, &with_scores](std::ostream& out,
                                                          const TranslationResult& result) {
@@ -373,8 +373,8 @@ namespace ctranslate2 {
                                          const size_t read_batch_size = 0,
                                          const BatchType batch_type = BatchType::Examples,
                                          bool with_token_scores = false) {
-      TokensReader<SourceTokenizer> source_reader(source_tokenizer);
-      TokensReader<TargetTokenizer> target_reader(target_tokenizer);
+      TextLineReader<SourceTokenizer> source_reader(source_tokenizer);
+      TextLineReader<TargetTokenizer> target_reader(target_tokenizer);
       TranslationStats stats;
 
       auto writer = [&target_detokenizer, &stats, with_token_scores](std::ostream& out,
@@ -672,26 +672,6 @@ namespace ctranslate2 {
                             const long max_queued_batches);
 
     std::unique_ptr<ThreadPool> _thread_pool;
-
-    template <typename Tokenizer>
-    class TokensReader {
-    public:
-      TokensReader(Tokenizer& tokenizer)
-        : _tokenizer(tokenizer)
-      {
-      }
-
-      bool operator()(std::istream& in, std::vector<std::string>& tokens) {
-        std::string line;
-        if (!std::getline(in, line))
-          return false;
-        tokens = _tokenizer(line);
-        return true;
-      }
-
-    private:
-      Tokenizer& _tokenizer;
-    };
   };
 
 }

include/ctranslate2/types.h

@@ -33,11 +33,15 @@ namespace ctranslate2 {
   ComputeType str_to_compute_type(const std::string& compute_type);
   std::string compute_type_to_str(const ComputeType compute_type);
 
+  bool mayiuse_float16(const Device device, const int device_index = 0);
+  bool mayiuse_int16(const Device device, const int device_index = 0);
+  bool mayiuse_int8(const Device device, const int device_index = 0);
+
   // Returns the final compute type based on model weights and device information.
   ComputeType resolve_compute_type(const ComputeType requested_compute_type,
                                    const ComputeType model_compute_type,
                                    const Device device,
-                                   const int device_index,
+                                   const int device_index = 0,
                                    const bool enable_fallback = false);
 
   // Gets the weights data type for the given compute type.
@@ -46,4 +50,9 @@ namespace ctranslate2 {
   // Gets the default floating point type for the given compute type.
   DataType get_default_float_type(const ComputeType compute_type);
 
+  // Gets the preferred size multiple for this compute type.
+  dim_t get_preferred_size_multiple(const ComputeType compute_type,
+                                    const Device device,
+                                    const int device_index = 0);
+
 }

include/ctranslate2/utils.h

@@ -1,40 +1,24 @@
 #pragma once
 
-#include <fstream>
-#include <random>
+#include <stdexcept>
 #include <string>
-#include <thread>
 #include <vector>
 
-#include "devices.h"
-#include "types.h"
-
 namespace ctranslate2 {
 
   bool string_to_bool(const std::string& str);
-  std::string read_string_from_env(const char* var, const std::string& default_value = "");
-  bool read_bool_from_env(const char* var, const bool default_value = false);
-  int read_int_from_env(const char* var, const int default_value = 0);
-
-  // Check feature support.
-  bool mayiuse_float16(Device device, int device_index = 0);
-  bool mayiuse_int16(Device device, int device_index = 0);
-  bool mayiuse_int8(Device device, int device_index = 0);
-  dim_t get_preferred_size_multiple(ComputeType compute_type,
-                                    Device device,
-                                    int device_index = 0);
-
-  int get_gpu_count();
 
   void set_num_threads(size_t num_threads);
-  void set_thread_affinity(std::thread& thread, int index);
 
   bool ends_with(const std::string& str, const std::string& suffix);
   bool starts_with(const std::string& str, const std::string& prefix);
 
   std::vector<std::string> split_string(const std::string& str, char delimiter);
   std::vector<std::string> split_string(const std::string& str, const std::string& delimiter);
 
+  std::vector<std::string> split_tokens(const std::string& text);
+  std::string join_tokens(const std::vector<std::string>& tokens);
+
   template <typename Stream>
   Stream open_file(const std::string& path) {
     Stream stream(path);
@@ -61,10 +45,6 @@ namespace ctranslate2 {
     }
   }
 
-  void set_random_seed(const unsigned int seed);
-  unsigned int get_random_seed();
-  std::mt19937& get_random_generator();
-
 #ifdef NDEBUG
 #  define THROW_EXCEPTION(EXCEPTION, MESSAGE) throw EXCEPTION(MESSAGE)
 #else

python/translator.cc

@@ -9,6 +9,7 @@
 #include <pybind11/stl.h>
 
 #include <ctranslate2/translator_pool.h>
+#include <ctranslate2/random.h>
 
 namespace py = pybind11;
 

src/cpu/backend.cc

@@ -6,6 +6,7 @@
 
 #include "ctranslate2/utils.h"
 #include "cpu_info.h"
+#include "env.h"
 
 namespace ctranslate2 {
   namespace cpu {

src/cpu/cpu_isa.cc

@@ -1,7 +1,9 @@
 #include "cpu_isa.h"
 
-#include "ctranslate2/utils.h"
+#include <stdexcept>
+
 #include "cpu_info.h"
+#include "env.h"
 
 namespace ctranslate2 {
   namespace cpu {

src/cuda/allocator.cc

@@ -5,6 +5,7 @@
 
 #include "ctranslate2/utils.h"
 #include "cuda/utils.h"
+#include "env.h"
 
 #include <cuda.h>
 #include <cub/util_allocator.cuh>

src/cuda/cublas_stub.cc

@@ -15,7 +15,7 @@
 
 #include <spdlog/spdlog.h>
 
-#include "ctranslate2/utils.h"
+#include "env.h"
 
 namespace ctranslate2 {
 

src/devices.cc

@@ -76,6 +76,14 @@ namespace ctranslate2 {
     DEVICE_DISPATCH(device, set_device_index<D>(index));
   }
 
+  int get_gpu_count() {
+#ifdef CT2_WITH_CUDA
+    return cuda::get_gpu_count();
+#else
+    return 0;
+#endif
+  }
+
   void synchronize_device(Device device, int index) {
 #ifdef CT2_WITH_CUDA
     if (device == Device::CUDA) {

src/env.cc

@@ -0,0 +1,27 @@
+#include "env.h"
+
+#include <cstdlib>
+
+#include "ctranslate2/utils.h"
+
+namespace ctranslate2 {
+
+  std::string read_string_from_env(const char* var, const std::string& default_value) {
+    const char* value = std::getenv(var);
+    if (!value)
+      return default_value;
+    return value;
+  }
+
+  bool read_bool_from_env(const char* var, const bool default_value) {
+    return string_to_bool(read_string_from_env(var, default_value ? "1" : "0"));
+  }
+
+  int read_int_from_env(const char* var, const int default_value) {
+    const std::string value = read_string_from_env(var);
+    if (value.empty())
+      return default_value;
+    return std::stoi(value);
+  }
+
+}

src/env.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <string>
+
+namespace ctranslate2 {
+
+  std::string read_string_from_env(const char* var, const std::string& default_value = "");
+  bool read_bool_from_env(const char* var, const bool default_value = false);
+  int read_int_from_env(const char* var, const int default_value = 0);
+
+}

src/layers/transformer.cc

@@ -1,5 +1,7 @@
 #include "ctranslate2/layers/transformer.h"
 
+#include <cmath>
+
 namespace ctranslate2 {
   namespace layers {
 

src/ops/multinomial_cpu.cc

@@ -1,6 +1,6 @@
 #include "ctranslate2/ops/multinomial.h"
 
-#include "ctranslate2/utils.h"
+#include "ctranslate2/random.h"
 
 #include "type_dispatch.h"
 

src/ops/multinomial_gpu.cu

@@ -6,7 +6,7 @@
 #include <cub/block/block_scan.cuh>
 #include <curand_kernel.h>
 
-#include "ctranslate2/utils.h"
+#include "ctranslate2/random.h"
 
 #include "type_dispatch.h"
 #include "cuda/helpers.h"

src/ops/quantize_cpu.cc

@@ -1,5 +1,7 @@
 #include "ctranslate2/ops/quantize.h"
 
+#include <cmath>
+
 #include "cpu/kernels.h"
 #include "cpu/parallel.h"
 

src/random.cc

@@ -0,0 +1,23 @@
+#include "ctranslate2/random.h"
+
+#include <atomic>
+
+namespace ctranslate2 {
+
+  constexpr unsigned int default_seed = static_cast<unsigned int>(-1);
+  static std::atomic<unsigned int> g_seed(default_seed);
+
+  void set_random_seed(const unsigned int seed) {
+    g_seed = seed;
+  }
+
+  unsigned int get_random_seed() {
+    return g_seed == default_seed ? std::random_device{}() : g_seed.load();
+  }
+
+  std::mt19937& get_random_generator() {
+    static thread_local std::mt19937 generator(get_random_seed());
+    return generator;
+  }
+
+}

src/thread_pool.cc

@@ -67,6 +67,24 @@ namespace ctranslate2 {
   }
 
 
+  static void set_thread_affinity(std::thread& thread, int index) {
+#if !defined(__linux__) || defined(_OPENMP)
+    (void)thread;
+    (void)index;
+    throw std::runtime_error("Setting thread affinity is only supported in Linux binaries built "
+                             "with -DOPENMP_RUNTIME=NONE");
+#else
+    cpu_set_t cpuset;
+    CPU_ZERO(&cpuset);
+    CPU_SET(index, &cpuset);
+    const int status = pthread_setaffinity_np(thread.native_handle(), sizeof (cpu_set_t), &cpuset);
+    if (status != 0) {
+      throw std::runtime_error("Error calling pthread_setaffinity_np: "
+                               + std::to_string(status));
+    }
+#endif
+  }
+
   void Worker::start(JobQueue& job_queue, int thread_affinity) {
     _thread = std::thread(&Worker::run, this, std::ref(job_queue));
     if (thread_affinity >= 0)