1 #ifndef CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_ 2 #define CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_ 6 #include "caffe2/core/common_omp.h" 7 #include "caffe2/core/context.h" 8 #include "caffe2/core/logging.h" 9 #include "caffe2/core/operator.h" 10 #include "caffe2/core/types.h" 11 #include "caffe2/utils/math.h" 17 template <
class Context>
20 USE_OPERATOR_CONTEXT_FUNCTIONS;
23 lengths_(OperatorBase::GetRepeatedArgument<int>(
"lengths")) {
24 CAFFE_ENFORCE_GT(lengths_.size(), 0,
"There has to be at least one length");
25 for (
auto length : lengths_) {
26 CAFFE_ENFORCE_GT(length, 0,
"Each length should be positive");
30 bool RunOnDevice()
override {
32 this, OperatorBase::Input<TensorCPU>(RANGES));
35 template <
typename Index>
36 bool DoRunWithType() {
37 auto& data = Input(DATA);
38 auto& ranges = Input(RANGES);
39 CAFFE_ENFORCE_EQ(data.ndim(), 1,
"Data has to be 1-D");
40 CAFFE_ENFORCE_EQ(ranges.ndim(), 3,
"Ranges has to be 3-D");
41 if (InputSize() == 3) {
42 auto& key = Input(KEY);
43 CAFFE_ENFORCE_EQ(key.ndim(), 1,
"Key has to be 1-D");
45 key.meta().template Match<int64_t>(),
"Key has to be type int64_t");
50 "Nummber of ranges should match number of lengths");
54 "Nummber of ranges should match number of outputs");
56 ranges.dim(2), 2,
"Ranges last dimension should be of size 2");
58 auto* rawData =
static_cast<const char*
>(data.raw_data());
59 auto* rangesData = ranges.template data<Index>();
60 int rangesDataOffset = 0;
61 auto itemsize = data.meta().itemsize();
63 auto batchSize = ranges.dim(0);
64 vector<TIndex> outputDims{batchSize, 0};
65 vector<char*> outputRawData;
66 for (
int i = 0; i < OutputSize(); ++i) {
67 auto* output = Output(i);
68 outputDims[1] = lengths_[i];
69 output->Resize(outputDims);
70 char* ptr =
static_cast<char*
>(output->raw_mutable_data(data.meta()));
71 memset(ptr, 0, output->nbytes());
72 outputRawData.push_back(ptr);
75 for (
int i = 0; i < batchSize; ++i) {
76 for (
int j = 0; j < OutputSize(); ++j) {
77 auto rangeStart = rangesData[rangesDataOffset++];
78 auto rangeLength = rangesData[rangesDataOffset++];
79 if (rangeLength == 0) {
86 "Range lengths missmatch for output #",
89 if (InputSize() == 2) {
90 context_.template CopyItems<Context, Context>(
93 rawData + rangeStart * itemsize,
94 outputRawData[j] + i * itemsize * lengths_[j]);
96 auto& key = Input(KEY);
97 auto* key_data = key.template data<int64_t>();
98 vector<std::pair<int64_t, const char*>> buffer;
99 for (
int b_i = 0; b_i < rangeLength; ++b_i) {
100 int64_t one_key_item = key_data[rangeStart + b_i];
101 auto* one_data_item = rawData + (rangeStart + b_i) * itemsize;
102 buffer.emplace_back(one_key_item, one_data_item);
107 [](
const std::pair<int64_t, const char*>& left,
108 const std::pair<int64_t, const char*>& right) {
109 return left.first < right.first;
111 for (
int b_i = 0; b_i < rangeLength; ++b_i) {
114 outputRawData[j] + (i * lengths_[j] + b_i) * itemsize,
121 CAFFE_ENFORCE_EQ(rangesDataOffset, ranges.size());
126 INPUT_TAGS(DATA, RANGES, KEY);
129 vector<int> lengths_;
134 #endif // CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...