Caffe2 - C++ API
A deep learning, cross platform ML framework
extend_tensor_op.cc
1 
17 #include <atomic>
18 #include <mutex>
19 #include <unordered_map>
20 #include <vector>
21 #include "caffe2/core/operator.h"
22 #include "caffe2/core/tensor.h"
23 #include "caffe2/utils/math.h"
24 
25 namespace caffe2 {
26 
27 namespace {
28 
29 template <class Context>
30 class ExtendTensorOp final : public Operator<Context> {
31  public:
32  USE_OPERATOR_CONTEXT_FUNCTIONS;
33  ExtendTensorOp(const OperatorDef& operator_def, Workspace* ws)
34  : Operator<Context>(operator_def, ws),
35  growthPct_(OperatorBase::GetSingleArgument<int>("growthPct", 40)) {}
36 
37  bool RunOnDevice() override {
38  auto& old_tensor = Input(0);
39  auto& indices = Input(1);
40  auto* new_tensor = Output(0);
41  CAFFE_ENFORCE(indices.ndim() >= 1);
42  CAFFE_ENFORCE(
43  &old_tensor == new_tensor, "First argument must be in-place.");
44  CAFFE_ENFORCE(new_tensor->ndim() == indices.ndim());
45  CAFFE_ENFORCE(indices.ndim() == new_tensor->ndim());
46 
47  auto oldSize = new_tensor->size();
48  auto maxElem = 1 +
49  *(std::max_element(
50  indices.template data<int>(),
51  indices.template data<int>() + indices.size()));
52 
53  auto extendSize = (TIndex)maxElem - oldSize;
54  if (extendSize > 0) {
55  new_tensor->Extend(extendSize, growthPct_, &context_);
56  if (!new_tensor->meta().ctor()) {
57  auto oldSizeBytes = oldSize * new_tensor->meta().itemsize();
58  auto* dst = (char*)new_tensor->raw_mutable_data() + oldSizeBytes;
59  math::Set<char, Context>(
60  new_tensor->nbytes() - oldSizeBytes, 0, dst, &context_);
61  }
62  }
63  return true;
64  }
65 
66  int growthPct_;
67 };
68 
69 REGISTER_CPU_OPERATOR(ExtendTensor, ExtendTensorOp<CPUContext>);
70 
71 OPERATOR_SCHEMA(ExtendTensor)
72  .NumInputs(2)
73  .NumOutputs(1)
74  .EnforceInplace({{0, 0}})
75  .SetDoc(R"DOC(
76 Extend input 0 if necessary based on max element in input 1.
77 Input 0 must be the same as output, that is, it is required to be in-place.
78 Input 0 may have to be re-allocated in order for accommodate to the new size.
79 Currently, an exponential growth ratio is used in order to ensure amortized
80 constant time complexity.
81 All except the outer-most dimension must be the same between input 0 and 1.
82 )DOC")
83  .Input(0, "tensor", "The tensor to be extended.")
84  .Input(
85  1,
86  "new_indices",
87  "The size of tensor will be extended based on max element in "
88  "new_indices.")
89  .Output(
90  0,
91  "extended_tensor",
92  "Same as input 0, representing the mutated tensor.");
93 }
94 } // namespace
Copyright (c) 2016-present, Facebook, Inc.