MNN/source/shape/SizeComputer.cpp

//
//  SizeComputer.cpp
//  MNN
//
//  Created by MNN on 2019/01/10.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "shape/SizeComputer.hpp"
#include <stdlib.h>
#include <mutex>
#include "core/Macro.h"
#include "core/TensorUtils.hpp"
//#define MNN_DEBUG_TENSOR_SIZE
namespace MNN {
void registerShapeOps();
SizeComputerSuite* SizeComputerSuite::gInstance = nullptr;

SizeComputerSuite::~SizeComputerSuite() {
    for (auto& iter : mRegistry) {
        delete iter.second;
    }
}

void SizeComputerSuite::init() {
    if (nullptr != gInstance) {
        return;
    }
    gInstance = new SizeComputerSuite;
    registerShapeOps();
}

SizeComputerSuite* SizeComputerSuite::get() {
    return gInstance;
}

void SizeComputerSuite::insert(SizeComputer* t, OpType type) {
    mRegistry.insert(std::make_pair(type, t));
}

SizeComputer* SizeComputerSuite::search(OpType name) {
    auto iter = mRegistry.find(name);
    if (iter == mRegistry.end()) {
        return nullptr;
    }
    return iter->second;
}
float SizeComputer::onComputeFlops(const MNN::Op* op, const std::vector<Tensor*>& inputs,
                                   const std::vector<Tensor*>& outputs) const {
    MNN_ASSERT(outputs.size() >= 1);
    return (float)outputs[0]->elementSize() / 1024.0f / 1024.0f;
}

float SizeComputer::computeFlops(const MNN::Op* op, const std::vector<Tensor*>& inputs,
                                 const std::vector<Tensor*>& outputs) {
    auto computeFactory = SizeComputerSuite::get();
    auto computer       = computeFactory->search(op->type());
    if (nullptr != computer) {
        return computer->onComputeFlops(op, inputs, outputs);
    }
    auto sumFlops = 0.0f;
    for (auto output : outputs) {
        sumFlops += (float)output->elementSize() / 1024.0f / 1024.0f;
    }
    return sumFlops;
}

bool SizeComputer::computeOutputSize(const MNN::Op* op, const std::vector<Tensor*>& inputs,
                                     const std::vector<Tensor*>& outputs) {
    auto computeFactory = SizeComputerSuite::get();
    // When op is nullptr, it means a copy op
    if (nullptr != op) {
        // Don't support compute shape for control flow op
        if (op->type() == OpType_While || op->type() == OpType_If) {
            return false;
        }
        auto computer = computeFactory->search(op->type());
        if (nullptr != computer) {
            bool ret = computer->onComputeSize(op, inputs, outputs);
#ifdef MNN_DEBUG_TENSOR_SIZE
            if (op->name() != nullptr) {
                MNN_PRINT("\t===> compute shape: %s, [%s]\n", op->name()->c_str(), MNN::EnumNameOpType(op->type()));
            } else {
                MNN_PRINT("\t===> compute shape:[%s]\n", MNN::EnumNameOpType(op->type()));
            }
            if (inputs.size()) {
                MNN_PRINT("Inputs:\n");
                for (auto o : inputs) {
                    if (o->dimensions() == 0) {
                        MNN_PRINT("\t*Scalar*");
                    }
                    for (int i = 0; i < o->dimensions(); ++i) {
                        MNN_PRINT("%d, ", o->length(i));
                    }
                    MNN_PRINT("\n");
                }
            }
            MNN_PRINT("Outputs:\n");
            for (auto o : outputs) {
                if (o->dimensions() == 0) {
                    MNN_PRINT("\t*Scalar*");
                }
                for (int i = 0; i < o->dimensions(); ++i) {
                    MNN_PRINT("%d, ", o->length(i));
                }
                MNN_PRINT("\n");
            }
#endif
            return ret;
        }
    }

    // Default Set to the same
    if (inputs.size() >= 1 && outputs.size() == 1) {
        if (inputs[0] == outputs[0]) {
            return true;
        }
        const auto& ib = inputs[0]->buffer();
        auto& ob       = outputs[0]->buffer();
        memcpy(ob.dim, ib.dim, sizeof(halide_dimension_t) * ib.dimensions);
        ob.dimensions                                         = ib.dimensions;
        ob.type                                               = ib.type;
        TensorUtils::getDescribe(outputs[0])->dimensionFormat = TensorUtils::getDescribe(inputs[0])->dimensionFormat;

        return true;
    }
    // Not Support
    MNN_PRINT("Can't compute size for %d, name=%s\n", op->type(), op->name() ? op->name()->c_str() : "");

    return false;
}

std::vector<int> SizeComputer::needInputContent(const MNN::Op* op, int inputSize) {
    auto computeFactory = SizeComputerSuite::get();
    // When op is nullptr, it means a copy op
    if (nullptr != op) {
        // when hasOutputShape = true, deconv last is outputShape
        if (op->type() == OpType_Deconvolution && op->main_as_Convolution2D() && op->main_as_Convolution2D()->common()) {
            if (op->main_as_Convolution2D()->common()->hasOutputShape()) {
                return std::vector<int>{ inputSize - 1 };
            }
        }
        auto computer = computeFactory->search(op->type());
        if (nullptr != computer) {
            return computer->mNeedContentInputIndex;
        }
    }
    return std::vector<int>{};
}
bool SizeComputer::computeBroadCastDims(const MNN::Op* op, const std::vector<Tensor*>& inputs,
                                 const std::vector<Tensor*>& outputs) {
    int maxDimensions = inputs[0]->dimensions();
    int maxIndex = 0;
    for (int index=1; index < inputs.size(); ++index) {
        if (inputs[index]->dimensions() > maxDimensions) {
            maxDimensions = inputs[index]->dimensions();
            maxIndex = index;
        }
    }
    int outputDims[MNN_MAX_TENSOR_DIM];
    for (int i = 0; i < maxDimensions; i++) {
        outputDims[i] = inputs[maxIndex]->length(i);
    }
    for (int index=0; index < inputs.size(); ++index) {
        if (index == maxIndex) {
            continue;
        }
        auto input1 = inputs[index];
        auto input0 = inputs[maxIndex];
        const int diffDimension = maxDimensions - input1->dimensions();
        for (int i = diffDimension; i < maxDimensions; i++) {
            const int input1Index = i - diffDimension;
            int dim1 = input1->buffer().dim[input1Index].extent;
            if (dim1 != outputDims[i] && (dim1 != 1 && outputDims[i] != 1)) {
                MNN_ERROR("Broad cast error, dim1 = %d, dim2 = %d\n", dim1, outputDims[i]);
                return false;
            }
            if (dim1 == outputDims[i]) {
                continue;
            }
            if (dim1 != outputDims[i] && (dim1 == 1 || outputDims[i] == 1)) {
                outputDims[i] = outputDims[i] * dim1;
            } else {
                return false;
            }
        }
    }
    auto& ob       = outputs[0]->buffer();
    ob.dimensions = maxDimensions;
    for (int i = 0; i < maxDimensions; i++) {
        ob.dim[i].extent = outputDims[i];
    }
    return true;
}
} // namespace MNN