Method onCompute

source/geometry/GeometryReduce.cpp:92–178

class GeometryReduce : 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 {
        MNN_ASSERT(1 == outputs.size());
        MNN_ASSERT(inputs.size() >= 1);
        auto reduct          = op->main_as_ReductionParam();
        auto reductOp        = reduct->operation();
        std::vector<int> axises;
        if (inputs.size() >= 2) {
            auto size = inputs[1]->elementSize();
            auto dims = inputs[1]->host<int32_t>();
            for (int i = 0; i < size; ++i) {
                axises.emplace_back(dims[i]);
            }
        } else {
            auto reduct = op->main_as_ReductionParam();
            if (nullptr != reduct->dim()) {
                for (int i = 0; i < reduct->dim()->size(); ++i) {
                    axises.emplace_back(reduct->dim()->data()[i]);
                }
            }
        }
        for (int i = 0; i < axises.size(); ++i) {
            if (axises[i] < 0) {
                axises[i] = inputs[0]->dimensions() + axises[i];
            }
        }
        if (1 == axises.size() && TensorUtils::getDescribe(inputs[0])->dimensionFormat != MNN_DATA_FORMAT_NC4HW4 && TensorUtils::getDescribe(outputs[0])->dimensionFormat != MNN_DATA_FORMAT_NC4HW4) {
            auto cmd = GeometryComputerUtils::makeReduce(reductOp, inputs[0], outputs[0], axises[0]);
            res.command.emplace_back(std::move(cmd));
            return true;
        }
        // prod([]) = 1
        if (inputs[0]->elementSize() == 0) {
            if(!context.allocTensor(outputs[0])) {
                return false;
            }
            float res;
            switch (reductOp) {
                case ReductionType_PROD:
                    res = 1.0f;
                    break;
                default:
                    res = 0.0f;
                    break;
            }
            if (outputs[0]->getType() == halide_type_of<float>()) {
                outputs[0]->host<float>()[0] = (float)res;
            } else {
                outputs[0]->host<int>()[0] = (int)res;
            }
            return true;
        }
        auto reduceDims = _computeReduceDims(inputs, axises);
        Tensor* currentInput = inputs[0];
        MNN_ASSERT(reduceDims.size() > 0);
        auto dimType = currentInput->getDimensionType();
        for (int i = 0; i < reduceDims.size(); ++i) {
            auto& iter   = reduceDims[i];