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

#include "ctranslate2/ops/gather.h"

#include "cpu/parallel.h"
#include "type_dispatch.h"

namespace ctranslate2 {
  namespace ops {

    template <Device D, typename T>
    void Gather::compute(const StorageView& data,
                         const StorageView& input,
                         const dim_t axis,
                         const dim_t batch_dims,
                         StorageView& output) const {
      const auto* indices = input.data<int32_t>();
      const T* src = data.data<T>();
      T* dst = output.data<T>();

      if (axis == batch_dims) {
        const dim_t copy_size = data.stride(axis);
        const dim_t batch_stride = axis > 0 ? data.stride(axis - 1) : data.size();
        const dim_t batch_size = data.size() / batch_stride;
        const dim_t num_indices = input.size();
        const dim_t num_indices_per_batch = num_indices / batch_size;

        const dim_t grain_size = cpu::get_minimum_batch_copies_per_thread<T>(copy_size);
        cpu::parallel_for(0, num_indices, grain_size, [&](dim_t begin, dim_t end) {
          for (dim_t i = begin; i < end; ++i) {
            const dim_t batch_index = i / num_indices_per_batch;
            const dim_t read_index = indices[i];
            primitives<Device::CPU>::copy(src + batch_index * batch_stride + read_index * copy_size,
                                          dst + i * copy_size,
                                          copy_size);
          }
        });

      } else {
        throw std::invalid_argument("Gather only supports indexing the first non batch dimension");
      }
    }

#define DECLARE_IMPL(T)                                         \
    template void                                               \
    Gather::compute<Device::CPU, T>(const StorageView& data,    \
                                    const StorageView& input,   \
                                    const dim_t axis,           \
                                    const dim_t batch_dims,     \
                                    StorageView& output) const;

    DECLARE_ALL_TYPES(DECLARE_IMPL)

  }
}