wehub-resource-sync
47 lines
1.4 KiB
C++
47 lines
1.4 KiB
C++
//
|
|
// IfTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/05/10.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(IfTorch);
|
|
|
|
MNN::OpType IfTorch::opType() {
|
|
return MNN::OpType_If;
|
|
}
|
|
MNN::OpParameter IfTorch::type() {
|
|
return MNN::OpParameter_IfParam;
|
|
}
|
|
std::vector<int> IfTorch::inputTensorIdx() {
|
|
return {0};
|
|
}
|
|
|
|
void IfTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::IfParamT;
|
|
const auto blocks = node->blocks();
|
|
param->then_graph = dstOp->name + "/then";
|
|
param->else_graph = dstOp->name + "/else";
|
|
// then
|
|
scope->buildSubGraph(blocks[0], param->then_graph);
|
|
// else
|
|
scope->buildSubGraph(blocks[1], param->else_graph);
|
|
scope->dealSubgraphDepsForOp(dstOp);
|
|
for (int idx : dstOp->inputIndexes) {
|
|
std::unique_ptr<MNN::StringVecT> inputT(new MNN::StringVecT);
|
|
inputT->data.emplace_back(scope->lookupTensorByIdx(idx));
|
|
param->aliases_inputs.emplace_back(std::move(inputT));
|
|
}
|
|
std::unique_ptr<MNN::StringVecT> outputPari(new MNN::StringVecT);
|
|
outputPari->data.emplace_back(blocks[0]->outputs()[0]->debugName());
|
|
outputPari->data.emplace_back(blocks[1]->outputs()[0]->debugName());
|
|
param->aliases_outputs.emplace_back(std::move(outputPari));
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(IfTorch, If);
|