1 #ifndef CAFFE2_OPERATORS_LENGTHS_TILE_OP_H_ 2 #define CAFFE2_OPERATORS_LENGTHS_TILE_OP_H_ 4 #include "caffe2/core/operator.h" 5 #include "caffe2/utils/math.h" 9 template <
class Context>
12 USE_OPERATOR_CONTEXT_FUNCTIONS;
15 bool RunOnDevice()
override {
16 auto& data = Input(DATA);
17 auto& lengths = Input(LENGTHS);
18 auto* output = Output(0);
20 CAFFE_ENFORCE_EQ(lengths.ndim(), 1,
"LENGTHS must be 1-D");
21 CAFFE_ENFORCE_GE(data.ndim(), 1,
"DATA should be at least 1-D");
22 CAFFE_ENFORCE_EQ(lengths.size(), data.dim(0));
27 lengths_host_.
CopyFrom(lengths, &cpuContext);
28 auto lengths_size = lengths_host_.
size();
29 auto* lengths_data = lengths_host_.
data<int32_t>();
31 int32_t total_length = 0;
32 math::Sum<int32_t, CPUContext>(
33 lengths_size, lengths_data, &total_length, &cpuContext);
35 auto shape = data.dims();
36 shape[0] = total_length;
37 output->Resize(shape);
39 auto block_bytesize = data.size_from_dim(1) * data.meta().itemsize();
40 auto src =
static_cast<const char*
>(data.raw_data());
41 auto out =
static_cast<char*
>(output->raw_mutable_data(data.meta()));
43 for (TIndex i = 0; i < lengths_size; ++i) {
44 auto length = lengths_data[i];
45 CAFFE_ENFORCE_GE(length, 0);
46 for (int32_t j = 0; j < length; ++j) {
47 context_.template CopyBytes<Context, Context>(block_bytesize, src, out);
48 out += block_bytesize;
50 src += block_bytesize;
55 INPUT_TAGS(DATA, LENGTHS);
63 #endif // CAFFE2_OPERATORS_LENGTHS_TILE_OP_H_ const T * data() const
Returns a typed pointer of the underlying storage.
void CopyFrom(const Tensor< SrcContext > &src, ContextForCopy *context)
Copies the data from a source tensor, with a contex provided to carry out the underlying memcpy opera...
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
TIndex size() const
Returns the size (i.e.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...