Function buildScatterND

source/geometry/GeometryScatter.cpp:13–152

#include "core/OpCommonUtils.hpp"
namespace MNN {
static bool buildScatterND(const Op* op, Tensor* indices, Tensor* updates, Tensor* data, Tensor* output,
                          int N, int D, int S, int totalSize, int reduction,
                          GeometryComputer::Context& context, CommandBuffer& res) {
    // get stride
    std::shared_ptr<Tensor> constStride(Tensor::createDevice<int>({D}));
    if (!context.allocTensor(constStride.get())) {
        return false;
    }
    int count = output->elementSize();
    for (int i = 0; i < D; ++i) {
        count = count / output->length(i);
        constStride->host<int>()[i] = count;
    }
    res.extras.emplace_back(constStride);
    std::shared_ptr<Tensor> broadcastStride(Tensor::createDevice<int>({N, D}));
    {
        // [D] => [N, D]
        auto des = TensorUtils::getDescribe(broadcastStride.get());
        des->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
        des->regions.resize(1);
        des->regions[0].origin = constStride.get();
        des->regions[0].size[0] = 1;
        des->regions[0].size[1] = N;
        des->regions[0].size[2] = D;
        des->regions[0].dst.stride[0] = N*D;
        des->regions[0].dst.stride[1] = D;
        des->regions[0].dst.stride[2] = 1;
        des->regions[0].src.stride[0] = 0;
        des->regions[0].src.stride[1] = 0;
        des->regions[0].src.stride[2] = 1;
        res.extras.emplace_back(broadcastStride);
    }
    // reshape indices: [dims1, D] -> [N, D]
    std::shared_ptr<Tensor> reshapeIndice(Tensor::createDevice<int>({N, D}));
    {
        auto des = TensorUtils::getDescribe(reshapeIndice.get());
        des->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
        des->regions = {GeometryComputerUtils::makeRawAddressRef(indices, 0, N * D)};
        res.extras.emplace_back(reshapeIndice);
    }
    // get index
    std::shared_ptr<Tensor> mulIndice(Tensor::createDevice<int>({N, D}));
    {
        // [N, D] * [N, D] => [N, D]
        auto cmd = GeometryComputerUtils::makeBinary(BinaryOpOperation_MUL, reshapeIndice.get(), broadcastStride.get(), mulIndice.get());
        res.extras.emplace_back(mulIndice);
        res.command.emplace_back(std::move(cmd));
    }
    std::shared_ptr<Tensor> indiceOneLine(Tensor::createDevice<int>({N, 1}));
    {
        // [N, D] => [N, 1]
        auto cmd = GeometryComputerUtils::makeReduce(ReductionType_SUM, mulIndice.get(), indiceOneLine.get());
        res.extras.emplace_back(indiceOneLine);
        res.command.emplace_back(std::move(cmd));
    }
    auto outputDes = TensorUtils::getDescribe(output);
    flatbuffers::FlatBufferBuilder builder;
    {