CTranslate2/src/cpu/cpu_isa.cc at 8e8f597473e8633877ba7e71cda9711b6a1bdab6 · OpenNMT/CTranslate2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#include "cpu_isa.h"

#include <stdexcept>

#include "cpu_info.h"
#include "env.h"

namespace ctranslate2 {
  namespace cpu {

    static CpuIsa try_isa(const std::string& name, CpuIsa cpu_isa, bool supported) {
#ifdef CT2_WITH_CPU_DISPATCH
      if (!supported) {
        throw std::invalid_argument("The CPU does not support " + name);
      } else {
        return cpu_isa;
      }
#else
      (void)cpu_isa;
      (void)supported;
      throw std::invalid_argument("This CTranslate2 binary was not compiled with "
                                  + name + " support");
#endif
    }

    std::string isa_to_str(CpuIsa isa) {
      switch (isa) {
#if defined(CT2_X86_BUILD)
      case CpuIsa::AVX:
        return "AVX";
      case CpuIsa::AVX2:
        return "AVX2";
      case CpuIsa::AVX512:
        return "AVX512";
#elif defined(CT2_ARM64_BUILD)
      case CpuIsa::NEON:
        return "NEON";
#elif defined(CT2_WITH_RVV)
      case CpuIsa::RVV:
        return "RVV";
#endif
      default:
        return "GENERIC";
      }
    }

    static CpuIsa init_isa() {
      const std::string env_isa = read_string_from_env("CT2_FORCE_CPU_ISA");
      if (!env_isa.empty()) {
#if defined(CT2_X86_BUILD)
        if (env_isa == "AVX512")
          return try_isa(env_isa, CpuIsa::AVX512, cpu_supports_avx512());
        if (env_isa == "AVX2")
          return try_isa(env_isa, CpuIsa::AVX2, cpu_supports_avx2());
        if (env_isa == "AVX")
          return try_isa(env_isa, CpuIsa::AVX, cpu_supports_avx());
#elif defined(CT2_ARM64_BUILD)
        if (env_isa == "NEON")
          return try_isa(env_isa, CpuIsa::NEON, cpu_supports_neon());
#elif defined(CT2_WITH_RVV)
        if (env_isa == "RVV")
          return try_isa(env_isa, CpuIsa::RVV, cpu_supports_rvv());
#endif
        if (env_isa == "GENERIC")
          return CpuIsa::GENERIC;

        throw std::invalid_argument("Invalid CPU ISA: " + env_isa);
      }

#ifdef CT2_WITH_CPU_DISPATCH
#  if defined(CT2_X86_BUILD)
      // Note that AVX512 can only be enabled with the environment variable at this time.
      if (cpu_supports_avx2())
        return CpuIsa::AVX2;
      if (cpu_supports_avx())
        return CpuIsa::AVX;
#  elif defined(CT2_ARM64_BUILD)
      if (cpu_supports_neon())
        return CpuIsa::NEON;
#  elif defined(CT2_WITH_RVV)
      if (cpu_supports_rvv())
        return CpuIsa::RVV;
#  endif
#endif

      return CpuIsa::GENERIC;
    }

    CpuIsa get_cpu_isa() {
      static const CpuIsa cpu_isa = init_isa();
      return cpu_isa;
    }

  }
}