| 40 | class GeometryShape : public GeometryComputer { |
| 41 | public: |
| 42 | virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, |
| 43 | Context& context, CommandBuffer& res) const override { |
| 44 | if (nullptr == TensorUtils::getDescribeOrigin(outputs[0])->mem.get()) { |
| 45 | auto originSize = outputs[0]->length(0); |
| 46 | outputs[0]->setLength(0, MNN_MAX_TENSOR_DIM); |
| 47 | if(!context.allocTensor(outputs[0])) { |
| 48 | return false; |
| 49 | } |
| 50 | outputs[0]->setLength(0, originSize); |
| 51 | } |
| 52 | auto& ib = inputs[0]->buffer(); |
| 53 | auto outputData = outputs[0]->host<int>(); |
| 54 | auto inputFormat = TensorUtils::getDescribe(inputs[0])->dimensionFormat; |
| 55 | int shapeData[MNN_MAX_TENSOR_DIM]; |
| 56 | int rank = ib.dimensions; |
| 57 | if ((inputFormat == MNN_DATA_FORMAT_NC4HW4) && TensorUtils::getDescribe(outputs[0])->dimensionFormat == MNN_DATA_FORMAT_NHWC) { |
| 58 | rank = 4; |
| 59 | shapeData[0] = ib.dim[0].extent; |
| 60 | shapeData[1] = ib.dim[2].extent; |
| 61 | shapeData[2] = ib.dim[3].extent; |
| 62 | shapeData[3] = ib.dim[1].extent; |
| 63 | } else { |
| 64 | for (int i = 0; i < ib.dimensions; i++) { |
| 65 | shapeData[i] = ib.dim[i].extent; |
| 66 | } |
| 67 | } |
| 68 | auto range = _resolveShapeRange(op, rank); |
| 69 | for (int i = range.first; i < range.second; ++i) { |
| 70 | outputData[i - range.first] = shapeData[i]; |
| 71 | } |
| 72 | return true; |
| 73 | } |
| 74 | }; |
| 75 | |
| 76 | class GeometryRank : public GeometryComputer { |
nothing calls this directly
no test coverage detected