gather_ranges_to_dense_op.h

#ifndef CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_
#define CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_

#include <math.h>

#include "caffe2/core/common_omp.h"
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/core/types.h"
#include "caffe2/utils/math.h"

#include <cstring>
#include <map>
#include <utility>

namespace caffe2 {
template <class Context>
class GatherRangesToDenseOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit GatherRangesToDenseOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        lengths_(this->template GetRepeatedArgument<int>("lengths")) {
    CAFFE_ENFORCE_GT(lengths_.size(), 0, "There has to be at least one length");
    for (auto length : lengths_) {
      CAFFE_ENFORCE_GT(length, 0, "Each length should be positive");
    }
  }

  bool RunOnDevice() override {
    return DispatchHelper<TensorTypes<int32_t, int64_t>>::call(
        this, this->template Input<Tensor>(RANGES, CPU));
  }

  template <typename Index>
  bool DoRunWithType() {
    auto& data = Input(DATA);
    auto& ranges = Input(RANGES);
    CAFFE_ENFORCE_EQ(data.dim(), 1, "Data has to be 1-D");
    CAFFE_ENFORCE_EQ(ranges.dim(), 3, "Ranges has to be 3-D");
    if (InputSize() == 3) {
      auto& key = Input(KEY);
      CAFFE_ENFORCE_EQ(key.dim(), 1, "Key has to be 1-D");
      CAFFE_ENFORCE(
          key.dtype().template Match<int64_t>(), "Key has to be type int64_t");
    }
    CAFFE_ENFORCE_EQ(
        ranges.size(1),
        lengths_.size(),
        "Nummber of ranges should match number of lengths");
    CAFFE_ENFORCE_EQ(
        ranges.size(1),
        OutputSize(),
        "Nummber of ranges should match number of outputs");
    CAFFE_ENFORCE_EQ(
        ranges.size(2), 2, "Ranges last dimension should be of size 2");

    auto* rawData = static_cast<const char*>(data.raw_data());
    auto* rangesData = ranges.template data<Index>();
    int rangesDataOffset = 0;
    auto itemsize = data.dtype().itemsize();

    auto batchSize = ranges.size(0);
    vector<int64_t> outputDims{batchSize, 0};
    vector<char*> outputRawData;
    for (int i = 0; i < OutputSize(); ++i) {
      auto* output = Output(i);
      outputDims[1] = lengths_[i];
      output->Resize(outputDims);
      char* ptr = static_cast<char*>(output->raw_mutable_data(data.dtype()));
      memset(ptr, 0, output->nbytes());
      outputRawData.push_back(ptr);
    }

    for (int i = 0; i < batchSize; ++i) {
      for (int j = 0; j < OutputSize(); ++j) {
        auto rangeStart = rangesData[rangesDataOffset++];
        auto rangeLength = rangesData[rangesDataOffset++];
        if (rangeLength == 0) {
          // empty range, will be filled with zeros
          continue;
        }
        CAFFE_ENFORCE_EQ(
            rangeLength,
            lengths_[j],
            "Range lengths missmatch for output #",
            j);

        if (InputSize() == 2) {
          context_.CopyItemsSameDevice(
              data.dtype(),
              rangeLength,
              rawData + rangeStart * itemsize,
              outputRawData[j] + i * itemsize * lengths_[j]);
        } else {
          auto& key = Input(KEY);
          auto* key_data = key.template data<int64_t>();
          vector<std::pair<int64_t, const char*>> buffer;
          for (int b_i = 0; b_i < rangeLength; ++b_i) {
            int64_t one_key_item = key_data[rangeStart + b_i];
            auto* one_data_item = rawData + (rangeStart + b_i) * itemsize;
            buffer.emplace_back(one_key_item, one_data_item);
          }
          std::sort(
              buffer.begin(),
              buffer.end(),
              [](const std::pair<int64_t, const char*>& left,
                 const std::pair<int64_t, const char*>& right) {
                return left.first < right.first;
              });
          for (int b_i = 0; b_i < rangeLength; ++b_i) {
            // Since this CPU only, directly copy to the destination.
            std::memcpy(
                outputRawData[j] + (i * lengths_[j] + b_i) * itemsize,
                buffer[b_i].second,
                itemsize);
          }
        }
      }
    }
    CAFFE_ENFORCE_EQ(rangesDataOffset, ranges.numel());

    return true;
  }

  INPUT_TAGS(DATA, RANGES, KEY);

 private:
  vector<int> lengths_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_GATHER_RANGES_TO_DENSE_OPS_H_