1 //
2 // GeometryConcat.cpp
3 // MNN
4 //
5 // Created by MNN on 2020/04/07.
6 // Copyright © 2018, Alibaba Group Holding Limited
7 //
8
9 #include "geometry/GeometryComputer.hpp"
10 namespace MNN {
11 class GeometryConcat : public GeometryComputer {
12 public:
onCompute(const Op * op,const std::vector<Tensor * > & inputs,const std::vector<Tensor * > & outputs,Context & context,CommandBuffer & res) const13 virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
14 Context& context, CommandBuffer& res) const override {
15 MNN_ASSERT(1 == outputs.size());
16 MNN_ASSERT(inputs.size() >= 1);
17 int basicAxis = 0;
18 bool inputFix = false;
19 if (op->type() == OpType_Concat) {
20 basicAxis = op->main_as_Axis()->axis();
21 } else if (op->type() == OpType_QuantizedConcat) {
22 basicAxis = op->main_as_QuantizedConcat()->axis();
23 } else if (op->type() == OpType_Pack) {
24 basicAxis = op->main_as_PackParam()->axis();
25 inputFix = true;
26 }
27 auto output = outputs[0];
28 int axis = basicAxis;
29 if (axis < 0) {
30 axis = output->dimensions() + axis;
31 }
32 auto outputDes = TensorUtils::getDescribe(output);
33 outputDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
34
35 int outside = 1;
36 int inside = 1;
37 for (int i = 0; i < axis; ++i) {
38 outside *= output->length(i);
39 }
40 for (int i = axis + 1; i < output->dimensions(); ++i) {
41 inside *= output->length(i);
42 }
43 int offset = 0;
44 outputDes->regions.clear();
45 if (output->elementSize() <= 0) {
46 return true;
47 }
48
49 for (int i = 0; i < inputs.size(); ++i) {
50 auto t = inputs[i];
51 if (t->elementSize() == 0) {
52 continue;
53 }
54 Tensor::InsideDescribe::Region dstSlice;
55 int basicStride = 1;
56 int basicOffset = 0;
57
58 dstSlice.origin = t;
59 dstSlice.src.offset = basicOffset;
60 dstSlice.dst.offset = offset * inside;
61 dstSlice.size[0] = outside;
62 if (inputFix) {
63 dstSlice.size[1] = 1;
64 } else {
65 dstSlice.size[1] = t->length(axis);
66 }
67 dstSlice.dst.stride[0] = inside * output->length(axis);
68 dstSlice.dst.stride[1] = inside;
69 dstSlice.dst.stride[2] = 1;
70 dstSlice.size[2] = inside;
71 offset += dstSlice.size[1];
72 dstSlice.src.stride[0] = basicStride * dstSlice.size[1] * dstSlice.size[2];
73 dstSlice.src.stride[1] = basicStride * dstSlice.size[2];
74 dstSlice.src.stride[2] = basicStride;
75 outputDes->regions.emplace_back(std::move(dstSlice));
76 }
77 return true;
78 }
79 };
_create()80 static void _create() {
81 std::shared_ptr<GeometryComputer> comp(new GeometryConcat);
82 GeometryComputer::registerGeometryComputer(comp, {OpType_QuantizedConcat, OpType_Concat, OpType_Pack});
83 }
84
85 REGISTER_GEOMETRY(GeometryConcat, _create);
86
87 } // namespace MNN
88