Method onCompute

source/geometry/GeometryTensorArray.cpp:290–370

class GeometryTensorArrayScatter : 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[3];
        auto inDes = TensorUtils::getDescribe(tensorArrayInput);
        if (inDes->tensorArrayAttr == nullptr) {
            return false;
        }
        int oldSize = inDes->tensorArrayAttr->arraySize;
        auto output    = outputs[0];
        int elemSize = getElemSize(output, 0).second;
        auto indicesTensor = inputs[1];
        // tag index write or not
        std::vector<bool> isWrite(oldSize, false);
        // write index
        std::vector<int> indices(indicesTensor->elementSize());
        // not write index
        std::vector<int> remains;
        for (int i = 0; i < indices.size(); i++) {
            indices[i] = indicesTensor->host<int>()[i];
            if (i < oldSize) {
                isWrite[i] = true;
            }
        }
        for (int i = 0; i < oldSize; i++) {
            if (!isWrite[i]) {
                remains.push_back(i);
            }
        }
        auto outputDes = TensorUtils::getDescribe(output);
        outputDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
        outputDes->regions.resize(indices.size() + remains.size());
        // write value by indices
        for (int i = 0; i < indices.size(); i++) {
            MNN_ASSERT(indices[i] < outputDes->tensorArrayAttr->arraySize);
            auto& reg = outputDes->regions[i];
            reg.origin = inputs[2];
            reg.src.offset = i * elemSize;
            reg.src.stride[0] = 1;
            reg.src.stride[1] = 1;
            reg.src.stride[2] = 1;
            reg.dst.offset = indices[i] * elemSize;
            reg.dst.stride[0] = 1;
            reg.dst.stride[1] = 1;
            reg.dst.stride[2] = 1;
            reg.size[0] = elemSize;
            reg.size[1] = 1;
            reg.size[2] = 1;
        }
        if (remains.empty()) {
            return true;
        }
        // first write data, set zero
        bool firstWrite = isFirstWrite(inDes);
        if (firstWrite) {
            auto type = tensorArrayInput->getType();
            auto zeroConst = context.allocConst(op, {}, type);
            if (type == halide_type_of<float>()) {
                zeroConst->host<float>()[0] = 0.0;