wehub-resource-sync
248 lines
9.6 KiB
C++
248 lines
9.6 KiB
C++
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <memory>
|
|
#include <cstdlib>
|
|
#include <ctime>
|
|
|
|
#include <MNN/expr/Module.hpp>
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
#include <MNN/expr/Executor.hpp>
|
|
#include "core/MNNFileUtils.h"
|
|
|
|
#include "rapidjson/document.h"
|
|
#include "rapidjson/istreamwrapper.h"
|
|
|
|
#include <limits>
|
|
|
|
static void saveInputOutputs(const MNN::Express::Module::Info* info, std::vector<MNN::Express::VARP> inputs, std::vector<MNN::Express::VARP> outputs, const std::string & outputDir, int index) {
|
|
MNN_ASSERT(info->inputNames.size() == inputs.size());
|
|
MNN_ASSERT(info->outputNames.size() == outputs.size());
|
|
for (int i=0; i<info->inputNames.size(); ++i) {
|
|
inputs[i].fix(MNN::Express::VARP::CONSTANT);
|
|
inputs[i]->setName(info->inputNames[i]);
|
|
}
|
|
for (int i=0; i<info->outputNames.size(); ++i) {
|
|
outputs[i]->setName(info->outputNames[i]);
|
|
}
|
|
auto subDir = MNNFilePathConcat(outputDir, std::to_string(index));
|
|
if (!(MNNCreateDir(subDir.c_str()))) {
|
|
MNN_PRINT("Failed to create dir %s.\n", outputDir.c_str());
|
|
}
|
|
|
|
std::string inputPath = MNNFilePathConcat(subDir, "input.mnn");
|
|
std::string outputPath = MNNFilePathConcat(subDir, "output.mnn");
|
|
MNN::Express::Variable::save(inputs, inputPath.c_str());
|
|
MNN::Express::Variable::save(outputs, outputPath.c_str());
|
|
MNN_PRINT("Successfully generate %s and %s.\n", inputPath.c_str(), outputPath.c_str());
|
|
}
|
|
|
|
static void createInputsForLLM(int seqLen, int hiddenSize, const std::string& attentionMaskType, bool lastLogit, std::vector<MNN::Express::VARP>& inputs) {
|
|
if (attentionMaskType != "float") {
|
|
MNN_ERROR("Don't support Attention Mask Type other than 'float', currently.\n");
|
|
return;
|
|
}
|
|
|
|
MNN::Express::VARP inputIdx = MNN::Express::_Input({seqLen, 1, hiddenSize}, MNN::Express::NCHW, halide_type_of<float>());
|
|
float * inputIdxData = inputIdx->writeMap<float>();
|
|
for (int i = 0; i < seqLen * hiddenSize; ++i) {
|
|
inputIdxData[i] = (float)(rand()) / RAND_MAX;
|
|
}
|
|
inputs.push_back(inputIdx);
|
|
|
|
MNN::Express::VARP attentionMask = MNN::Express::_Input({1, 1, seqLen, seqLen}, MNN::Express::NCHW, halide_type_of<float>());
|
|
float * attentionMaskData = attentionMask->writeMap<float>();
|
|
for (int i = 0; i < seqLen; ++i) {
|
|
for (int j = 0; j < seqLen; ++j) {
|
|
attentionMaskData[i * seqLen + j] = (j > i) * std::numeric_limits<float>::lowest();
|
|
}
|
|
}
|
|
inputs.push_back(attentionMask);
|
|
|
|
MNN::Express::VARP positionIds = MNN::Express::_Input({1, seqLen}, MNN::Express::NCHW, halide_type_of<int>());
|
|
int * positionIdsData = positionIds->writeMap<int>();
|
|
for (int i = 0; i < seqLen; i++) {
|
|
positionIdsData[i] = i;
|
|
}
|
|
inputs.push_back(positionIds);
|
|
|
|
int logitsIndexValue = lastLogit ? -1 : 0;
|
|
MNN::Express::VARP logitsIndex = MNN::Express::_Const((const void *) &logitsIndexValue, {1}, MNN::Express::NHWC, halide_type_of<int>());
|
|
inputs.push_back(logitsIndex);
|
|
|
|
return;
|
|
}
|
|
|
|
static void createInputsForEmbedding(int seqLen, int hiddenSize, const std::string& attentionMaskType, std::vector<MNN::Express::VARP>& inputs) {
|
|
MNN::Express::VARP inputEmbeds = MNN::Express::_Input({seqLen, 1, hiddenSize}, MNN::Express::NCHW, halide_type_of<float>());
|
|
float* inputEmbedsData = inputEmbeds->writeMap<float>();
|
|
for (int i = 0; i < seqLen * hiddenSize; ++i) {
|
|
inputEmbedsData[i] = (float)(rand()) / RAND_MAX;
|
|
}
|
|
inputs.push_back(inputEmbeds);
|
|
|
|
MNN::Express::VARP attentionMask = MNN::Express::_Input({1, 1, seqLen, seqLen}, MNN::Express::NCHW, halide_type_of<float>());
|
|
float* attentionMaskData = attentionMask->writeMap<float>();
|
|
for (int i = 0; i < seqLen; ++i) {
|
|
for (int j = 0; j < seqLen; ++j) {
|
|
if (attentionMaskType == "float") {
|
|
attentionMaskData[i * seqLen + j] = (j > i) ? std::numeric_limits<float>::lowest() : 0.0f;
|
|
} else {
|
|
attentionMaskData[i * seqLen + j] = 1.0f;
|
|
}
|
|
}
|
|
}
|
|
inputs.push_back(attentionMask);
|
|
|
|
MNN::Express::VARP positionIds = MNN::Express::_Input({1, seqLen}, MNN::Express::NCHW, halide_type_of<int>());
|
|
int* positionIdsData = positionIds->writeMap<int>();
|
|
for (int i = 0; i < seqLen; ++i) {
|
|
positionIdsData[i] = i;
|
|
}
|
|
inputs.push_back(positionIds);
|
|
}
|
|
|
|
static bool isEmbeddingModel(const rapidjson::Document& doc) {
|
|
if (doc.HasMember("output_names") && doc["output_names"].IsArray()) {
|
|
for (auto iter = doc["output_names"].Begin(); iter != doc["output_names"].End(); ++iter) {
|
|
if (iter->IsString() && std::string(iter->GetString()) == "sentence_embeddings") {
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
auto modelType = std::string(doc.HasMember("model_type") && doc["model_type"].IsString() ? doc["model_type"].GetString() : "");
|
|
if (modelType == "bert" || modelType == "new" || modelType == "qwen3") {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static bool generateForModel(const std::string& modelPath, const std::string& outputDir, const std::string& jsonPath, int blockSize) {
|
|
std::shared_ptr<MNN::Express::Module> net;
|
|
std::vector<std::string> inputNames;
|
|
std::vector<std::string> outputNames;
|
|
bool isEmbedding = false;
|
|
|
|
int hiddenSize;
|
|
std::string attentionMaskType;
|
|
{
|
|
std::ifstream ifs(jsonPath);
|
|
if (!ifs.is_open()) {
|
|
MNN_ERROR("Failed to open JSON config file: %s.\n", jsonPath.c_str());
|
|
return false;
|
|
}
|
|
rapidjson::IStreamWrapper isw(ifs);
|
|
rapidjson::Document doc;
|
|
doc.ParseStream(isw);
|
|
|
|
if (doc.HasParseError() || !doc.IsObject()) {
|
|
MNN_ERROR("Failed to parse JSON config file: %s.\n", jsonPath.c_str());
|
|
return false;
|
|
}
|
|
|
|
if (!doc.HasMember("hidden_size") || !doc["hidden_size"].IsInt()) {
|
|
MNN_ERROR("'hidden_size' not found or not an integer in %s\n", jsonPath.c_str());
|
|
return false;
|
|
}
|
|
hiddenSize = doc["hidden_size"].GetInt();
|
|
|
|
if (!doc.HasMember("attention_mask") || !doc["attention_mask"].IsString()) {
|
|
MNN_ERROR("'attention_mask' not found or not a string in %s\n", jsonPath.c_str());
|
|
return false;
|
|
}
|
|
attentionMaskType = doc["attention_mask"].GetString();
|
|
|
|
isEmbedding = isEmbeddingModel(doc);
|
|
}
|
|
|
|
MNN::ScheduleConfig config;
|
|
std::shared_ptr<MNN::Express::Executor::RuntimeManager> rtmgr(MNN::Express::Executor::RuntimeManager::createRuntimeManager(config));
|
|
rtmgr->setExternalFile((modelPath + ".weight").c_str());
|
|
|
|
if (isEmbedding) {
|
|
inputNames = {"input_ids", "attention_mask", "position_ids"};
|
|
outputNames = {"sentence_embeddings"};
|
|
} else {
|
|
inputNames = {"input_ids", "attention_mask", "position_ids", "logits_index"};
|
|
outputNames = {"logits"};
|
|
}
|
|
net.reset(MNN::Express::Module::load(inputNames, outputNames, modelPath.c_str(), rtmgr), MNN::Express::Module::destroy);
|
|
if (nullptr == net.get()) {
|
|
MNN_ERROR("Failed to load module for QNN IO generation as %s model.\n", isEmbedding ? "embedding" : "llm");
|
|
return false;
|
|
}
|
|
|
|
{
|
|
std::vector<MNN::Express::VARP> inputs;
|
|
std::vector<MNN::Express::VARP> outputs;
|
|
if (isEmbedding) {
|
|
createInputsForEmbedding(blockSize, hiddenSize, attentionMaskType, inputs);
|
|
} else {
|
|
createInputsForLLM(blockSize, hiddenSize, attentionMaskType, false, inputs);
|
|
}
|
|
outputs = net->onForward(inputs);
|
|
if (outputs.empty()) {
|
|
MNN_ERROR("Failed to run forward for QNN IO generation.\n");
|
|
return false;
|
|
}
|
|
saveInputOutputs(net->getInfo(), inputs, outputs, outputDir, blockSize);
|
|
}
|
|
|
|
if (!isEmbedding) {
|
|
std::vector<MNN::Express::VARP> inputs;
|
|
std::vector<MNN::Express::VARP> outputs;
|
|
createInputsForLLM(1, hiddenSize, attentionMaskType, true, inputs);
|
|
outputs = net->onForward(inputs);
|
|
if (outputs.empty()) {
|
|
MNN_ERROR("Failed to run decode forward for QNN IO generation.\n");
|
|
return false;
|
|
}
|
|
saveInputOutputs(net->getInfo(), inputs, outputs, outputDir, 1);
|
|
}
|
|
|
|
if (isEmbedding) {
|
|
std::vector<MNN::Express::VARP> inputs;
|
|
std::vector<MNN::Express::VARP> outputs;
|
|
createInputsForEmbedding(1, hiddenSize, attentionMaskType, inputs);
|
|
outputs = net->onForward(inputs);
|
|
if (outputs.empty()) {
|
|
MNN_ERROR("Failed to run single token embedding forward for QNN IO generation.\n");
|
|
return false;
|
|
}
|
|
saveInputOutputs(net->getInfo(), inputs, outputs, outputDir, 1);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
int main(int argc, char* argv[]) {
|
|
if (argc < 3) {
|
|
MNN_PRINT("Usage: ./generateLlmIO model/config.json outputDir [blocksize]\n");
|
|
MNN_PRINT("This program generates IO test data for QNN export. It supports both generation models and embedding models exported by llmexport.\n");
|
|
return 1;
|
|
}
|
|
|
|
srand(time(NULL));
|
|
int blockSize = 128;
|
|
if (argc >= 4) {
|
|
blockSize = atoi(argv[3]);
|
|
}
|
|
FUNC_PRINT(blockSize);
|
|
|
|
std::string modelPath = std::string(argv[1]) + "/llm.mnn";
|
|
std::string llmConfigPath = std::string(argv[1]) + "/llm_config.json";
|
|
FUNC_PRINT_ALL(modelPath.c_str(), s);
|
|
FUNC_PRINT_ALL(llmConfigPath.c_str(), s);
|
|
std::string outputDir = argv[2];
|
|
|
|
if (!(MNNCreateDir(outputDir.c_str()))) {
|
|
MNN_PRINT("Failed to create dir %s.\n", outputDir.c_str());
|
|
}
|
|
|
|
if (!generateForModel(modelPath, outputDir, llmConfigPath, blockSize)) {
|
|
return 1;
|
|
}
|
|
|
|
return 0;
|
|
}
|