Method onCompute

source/geometry/GeometryTensorArray.cpp:248–285

class GeometryTensorArrayGather : public GeometryComputer {
public:
    virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                           Context& context, CommandBuffer& res) const override {
        auto tensorArrayInput = inputs[2];
        auto inDes = TensorUtils::getDescribe(tensorArrayInput);
        if (inDes->tensorArrayAttr == nullptr) {
            return false;
        }
        auto indicesTensor = inputs[1];
        std::vector<int> indices(indicesTensor->elementSize());
        for (int i = 0; i < indices.size(); i++) {
            indices[i] = indicesTensor->host<int>()[i];
        }
        auto output    = outputs[0];
        auto outputDes = TensorUtils::getDescribe(output);
        outputDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
        outputDes->regions.resize(indices.size());
        int arraySize = inDes->tensorArrayAttr->arraySize;
        int dstOffset = 0;
        for (int i = 0; i < indices.size(); i++) {
            MNN_ASSERT(indices[i] < arraySize);
            auto elemSize = getElemSize(tensorArrayInput, indices[i]);
            auto& reg = outputDes->regions[i];
            reg.origin = tensorArrayInput;
            reg.src.offset = elemSize.first;
            reg.src.stride[0] = 1;
            reg.src.stride[1] = 1;
            reg.src.stride[2] = 1;
            reg.dst.offset = dstOffset;
            reg.dst.stride[0] = 1;
            reg.dst.stride[1] = 1;
            reg.dst.stride[2] = 1;
            reg.size[0] = elemSize.second;
            reg.size[1] = 1;
            reg.size[2] = 1;
            dstOffset += elemSize.second;
        }
        return true;
    }
};

class GeometryTensorArrayScatter : public GeometryComputer {