ct2rs 0.9.19

Rust bindings for OpenNMT/CTranslate2
Documentation 
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

#include "ctranslate2/allocator.h"

#ifdef _WIN32
#  include <malloc.h>
#else
#  include <cstdlib>
#endif

#ifdef CT2_WITH_MKL
#  include <mkl.h>
#endif

namespace ctranslate2 {
  namespace cpu {

    class AlignedAllocator : public Allocator {
    public:
      AlignedAllocator(size_t alignment)
        : _alignment(alignment)
      {
      }

      void* allocate(size_t size, int) override {
        void* ptr = nullptr;
#ifdef _WIN32
        ptr = _aligned_malloc(size, _alignment);
#else
        if (posix_memalign(&ptr, _alignment, size) != 0)
          ptr = nullptr;
#endif
        if (!ptr)
          throw std::runtime_error("aligned_alloc: failed to allocate memory");
        return ptr;
      }

      void free(void* ptr, int) override {
#ifdef _WIN32
        _aligned_free(ptr);
#else
        std::free(ptr);
#endif
      }

    private:
      size_t _alignment;
    };

#ifdef CT2_WITH_MKL
    class MklAllocator : public Allocator {
    public:
      MklAllocator(size_t alignment)
        : _alignment(alignment)
      {
      }

      void* allocate(size_t size, int) override {
        void* ptr = mkl_malloc(size, _alignment);
        if (!ptr)
          throw std::runtime_error("mkl_malloc: failed to allocate memory");
        return ptr;
      }

      void free(void* ptr, int) override {
        mkl_free(ptr);
      }

      void clear_cache() override {
        mkl_free_buffers();
      }

    private:
      size_t _alignment;
    };
#endif

  }

  template<>
  Allocator& get_allocator<Device::CPU>() {
    constexpr size_t alignment = 64;
#ifdef CT2_WITH_MKL
    static cpu::MklAllocator allocator(alignment);
#else
    static cpu::AlignedAllocator allocator(alignment);
#endif
    return allocator;
  }

}