LLM QNN 插件推理流程

按当前仓库代码梳理 MNN LLM 从 response 到 QNN graphExecute 的运行时路径

LLM QNN 插件推理流程

本篇梳理移动设备侧运行 QNN 后端主图 qnn/llm.mnn 时的代码路径。离线编译阶段会把部分子图编译成 QNN 后端使用的 qnn/*.bin 二进制文件,并把主图改写成调用插件算子 OpType_Plugin(type="QNN") 的形式。推理时,核心调用链路如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
[LLM 生成层]
llm->response()
└── generate(input_ids)
    ├── embedding()
    └── forwardVec()
        └── forwardRaw()
[MNN Express Module 层]
Module::onForward()
├── [阶段 1: resize / 形状推导]
│   └── StaticModule::_resize()
│       └── ShapePlugin
│           └── PluginShapeRaw::compute()
│               └── 根据 allInputShape 选择 shapeIndex,并读取 o_<shapeIndex>_* 输出描述
└── [阶段 2: execute / 真正执行]
    └── StaticModule::_execute()
        └── CPUPlugin::onExecute()
            └── PluginExecuteRaw::compute()
                └── RawExecutorWrapper::invokModel()
                    └── QNN graphExecute()

1. config_qnn.json 配置

QNN 模型运行时一般使用 config_qnn.json 配置文件:

1
2
3
4
5
{
  "llm_model": "qnn/llm.mnn",
  "backend_type": "cpu",
  "chunk_limits": [128, 1]
}

LlmConfig::llm_model() 会把 llm_model 拼到模型目录下:

1
2
3
4
// transformers/llm/engine/src/llmconfig.hpp:279
std::string llm_model() const {
    return base_dir_ + config_.value("llm_model", "llm.mnn");
}

所以,最终加载主图路径就是 <model_dir>/qnn/llm.mnn,也就是 Llm::load() 里会用这个路径加载 Module

1
2
3
4
5
6
// transformers/llm/engine/src/llm.cpp:280
std::string model_path = mConfig->llm_model();
...
mRuntimeManager->setExternalFile(mConfig->llm_weight());
mModule.reset(Module::load(inputNames, outputNames, model_path.c_str(), mRuntimeManager, &module_config));
mRuntimeManager->setExternalFile("");

这里加载的 qnn/llm.mnn 是一个带 QNN 插件算子的主图。同时,Llm::setRuntimeHint() 会设置 NPU 文件目录:

1
2
// transformers/llm/engine/src/llm.cpp:153
rtg->setExternalPath(mConfig->npu_model_dir(), MNN::Interpreter::EXTERNAL_NPU_FILE_DIR);

npu_model_dir() 默认返回模型的根目录:

1
2
3
4
// transformers/llm/engine/src/llmconfig.hpp:311
std::string npu_model_dir() const {
    return base_dir_ + config_.value("npu_model_dir", "");
}

这个路径会先写入 RuntimeManager 中:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// transformers/llm/engine/src/llm.cpp:153
rtg->setExternalPath(mConfig->npu_model_dir(), MNN::Interpreter::EXTERNAL_NPU_FILE_DIR);

// express/Executor.cpp:242
void Executor::RuntimeManager::setExternalPath(std::string path, int type) {
    if (type == MNN::Interpreter::EXTERNAL_NPU_FILE_DIR) {
        mInside->mContent->mNpuDir = path;
        return;
    }
    mInside->mContent->modes.setExternalPath(path, type);
}

后面 StaticModule 构造时会把它写到 CPU 后端的 pNPUModelDirPath

1
2
3
// express/module/StaticModule.cpp:340
bnCache.cache.first->pNPUModelDirPath = rtm->getInside()->mContent->mNpuDir;
bnCache.cache.second->pNPUModelDirPath = rtm->getInside()->mContent->mNpuDir;

并在 QNN 插件读取 path 属性时,使用该目录拼出完整 .bin 路径,例如:

1
<model_dir>/qnn/graph0.bin

2. LLM 生成链路

文本输入最常见的入口是:

1
llm->response(prompt);

对应源码里,response(string) 会先按配置套 chat template,再使用 tokenizer 分词为 token id,最后转到 response(input_ids) 方法:

1
2
3
4
5
6
7
8
9
// transformers/llm/engine/src/llm.cpp:912
void Llm::response(const std::string& user_content, std::ostream* os, const char* end_with, int max_new_tokens) {
    auto prompt = user_content;
    if (mConfig->use_template()) {
        prompt = mPrompt->applyTemplate(user_content, true);
    }
    std::vector<int> input_ids = tokenizer_encode(prompt);
    response(input_ids, os, end_with, max_new_tokens);
}

response(input_ids) 方法中,会初始化一些记录的数据,随后继续走 generate(input_ids) 接口生成输出:

1
2
3
4
5
6
// transformers/llm/engine/src/llm.cpp:900
void Llm::response(const std::vector<int>& input_ids, std::ostream* os, const char* end_with, int max_new_tokens) {
    if (!end_with) { end_with = "\n"; }
    generate_init(os, end_with);
    generate(input_ids, max_new_tokens);
}

generate(input_ids) 里会根据 prompt 长度决定是一次 embedding,还是按 mBlockSize 分块:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
// transformers/llm/engine/src/llm.cpp:786
mContext->history_tokens.insert(mContext->history_tokens.end(), input_ids.begin(), input_ids.end());
if(!passExecute) {
    if (0 == mBlockSize || input_ids.size() <= mBlockSize) {
        auto hidden_states = embedding(input_ids);
        return generate(hidden_states, max_tokens);
    }
    int total_size = (int)input_ids.size();
    int loop_size = UP_DIV(total_size, mBlockSize);
    for (int i = 0; i < loop_size; i++) {
        ...
        auto input_embeds = embedding(chunk_ids);
        generate(input_embeds, 0);
    }
}
generate(max_tokens);

3. prefill / decode 形状

Llm::set_config() 读取 chunk_limits 后,会排序并把最大值作为 mBlockSize,代码块如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// transformers/llm/engine/src/llm.cpp:104
if (mConfig->config_.document.HasMember("chunk_limits")) {
    auto& size_limit = mConfig->config_.document["chunk_limits"];
    ...
    for (auto iter = size_limit.GetArray().begin(); iter != size_limit.GetArray().end(); iter++) {
        mValidBlockSize.emplace_back(iter->GetInt());
    }
    ...
    std::sort(mValidBlockSize.begin(), mValidBlockSize.end());
    mBlockSize = mValidBlockSize[mValidBlockSize.size()-1];
}

chunk_limits: [128, 1] 为例,内部结果是:

1
2
mValidBlockSize = [1, 128]
mBlockSize = 128

这和离线导出阶段的两组形状对上:

阶段 seq len 作用
prefill 128 prompt 分块或 pad 后的大块前向
decode 1 自回归单 token 前向

forwardVec(input_embeds) 会按 mBlockSize 拆分长 prompt

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
// transformers/llm/engine/src/llm.cpp:572
std::vector<VARP> Llm::forwardVec(MNN::Express::VARP input_embeds) {
    int seq_len = input_embeds->getInfo()->dim[mSeqLenIndex];
    if (0 == mBlockSize) {
        mMeta->add = seq_len;
        auto attention_mask = gen_attention_mask(seq_len);
        auto position_ids = gen_position_ids(seq_len);
        auto res = forwardRaw(input_embeds, attention_mask, position_ids);
        return res;
    }
    auto blockNumber = seq_len / mBlockSize;
    auto blockRemain = seq_len % mBlockSize;

整块部分直接按 mBlockSizeforwardRaw()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// transformers/llm/engine/src/llm.cpp:604
for (int i=0; i<blockNumber; ++i) {
    logits.clear();
    mMeta->add = blockSize;
    auto embed = embeddings[i];
    auto attention_mask = gen_attention_mask(blockSize);
    auto position_ids = gen_position_ids(blockSize);
    logits = forwardRaw(embed, attention_mask, position_ids);
    ...
    updateContext(blockSize, 0);
}

剩余长度如果不是合法形状,会 padmValidBlockSize 中最近的合法大小:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
// transformers/llm/engine/src/llm.cpp:623
if (!mValidBlockSize.empty()) {
    forwardSize = mValidBlockSize[mValidBlockSize.size()-1];
    for (int j=mValidBlockSize.size()-2; j>=0; --j) {
        if (mValidBlockSize[j] < blockRemain) {
            break;
        }
        forwardSize = mValidBlockSize[j];
    }
    if (blockRemain < forwardSize) {
        hasPad = true;
        auto dim = input_embeds->getInfo()->dim;
        dim[mSeqLenIndex] = forwardSize;
        auto newEmbed = _Input(dim, NCHW);
        ::memcpy(newEmbed->writeMap<void>(), input_embeds->readMap<void>(), input_embeds->getInfo()->size * sizeof(float));
        ::memset(newEmbed->writeMap<float>() + input_embeds->getInfo()->size, 0, (newEmbed->getInfo()->size - input_embeds->getInfo()->size) * sizeof(float));
        input_embeds = newEmbed;
    }
}

比如剩余 prompt 长度是 80,会 pad128decode 阶段长度是 1,直接匹配 1

这个形状需要匹配离线导出时 compilefornpu 写入插件属性的 allInputShape

4. 算子前向

forwardRaw() 构造主模型输入:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
// transformers/llm/engine/src/llm.cpp:448
std::vector<Express::VARP> Llm::forwardRaw(Express::VARP hiddenState, Express::VARP mask, Express::VARP inputPos, Express::VARPS extraArgs) {
    Express::VARP logitsIndex;
    bool inDecode = mContext->gen_seq_len > 0;
    bool isAllLogists = mConfig->all_logits() ? true : (inDecode ? mInSpec : false);
    auto seqLen = hiddenState->getInfo()->dim[mSeqLenIndex];
    int seqLenKey = inDecode ? hiddenState->getInfo()->dim[mSeqLenIndex] : mPrefillKey;
    ...
std::vector<Express::VARP> inputs {
    hiddenState,
    mask,
    inputPos,
    logitsIndex
};

然后选择 Module 模块:

1
2
3
4
// transformers/llm/engine/src/llm.cpp:480
std::vector<Express::VARP> inputs {hiddenState, mask, inputPos, logitsIndex};
inputs.insert(inputs.end(), extraArgs.begin(), extraArgs.end());
std::vector<Express::VARP> outputs = selectModule->onForward(inputs);

这里的 selectModule 通常有两种:

  • prefill 走主 mModule
  • decode 会按 (seqLenKey, isAllLogits)mModulePool 里取 clone 出来的模块。

对应代码里,默认先使用 mModule,在没有 mValidBlockSize 的路径下会按输入长度等 key 复制模块:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// transformers/llm/engine/src/llm.cpp:455
auto moduleKey = std::make_pair(seqLenKey, isAllLogists);
std::shared_ptr<Module> selectModule = mModule;
if (mValidBlockSize.empty()) {
    if(mModulePool.find(moduleKey) == mModulePool.end()) {
        MNN_PRINT("Warning: module need new clone, cloning now.\n");
        mRuntimeManager->setHintPtr(Interpreter::KVCACHE_INFO, mMeta.get());
        mModulePool[moduleKey].reset(Module::clone(mModule.get()));
    }
    selectModule = mModulePool[moduleKey];
}

进入 Module::onForward() 后,实际对象通常是 StaticModuleStaticModule::onForward() 的核心顺序是:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// express/module/StaticModule.cpp:566
std::vector<Express::VARP> StaticModule::onForward(const std::vector<Express::VARP>& inputs) {
    ...
    Variable::compute(inputs);
    ...
    if (runResize) {
        code = _resize(inputs);
    }
    if (NO_ERROR == code && runCompute) {
        code = _execute();
    }

_resize() 会把输入 VARP 绑定到 Session 输入张量,并触发 mSession->resize(),从而向下执行 Pipeline::encode()Pipeline::allocMemory(),完成创建算子 Execution 类的准备。对于 QNN 图来说,这里是在准备 OpType_Plugin 算子:

1
2
3
4
5
6
7
8
9
// express/module/StaticModule.cpp:521
for (int i = 0; i < inputs.size(); ++i) {
    ...
    if (_resizeTensor(mInputTensors[i], inputTensor, mSession.get(), nullptr)) {
        mSession->setNeedResize();
    }
}
mSession->getInfo(Interpreter::RESIZE_STATUS, &curStatus);
code = mSession->resize();

最终 _execute() 会调用 Sessionrun() 方法,执行流水线中的算子。这里使用 QNN 图时,对应的是 OpType_Plugin 算子的执行:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// express/module/StaticModule.cpp:550
ErrorCode StaticModule::_execute() {
    ErrorCode code;
    if (mResource->mModes.callBackMode == Interpreter::Session_Debug) {
        ...
        code = mSession->runWithCallBack(debug->before, debug->after);
    } else {
        code = mSession->run();
    }
    return code;
}

所以 QNN 插件会经历两个不同阶段:

  • resize 阶段:走 ShapePlugin -> PluginShapeRaw::compute(),确定输出形状;
  • execute 阶段:走 CPUPlugin::onExecute() -> PluginExecuteRaw::compute(),真正调用 QNN 计算图。

5. resize 阶段

上一节已经走到 StaticModule::_resize()

1
2
3
4
5
6
StaticModule::onForward()
└── _resize(inputs)
    └── mSession->resize()
        └── Pipeline::encode()
            ├── ShapePlugin / PluginShapeRaw::compute()
            └── CPUPluginCreator -> CPUPlugin -> PluginExecuteRaw::init()

这里要分清两件事:

  • ShapePlugin 负责输出形状推导;
  • CPUPluginCreator 负责创建真正参与后续 executeExecution

这两件事都发生在 resize 这条链路上,所以第 5 节先看形状推导,再看 QNN 内核的注册和创建。

5.1 ShapePlugin

通用插件形状入口在 source/shape/ShapePlugin.cpp。它会从 Plugin 参数里读出 type,再用这个 type 找到对应的形状推导内核:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
// source/shape/ShapePlugin.cpp:35
const Plugin* plugin_param = op->main_as<Plugin>();
std::shared_ptr<plugin::InferShapeKernel> kernel =
    getInferShapeKernel(plugin_param->type()->str());
...
plugin::InferShapeContext ctx(inputs, outputs);
for (const Attribute* attr : *(plugin_param->attr())) {
    ctx.setAttr(attr->key()->str(), attr);
}
bool status = kernel->compute(&ctx);

QNN 插件的 type"QNN"。这个 type 对应的形状推导内核和计算内核都在 QNNBackend.cpp 里注册:

1
2
3
4
5
6
7
8
9
// source/backend/qnn/backend/QNNBackend.cpp:1837
#ifdef MNN_WITH_PLUGIN
plugin::InferShapeKernelRegister::add("QNN", []() {
    return new plugin::shape_inference::PluginShapeRaw;
});
plugin::ComputeKernelRegistry<plugin::backend::PluginExecuteRaw::KernelT>::add("QNN", []() {
    return new plugin::backend::PluginExecuteRaw;
});
#endif

这里注册了两条不同阶段的入口:

注册项 阶段 作用
InferShapeKernelRegister::add("QNN", ...) resize / 形状推导 allInputShape 匹配当前输入形状,并写出输出形状
ComputeKernelRegistry::add("QNN", ...) resize / execute 创建 PluginExecuteRaw,后面加载 QNN 上下文二进制并执行计算图

QNN 注册的 InferShapeKernelRegisterPluginShapeRaw::compute(),它会先调用 computeIndex()。这一步的输入是当前插件算子的输入张量形状:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
// source/backend/qnn/backend/QNNBackend.cpp:361
static bool computeIndex(PluginContext* ctx, int & index) {
    const std::vector<Tensor *> & inputs = ctx->inputs();
    auto attrAllShape = ctx->getAttr("allInputShape");
    if (nullptr == attrAllShape || nullptr == attrAllShape->list() || nullptr == attrAllShape->list()->i()) {
        MNN_ERROR("MNN_QNN: Incorrect Plugin Op, can't find 'allInputShape' attr.\n");
        return false;
    }
    int dimSum = 0;
    for (int i = 0; i < inputs.size(); i++) {
        auto inputDim = inputs[i]->dimensions();
        dimSum += inputDim;
    }
    ...
}

随后与离线 compilefornpu 写进插件属性的多组合法输入 allInputShape 对比,找到匹配当前输入形状的 shapeIndex

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// source/backend/qnn/backend/QNNBackend.cpp:378
auto indexNumber = attrAllShape->list()->i()->size() / dimSum;
for (int si=0; si<indexNumber; ++si) {
    auto dstSi = attrAllShape->list()->i()->data() + si * dimSum;
    bool valid = true;
    for (int i=0; i<inputs.size(); ++i) {
        auto inputDim = inputs[i]->dimensions();
        for (int j = 0; j < inputDim; j++) {
            if (inputs[i]->length(j) != dstSi[j]) {
                valid = false;
                break;
            }
        }
        dstSi += inputDim;
        if (!valid) {
            break;
        }
    }
    if (valid) {
        index = si;
        return true;
    }
}

匹配成功后得到的 shapeIndex 是后面 QNN 路径的关键索引:

1
2
3
4
shapeIndex
├── resize 阶段: 读取 o_<shapeIndex>_<outputIndex> 写输出形状
├── PluginExecuteRaw::resize(): 记录 mShapeIndex
└── execute 阶段: 选择 mQnnGraphHandleVec[mShapeIndex]

因此,运行时输入形状必须能在 allInputShape 里找到,否则 QNN 插件的形状推导会失败。

当匹配到 shapeIndex 后,PluginShapeRaw::compute() 会读取 o_<shapeIndex>_<outputIndex>,把离线记录的输出类型、形状和格式写回当前输出张量:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// source/backend/qnn/backend/QNNBackend.cpp:403
bool PluginShapeRaw::compute(InferShapeContext* ctx) {
    ...
    int shapeIndex = 0;
    if (!(computeIndex(ctx, shapeIndex))) {
        MNN_ERROR("MNN_QNN: Failed to compute shape for Plugin Op.\n");
        return false;
    }

    std::string prefix = "o_" + std::to_string(shapeIndex) + "_";
    for (int i=0; i<ctx->outputs().size(); ++i) {
        auto dst = ctx->output(i);
        std::string key = prefix + std::to_string(i);
        auto attr = ctx->getAttr(key.c_str());
        ...
        auto blob = attr->tensor();
        dst->setType(blob->dataType());
        if (nullptr != blob->dims()) {
            dst->buffer().dimensions = blob->dims()->size();
            for (int j=0; j<blob->dims()->size(); ++j) {
                dst->setLength(j, blob->dims()->data()[j]);
            }
        } else {
            dst->buffer().dimensions = 0;
        }
        TensorUtils::getDescribe(dst)->dimensionFormat = blob->dataFormat();
    }
    return true;
}

5.2 创建 QNN 算子

形状推导完成后,Pipeline::encode() 还会为每个算子创建后端 Execution。这里不是直接手写调用 new CPUPlugin(...),而是走 MNN 后端的通用算子创建路径:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
Pipeline::encode()
└── OpCommonUtils::createExecutionWithExternal(...)
    └── backend->onCreate(inputs, outputs, op)
        └── CPUBackend::onCreate(...)
            └── gCreator[OpType_Plugin]->onCreate(...)
                └── CPUPluginCreator::onCreate(...)
                    └── new CPUPlugin(ctx)
                        └── getCPUComputeKernel("QNN")
                            └── ComputeKernelRegistry<CPUComputeKernel>::get("QNN")
                                └── new PluginExecuteRaw

Pipeline::encode() 里,当前 command 还没有 Execution 时,会调用 createExecutionWithExternal(...)

1
2
3
4
5
6
// source/core/Pipeline.cpp:563
std::shared_ptr<BufferStorage> tmpStorage;
if (nullptr == iter.execution) {
    iter.execution.reset(OpCommonUtils::createExecutionWithExternal(
        mBackend.get(), iter.inputs, iter.outputs, iter.op, &loader, tmpStorage));
}

OpCommonUtils::createExecutionWithExternal(...) 对没有外部权重重建需求的算子,会直接转到后端的 onCreate(...)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
// source/core/OpCommonUtils.cpp:664
Execution* OpCommonUtils::createExecutionWithExternal(
    Backend* backend,
    const std::vector<Tensor*>& inputs,
    const std::vector<Tensor*>& outputs,
    const MNN::Op* op,
    FileLoader* externalFile,
    std::shared_ptr<BufferStorage>& tmpstore) {
    ...
    if (!hasExternal) {
        return backend->onCreate(inputs, outputs, op);
    }

当前模型运行在 CPU 后端上,所以这里进入 CPUBackend::onCreate(...)。它用 op->type()gCreator 表里找创建器:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
// source/backend/cpu/CPUBackend.cpp:720
Execution* CPUBackend::onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                                const MNN::Op* op) {
    ...
    auto opType = op->type();
    ...
    auto map  = gCreator;
    auto iter = map->find(opType);
    if (iter == map->end() ) {
        MNN_PRINT("Don't support type [%s]\n", MNN::EnumNameOpType(op->type()));
        return nullptr;
    }
    Execution* exe = nullptr;
    if (exe == nullptr) {
        exe = iter->second->onCreate(inputs, outputs, op, this);
    }
    return exe;
}

OpType_Plugin 对应的创建器在 CPUPlugin.cpp 文件尾部注册:

1
2
// source/backend/cpu/CPUPlugin.cpp:98
REGISTER_CPU_OP_CREATOR(CPUPluginCreator, OpType_Plugin);

这个宏展开后的核心动作是把 CPUPluginCreator 放进 CPUBackend::gCreator

1
2
3
4
5
6
// source/backend/cpu/CPUBackend.hpp:218
#define REGISTER_CPU_OP_CREATOR(name, opType)     \
    void ___##name##__##opType##__() {            \
        static name _temp;\
        CPUBackend::addCreator(opType, &_temp); \
    }

所以对于 OpType_Plugin(type="QNN"),最终会调用 CPU 后端里的 CPUPluginCreator::onCreate(...)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
// source/backend/cpu/CPUPlugin.cpp:70
class CPUPluginCreator : public CPUBackend::Creator {
public:
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs,
                                const std::vector<Tensor*>& outputs,
                                const MNN::Op* op, Backend* backend) const {
        ...
        const Plugin* plugin_param = op->main_as<Plugin>();
        const std::string& op_type = plugin_param->type()->str();
        std::unique_ptr<plugin::CPUKernelContext> ctx(
            new plugin::CPUKernelContext(op_type, backend, inputs, outputs, static_cast<CPUBackend *>(backend)->pNPUModelDirPath));

        for (const Attribute* attr : *(plugin_param->attr())) {
            ctx->setAttr(attr->key()->str(), attr);
        }
        return new CPUPlugin(std::move(ctx));
    }
};

这里把三类信息塞进了 CPUKernelContext

  • op_type:这里是 "QNN"
  • pNPUModelDirPath:模型目录,后面用来拼 QNN 计算图 qnn/graph0.bin
  • 插件属性:包括 pathinputsoutputsallGraphNameallInputShape

CPUPluginCreator::onCreate(...) 最后一行真正进入 CPUPlugin 构造函数。构造函数里会通过 op_type 找到 QNN 计算内核,并立刻调用 kernel_->init(...)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// source/backend/cpu/CPUPlugin.cpp:27
class CPUPlugin : public Execution {
public:
    CPUPlugin(std::unique_ptr<plugin::CPUKernelContext> ctx)
        : Execution(ctx->backend()), ctx_(std::move(ctx)) {
        kernel_ = getCPUComputeKernel(ctx_->op_type());
        MNN_CHECK(nullptr != kernel_.get(),
                  "CPU compute kernel has not been registered for plugin op.");
        kernel_->init(ctx_.get());
        mNeedAllocIO = kernel_->needAllocIO();
    }

getCPUComputeKernel(...) 本身只是从 ComputeKernelRegistry 里按名字取工厂函数:

1
2
3
4
5
6
// source/backend/cpu/CPUPlugin.cpp:21
static std::shared_ptr<plugin::CPUComputeKernel> getCPUComputeKernel(
    const std::string& name) {
    return std::shared_ptr<plugin::CPUComputeKernel>(
        plugin::ComputeKernelRegistry<plugin::CPUComputeKernel>::get(name));
}

get(name) 会查 gFactoryMap,然后执行对应工厂函数:

1
2
3
4
5
6
7
8
9
// source/plugin/PluginKernel.cpp:48
PluginKernelT* ComputeKernelRegistry<PluginKernelT>::get(const std::string& name) {
    auto* gFactoryMap = ComputeKernelRegistry<PluginKernelT>::getFactoryMap();
    if (!gFactoryMap->count(name)) {
        MNN_PRINT("Factory has not been registered for name %s.", name.c_str());
        return nullptr;
    }
    return gFactoryMap->at(name)();
}

QNN 在前面注册过 "QNN" 这个工厂函数:

1
2
3
4
// source/backend/qnn/backend/QNNBackend.cpp:1841
plugin::ComputeKernelRegistry<plugin::backend::PluginExecuteRaw::KernelT>::add("QNN", []() {
    return new plugin::backend::PluginExecuteRaw;
});

所以这里的实际对象就是 PluginExecuteRaw。因此 CPUPlugin 构造函数里的 kernel_->init(ctx_.get()) 会动态分发到 PluginExecuteRaw::init(...),执行加载 qnn/graph*.bin

5.3 加载 QNN

PluginExecuteRaw::init() 先确保 QNN 全局上下文已经创建,然后用模型目录和插件属性 path 拼出 QNN 图的上下文二进制路径:

1
2
3
4
5
6
// source/backend/qnn/backend/QNNBackend.cpp:907
bool init(CPUKernelContext* ctx) override {
    if (QNN::gContext.deviceHandle == nullptr){
        QNN::createQnnContext();
    }
    auto path = MNNFilePathConcat(ctx->dir_path(), ctx->getAttr("path")->s()->str());

这里的路径关系是:

1
2
3
ctx->dir_path()              = <model_dir>
ctx->getAttr("path")->s()    = qnn/graph0.bin
最终 path                    = <model_dir>/qnn/graph0.bin

随后读取 allGraphName。一个 .bin 里可能包含多个计算图,比如 prefilldecode,所以运行时必须知道计算图名字顺序:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
// source/backend/qnn/backend/QNNBackend.cpp:913
std::vector<std::string> allGraphName;
auto allGraphNameAttr = ctx->getAttr("allGraphName");
if (allGraphNameAttr && allGraphNameAttr->list() && allGraphNameAttr->list()->s()) {
    auto graphNames = allGraphNameAttr->list()->s();
    for (int i = 0; i < graphNames->size(); ++i) {
        allGraphName.push_back(graphNames->GetAsString(i)->str());
    }
} else {
    MNN_ERROR("MNN_QNN: Incorrect Plugin Op, can't find 'allGraphName' attr.\n");
    return false;
}

最后创建 RawExecutorWrapper,把上下文二进制交给 compileModel(...)

1
2
3
// source/backend/qnn/backend/QNNBackend.cpp:952
mRawExecutor.reset(new RawExecutorWrapper());
return mRawExecutor->compileModel(path, binaryOffset, binarySize, allGraphName);

compileModel(...) 先读入 .bin。如果插件属性里有 offset / size,就从文件指定位置读取;否则直接 mmap 整个文件:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
// source/backend/qnn/backend/QNNBackend.cpp:772
bool compileModel(const std::string& path, size_t offset, size_t size, const std::vector<std::string>& allGraphName) {
    mPath = path;
    void* buffer = nullptr;
    std::vector<char> bufferVec(size, 0);
    MMapReader reader;
    if (size > 0) {
        buffer = bufferVec.data();
        std::unique_ptr<FileLoader> binaryFile(new FileLoader(path.c_str()));
        binaryFile->offset((int64_t)offset);
        binaryFile->read((char *)buffer, (int64_t)size);
    } else {
        reader.open(path.c_str());
        buffer = reader.addr();
        size = reader.size();
    }

接着通过 QNN system API 解析二进制里的计算图元信息:

1
2
3
4
5
6
7
8
9
// source/backend/qnn/backend/QNNBackend.cpp:790
QnnSystemContext_Handle_t systemContextHandle = nullptr;
QNN::gContext.systemInterface.systemContextCreate(&systemContextHandle);
Qnn_ContextBinarySize_t binarySize = 0;
const QnnSystemContext_BinaryInfo_t* binaryInfo = nullptr;
CALL_QNN(QNN::gContext.systemInterface.systemContextGetBinaryInfo(
    systemContextHandle, buffer, size, &binaryInfo, &binarySize));
copyMetadataToGraphsInfo(binaryInfo, mGraphsInfo, mGraphCount);
QNN::gContext.systemInterface.systemContextFree(systemContextHandle);

然后校验二进制,并从二进制创建 QNN 上下文:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
// source/backend/qnn/backend/QNNBackend.cpp:807
auto error = QNN::gContext.interface.contextValidateBinary(
    QNN::gContext.backendHandle,
    QNN::gContext.deviceHandle,
    mQnnContextConfig,
    buffer,
    size);
...
CALL_QNN(QNN::gContext.interface.contextCreateFromBinary(
    QNN::gContext.backendHandle,
    QNN::gContext.deviceHandle,
    mQnnContextConfig,
    buffer,
    size,
    &mQnnContextHandle,
    mQnnProfileHandle));

最后按 allGraphName 的顺序整理 mGraphsInfo,并取回每个计算图句柄:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
// source/backend/qnn/backend/QNNBackend.cpp:820
std::vector<GraphInfo*> sortedGraphsInfo(mGraphCount, nullptr);
std::map<std::string, GraphInfo*> graphInfoMap;
for (int i = 0; i < mGraphCount; ++i) {
    graphInfoMap[mGraphsInfo[i]->graphName] = mGraphsInfo[i];
}

for (int i = 0; i < mGraphCount; ++i) {
    auto it = graphInfoMap.find(allGraphName[i]);
    MNN_ASSERT(it != graphInfoMap.end());
    sortedGraphsInfo[i] = it->second;
}
for (int i = 0; i < mGraphCount; i++) {
    CALL_QNN(QNN::gContext.interface.graphRetrieve(
        mQnnContextHandle,
        mGraphsInfo[i]->graphName,
        &(mQnnGraphHandleVec[i])));
}

这一步结束后,QNN 插件已经持有:

1
2
3
4
RawExecutorWrapper
├── mQnnContextHandle
├── mGraphsInfo[shapeIndex]
└── mQnnGraphHandleVec[shapeIndex]

后面 execute 阶段不再重新创建 QNN 计算图,而是根据 shapeIndex 选择已经取回的计算图句柄。

5.4 PluginExecuteRaw::resize()

CPUPlugin::onResize() 会把当前输入输出张量重置到 CPUKernelContext,再调用 QNN 内核的 resize()

1
2
3
4
5
6
7
// source/backend/cpu/CPUPlugin.cpp:49
ErrorCode CPUPlugin::onResize(const std::vector<Tensor*>& inputs,
                              const std::vector<Tensor*>& outputs) {
    ctx_->reset(inputs, outputs);
    auto success = kernel_->resize(ctx_.get());
    ...
}

PluginExecuteRaw::resize() 会再次通过 computeIndex() 得到当前 shapeIndex,并保存到成员变量 mShapeIndex

1
2
3
4
5
6
7
8
// source/backend/qnn/backend/QNNBackend.cpp:956
bool resize(CPUKernelContext* ctx) override {
    int shapeIndex = 0;
    if (!(shape_inference::computeIndex(ctx, shapeIndex))) {
        MNN_ERROR("MNN_QNN: Failed to execute Plugin Op.\n");
        return false;
    }
    mShapeIndex = shapeIndex;

随后读取插件属性里的 inputs / outputs。这些名字一般是 t<tensor_index>,由离线 compilefornpu 按主图张量索引写入:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// source/backend/qnn/backend/QNNBackend.cpp:964
auto inputs = ctx->getAttr("inputs")->list();
auto inputTensor = ctx->inputs();
MNN_ASSERT(inputs->s()->size() == inputTensor.size());
mInputs.resize(inputs->s()->size());
mRealInputs.resize(inputTensor.size());
for (int i=0; i<inputs->s()->size(); ++i) {
    mRealInputs[i].reset(new Tensor(inputTensor[i], Tensor::CAFFE));
    mInputs[i].second = inputs->s()->GetAsString(i)->str();
    mInputs[i].first = mRealInputs[i].get();
}

输出也是同样处理:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// source/backend/qnn/backend/QNNBackend.cpp:974
auto outputs = ctx->getAttr("outputs")->list();
auto outputTensor = ctx->outputs();
mOutputs.resize(outputs->s()->size());
MNN_ASSERT(outputs->s()->size() == outputTensor.size());
mRealOutputs.resize(outputTensor.size());
for (int i=0; i<outputs->s()->size(); ++i) {
    mRealOutputs[i].reset(new Tensor(outputTensor[i], Tensor::CAFFE));
    mOutputs[i].second = outputs->s()->GetAsString(i)->str();
    mOutputs[i].first = mRealOutputs[i].get();
}

这里保存的是 MNN 张量和 QNN 张量名的对应关系:

1
2
3
mRealInputs / mRealOutputs: 临时 Tensor,后面交给 QNN graphExecute 使用
mInputs / mOutputs: (临时 Tensor 指针, QNN tensor name)
mShapeIndex: 当前输入 shape 对应的 graph index

到这里,resize 阶段完成了三件事:输出形状已经确定,QNN 上下文二进制已经加载,当前形状的输入输出映射也准备好了。

6. execute 阶段

StaticModule::_execute() 最后会走到 Session::run()。当流水线执行到 OpType_Plugin(type="QNN") 时,CPU 后端调用的是 CPUPlugin::onExecute()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
// source/backend/cpu/CPUPlugin.cpp:59
ErrorCode CPUPlugin::onExecute(const std::vector<Tensor*>& inputs,
                               const std::vector<Tensor*>& outputs) {
    if (kernel_->compute(ctx_.get())) {
        return NO_ERROR;
    } else {
        MNN_ERROR("Plugin kernel compute failed with false returned.\n");
        return INVALID_VALUE;
    }
}

QNN 插件来说,这里的 kernel_ 就是前面 init 阶段创建好的 PluginExecuteRaw

6.1 重新确认 shapeIndex

PluginExecuteRaw::compute() 开头会重新调用 computeIndex(),确认当前输入形状对应哪个计算图:

1
2
3
4
5
6
7
8
9
// source/backend/qnn/backend/QNNBackend.cpp:988
bool compute(CPUKernelContext* ctx) override {
    AUTOTIME;
    int shapeIndex = 0;
    if (!(shape_inference::computeIndex(ctx, shapeIndex))) {
        MNN_ERROR("MNN_QNN: Failed to execute Plugin Op.\n");
        return false;
    }
    std::string graphName = ctx->getAttr("allGraphName")->list()->s()->GetAsString(shapeIndex)->str();

当前源码后面传给 invokModel() 的是成员变量 mShapeIndex,也就是 resize 阶段记录的 shape indexcompute() 里重新算出的局部 shapeIndex 主要用于读取 allGraphName 和调试打印。

6.2 准备输入输出 Tensor

compute() 里会重新创建一次实时输入输出张量映射。源码注释写的是 in-time inputs and outputs tensor

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
// source/backend/qnn/backend/QNNBackend.cpp:997
// compute and alloc real in-time inputs and outputs tensor
{
    auto inputs = ctx->getAttr("inputs")->list();
    auto inputTensor = ctx->inputs();
    MNN_ASSERT(inputs->s()->size() == inputTensor.size());
    mInputs.resize(inputs->s()->size());
    mRealInputs.resize(inputTensor.size());
    for (int i=0; i<inputs->s()->size(); ++i) {
        mRealInputs[i].reset(new Tensor(inputTensor[i], Tensor::CAFFE));
        mInputs[i].second = inputs->s()->GetAsString(i)->str();
        mInputs[i].first = mRealInputs[i].get();
    }
    ...
}

这和 resize 阶段的映射逻辑有重复。resize 阶段先建立一次映射并记录 mShapeIndexexecute 阶段再基于当前 ctx->inputs() / ctx->outputs() 建立实时 Tensor,保证后面的拷贝和 QNNclientBuf 绑定使用的是本次运行的缓冲区。

6.3 拷贝输入、执行 QNN、拷回输出

真正执行前,MNN 会先把当前输入张量拷到 mRealInputs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
// source/backend/qnn/backend/QNNBackend.cpp:1024
auto inputTensor = ctx->inputs();
auto outputTensor = ctx->outputs();

for (int i=0; i<mInputs.size(); ++i) {
    ctx->backend()->onCopyBuffer(inputTensor[i], mRealInputs[i].get());
}
mRawExecutor->invokModel(mInputs, mOutputs, mShapeIndex);
for (int i=0; i<mOutputs.size(); ++i) {
    ctx->backend()->onCopyBuffer(mRealOutputs[i].get(), outputTensor[i]);
}

随后在 invokModel() 里执行,最后把结果从 mRealOutputs 拷回 ctx->outputs()

6.4 invokModel()执行入口

RawExecutorWrapper::invokModel() 先根据 shapeIndex 选择 QNN 计算图:

1
2
3
4
// source/backend/qnn/backend/QNNBackend.cpp:844
void invokModel(const std::vector<std::pair<const MNN::Tensor *, std::string>>& inputs, std::vector<std::pair<const MNN::Tensor *, std::string>>& outputs, int shapeIndex) {
    GraphInfo* graph = mGraphsInfo[shapeIndex];
    Qnn_GraphHandle_t qnnGraphHandle = mQnnGraphHandleVec[shapeIndex];

然后按张量名把 MNN 缓冲区绑定到 QNN 输入张量的 clientBuf

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
// source/backend/qnn/backend/QNNBackend.cpp:849
for (int i=0; i<inputs.size(); ++i) {
    auto t = inputs[i].first;
    bool find = false;
    for (int j=0; j<graph->numInputTensors; ++j) {
        auto& dstT = graph->inputTensors[j];
        if (inputs[i].second == dstT.v1.name) {
            dstT.v1.clientBuf.data = t->host<void>();
            dstT.v1.clientBuf.dataSize = t->usize();
            find = true;
            break;
        }
    }
}

输出张量也是同样按名字绑定:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
// source/backend/qnn/backend/QNNBackend.cpp:868
for (int i=0; i<outputs.size(); ++i) {
    auto t = outputs[i].first;
    bool find = false;
    for (int j=0; j<graph->numOutputTensors; ++j) {
        auto& dstT = graph->outputTensors[j];
        if (outputs[i].second == dstT.v1.name) {
            dstT.v1.clientBuf.data = t->host<void>();
            dstT.v1.clientBuf.dataSize = t->usize();
            find = true;
            break;
        }
    }
}

这里的名字来自离线阶段写进插件属性的 inputs / outputs

6.5 graphExecute() 执行点

绑定完输入输出缓冲区后,RawExecutorWrapper::invokModel() 会调用 QNN 接口 graphExecute() 执行计算图:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
// source/backend/qnn/backend/QNNBackend.cpp:887
CALL_QNN(QNN::gContext.interface.graphExecute(
    qnnGraphHandle,
    graph->inputTensors,
    graph->numInputTensors,
    graph->outputTensors,
    graph->numOutputTensors,
    mQnnProfileHandle,
    nullptr));
MNN::QNN::doProfile(QNN::gContext.interface, mQnnProfileHandle);

到这一行,计算已经交给了 QNN 后端:

  • 通过 shapeIndex 选中了正确的 QNN 计算图句柄;
  • 通过 inputs / outputs 名字把缓冲区绑到 QNN 张量;
  • 把执行交给 QNN 运行时。

.bin 里面的 HTP 计算图执行由 QNN 后端负责。MNN 运行时在这里主要做计算图选择、缓冲区绑定和结果回拷。

7. 小结

把第 4 到第 6 节串起来看,QNN 插件的运行时路径是:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
StaticModule::onForward()
├── _resize()
│   ├── ShapePlugin
│   │   └── PluginShapeRaw::compute()
│   │       ├── allInputShape -> shapeIndex
│   │       └── o_<shapeIndex>_* -> 输出 shape
│   │
│   └── Pipeline::encode() 创建 Plugin Execution
│       └── CPUPluginCreator
│           └── CPUPlugin
│               └── PluginExecuteRaw::init()
│                   └── RawExecutorWrapper::compileModel()
│                       ├── mmap/read qnn/graph*.bin
│                       ├── contextCreateFromBinary()
│                       └── graphRetrieve()
└── _execute()
    └── CPUPlugin::onExecute()
        └── PluginExecuteRaw::compute()
            ├── 准备实时输入输出 Tensor
            ├── onCopyBuffer(input -> mRealInputs)
            ├── RawExecutorWrapper::invokModel(mShapeIndex)
            │   ├── mQnnGraphHandleVec[mShapeIndex]
            │   ├── 绑定 QNN input/output clientBuf
            │   └── graphExecute()
            └── onCopyBuffer(mRealOutputs -> output)