MNN/tools/cpp/GpuInterTest.cpp

//
//  ModuleBasic.cpp
//  MNN
//
//  Created by MNN on 2021/10/15.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "MNN_generated.h"
#include <MNN/expr/Expr.hpp>
#include <MNN/expr/Module.hpp>
#include <MNN/expr/ExprCreator.hpp>
#define MNN_OPEN_TIME_TRACE
#include <MNN/AutoTime.hpp>
#include "rapidjson/document.h"
#include "core/MemoryFormater.h"
#include <fstream>
#include <sstream>
#include <numeric>
#include "ExprDebug.hpp"
#include "MNN/MNNSharedContext.h"
using namespace MNN::Express;
using namespace MNN;

#ifdef __ANDROID__
#include <android/hardware_buffer.h>
    
static AHardwareBuffer* creatAHardwareBuffer(int width, int height, void *data){
    // 创建和初始化硬件缓冲区
    AHardwareBuffer_Desc bufferDesc = {};
    bufferDesc.width = width;
    bufferDesc.height = height;
    bufferDesc.layers = 1;
    bufferDesc.format = AHARDWAREBUFFER_FORMAT_R8G8B8A8_UNORM;
    bufferDesc.usage = AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN | AHARDWAREBUFFER_USAGE_GPU_SAMPLED_IMAGE;

    AHardwareBuffer* buffer = nullptr;
    int result = AHardwareBuffer_allocate(&bufferDesc, &buffer);
    if(result != 0) {
        // Handle allocation error
        MNN_PRINT("alloc AHardwareBuffer failed   %d\n", result);
    }
        
    if(nullptr != data){
        void* map = nullptr;
        ARect rect = { 0, 0, width, height };  // Define the region to lock
        result = AHardwareBuffer_lock(buffer, AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN, -1, &rect, &map);
        if (result != 0) {
            // Handle lock failure
            MNN_PRINT("Handle lock failed\n");
        }
        if (map) {
            // Now write your pixel data to 'data'
            // For example, fill it with a solid color:
            memcpy(map, data, width * height * 4); // Assuming RGBA8888 format
        }
            
        AHardwareBuffer_unlock(buffer, nullptr);
    }
    return buffer;
}
static void copyDataFromAHardWareBuffer(AHardwareBuffer* buffer, int width, int height, void *data){
    int result = 0;
    if(nullptr != data){
        void* map = nullptr;
        ARect rect = { 0, 0, width, height };  // Define the region to lock
        result = AHardwareBuffer_lock(buffer, AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN, -1, &rect, &map);
        if (result != 0) {
            MNN_PRINT("Handle lock failed\n");
        }
        if (map) {
            memcpy(data, map, width * height * 4);
        }
            
        AHardwareBuffer_unlock(buffer, nullptr);
    }
}
static void ReleaseAHardWareBuffer(AHardwareBuffer* buffer){
    if(buffer != nullptr){
        AHardwareBuffer_release(buffer);
    }
}
#endif

int main(int argc, char *argv[]) {
    if (argc < 3) {
        MNN_ERROR("Usage: ./GpuInterTest.out ${test.mnn} ${Dir} [testMode] [forwardType] [numberThread] [precision | memory]\n");
        return 0;
    }
    std::string modelName = argv[1];
    std::string directName = argv[2];
    MNN_PRINT("Test %s from input info: %s\n", modelName.c_str(), directName.c_str());
    std::map<std::string, float> inputInfo;
    std::map<std::string, std::vector<int>> inputShape;
    std::vector<std::string> inputNames;
    std::vector<std::string> outputNames;
    int repeatNumber = 1;
    bool shapeMutable = true;
    std::vector<VARP> inputs;
    std::vector<VARP> outputs;
    if (inputNames.empty()) {
        rapidjson::Document document;
        std::ostringstream jsonNameOs;
        jsonNameOs << directName << "/input.json";
        std::ifstream fileNames(jsonNameOs.str().c_str());
        std::ostringstream output;
        output << fileNames.rdbuf();
        auto outputStr = output.str();
        document.Parse(outputStr.c_str());
        if (document.HasParseError()) {
            MNN_ERROR("Invalid json\n");
            return 0;
        }
        if (document.HasMember("inputs")) {
            auto inputsInfo = document["inputs"].GetArray();
            for (auto iter = inputsInfo.begin(); iter !=inputsInfo.end(); iter++) {
                auto obj = iter->GetObject();
                std::string name = obj["name"].GetString();
                inputNames.emplace_back(name);
                MNN_PRINT("%s\n", name.c_str());
                if (obj.HasMember("value")) {
                    float value = obj["value"].GetFloat();
                    inputInfo.insert(std::make_pair(name, value));
                }
                if (obj.HasMember("shape")) {
                    auto dims = obj["shape"].GetArray();
                    std::vector<int> shapes;
                    for (auto iter = dims.begin(); iter != dims.end(); iter++) {
                        shapes.emplace_back(iter->GetInt());
                    }
                    inputShape.insert(std::make_pair(name, shapes));
                }
            }
        }
        if (document.HasMember("outputs")) {
            auto array = document["outputs"].GetArray();
            for (auto iter = array.begin(); iter !=array.end(); iter++) {
                std::string name = iter->GetString();
                MNN_PRINT("output: %s\n", name.c_str());
                outputNames.emplace_back(name);
            }
        }
        if (document.HasMember("shapeMutable")) {
            shapeMutable = document["shapeMutable"].GetBool();
        }
        if (document.HasMember("repeat")) {
            repeatNumber = document["repeat"].GetInt();
        }
    }
    int testMode = 0;
    //testMode = 0 AhardwareBuffer
    if(argc > 3){
        testMode = atoi(argv[3]);
        MNN_PRINT("Use extra forward type: %d(0:AhardwareBuffer)\n", testMode);
    }

    auto type = MNN_FORWARD_CPU;
    if (argc > 4) {
        type = (MNNForwardType)atoi(argv[4]);
        MNN_PRINT("Use extra forward type: %d\n", type);
    }

    // Default single thread
    int modeNum = 1;
    if (argc > 5) {
        modeNum = ::atoi(argv[5]);
    }

    int precision = BackendConfig::Precision_Normal;
    int memory = BackendConfig::Memory_Normal;
    if (argc > 6) {
        int mask = atoi(argv[6]);
        precision = mask % 4;
        memory = (mask / 4) % 4;
    }
    const char* cacheFileName = ".tempcache";
    FUNC_PRINT(precision);
    FUNC_PRINT(memory);
    FUNC_PRINT_ALL(cacheFileName, s);
    // create session
    MNN::ScheduleConfig config;
    config.type      = type;
    /*modeNum means gpuMode for GPU usage, Or means numThread for CPU usage.*/
    config.numThread = modeNum;
    // If type not fount, let it failed
    config.backupType = type;
    BackendConfig backendConfig;
    backendConfig.precision = static_cast<MNN::BackendConfig::PrecisionMode>(precision);
    backendConfig.memory = static_cast<MNN::BackendConfig::MemoryMode>(memory);
    config.backendConfig     = &backendConfig;

    MNN::Express::Module::Config mConfig;
    mConfig.shapeMutable = shapeMutable;
#ifdef __ANDROID__
    std::vector<AHardwareBuffer*> AHardwarePtrInputVec;
    std::vector<AHardwareBuffer*> AHardwarePtrOutputVec;
#endif

    std::shared_ptr<Executor::RuntimeManager> rtmgr(Executor::RuntimeManager::createRuntimeManager(config));
    rtmgr->setCache(cacheFileName);
    std::shared_ptr<Module> net;
    {
        AUTOTIME;
        net.reset(Module::load(inputNames, outputNames, modelName.c_str(), rtmgr, &mConfig));
        if (net == nullptr) {
            MNN_PRINT("Error: can't load module\n");
            return 0;
        }
    }
    auto mInfo = net->getInfo();
#ifdef __ANDROID__
    AHardwarePtrInputVec.resize(mInfo->inputs.size());
    AHardwarePtrOutputVec.resize(outputNames.size());
#endif
    if (inputs.empty()) {
        inputs.resize(mInfo->inputs.size());
        for (int i=0; i<inputs.size(); ++i) {
            inputs[i] = _Input(mInfo->inputs[i].dim, mInfo->inputs[i].order, mInfo->inputs[i].type);
        }
        for (int i=0; i<inputs.size(); ++i) {
            auto inputName = inputNames[i];
            // Resize
            auto info = inputs[i]->getInfo();
            int width = info->dim[3], height = info->dim[2], channel = info->dim[1];
            auto shapeIter = inputShape.find(inputName);
            if (shapeIter != inputShape.end()) {
                auto s = shapeIter->second;
                inputs[i] = _Input(s, mInfo->inputs[i].order, mInfo->inputs[i].type);
                width = s[3];
                height = s[2];
                channel = s[1];
            }
            // set input device ptr
#ifdef __ANDROID__
            // OpenGL Texture defaultFormat NC4HW4
            if(testMode == 0){
                width = width * ((channel + 3) / 4);
                AHardwarePtrInputVec[i] = creatAHardwareBuffer(width,height,nullptr);
                volatile uint64_t value = (uint64_t)AHardwarePtrInputVec[i];
                inputs[i]->setDevicePtr((void*)value, MNN_MEMORY_AHARDWAREBUFFER);
            }
#endif
        }
    }

    bool modelError = false;
    for (int repeat = 0; repeat < repeatNumber; ++repeat) {
        AUTOTIME;
        auto outputs = net->onForward(inputs);
        if (outputs.empty()) {
            MNN_ERROR("Error in forward\n");
            return 0;
        }
        for (int i=0; i<outputNames.size(); ++i) {
            auto info = outputs[i]->getInfo();
            int width = info->dim[3], height = info->dim[2], channel = info->dim[1];
            // copy output to device ptr
#ifdef __ANDROID__
            if(testMode == 0){
                AHardwarePtrOutputVec[i] = creatAHardwareBuffer(width,height,nullptr);
                volatile uint64_t value = (uint64_t)AHardwarePtrOutputVec[i];
                outputs[i]->copyToDevicePtr((void*)value, MNN_MEMORY_AHARDWAREBUFFER);
            }
#endif
        }

        // Print module's memory
        float memoryInMB = 0.0f;
        rtmgr->getInfo(Interpreter::MEMORY, &memoryInMB);
        FUNC_PRINT_ALL(memoryInMB, f);
    }
#ifdef __ANDROID__
    if(testMode == 1){
        for(int i = 0; i < AHardwarePtrInputVec.size(); ++i){
            ReleaseAHardWareBuffer(AHardwarePtrInputVec[i]);
        }
        for(int i = 0; i < AHardwarePtrOutputVec.size(); ++i){
            ReleaseAHardWareBuffer(AHardwarePtrOutputVec[i]);
        }
    }
#endif
    rtmgr->updateCache();
    return 0;
}