1 #include "nomnigraph/Converters/Caffe2.h" 2 #include "nomnigraph/Graph/Algorithms.h" 4 #include "nomnigraph/Support/Casting.h" 5 #include "nomnigraph/Support/Pointer.h" 7 std::map<std::string, caffe2::Argument>
8 getArgumentsFromOperator(caffe2::OperatorDef op) {
9 std::map<std::string, caffe2::Argument> argMap;
10 for (
auto arg : op.arg()) {
11 argMap[arg.name()] = arg;
17 namespace converters {
19 std::unique_ptr<repr::NeuralNetOperator>
20 convertOperator(caffe2::OperatorDef op) {
21 auto argMap = getArgumentsFromOperator(op);
23 if (op.type() ==
"Conv") {
25 std::vector<int> kernelShape;
26 if (argMap.count(
"kernel")) {
27 assert(argMap[
"kernel"].has_i() &&
"Invalid kernel argument passed to Conv");
28 int kernel =
static_cast<int>(argMap[
"kernel"].i());
29 kernelShape = {kernel, kernel};
30 }
else if (argMap.count(
"kernels")) {
31 for (
auto i : argMap[
"kernels"].ints()) {
32 kernelShape.push_back(static_cast<int>(i));
34 }
else if (argMap.count(
"kernel_h") && argMap.count(
"kernel_w")) {
35 assert(argMap[
"kernel_h"].has_i() &&
"Invalid kernel argument passed to Conv");
36 assert(argMap[
"kernel_w"].has_i() &&
"Invalid kernel argument passed to Conv");
37 int kernelH =
static_cast<int>(argMap[
"kernel_h"].i());
38 int kernelW =
static_cast<int>(argMap[
"kernel_w"].i());
39 kernelShape = {kernelH, kernelW};
44 auto c = util::make_unique<repr::Conv>(kernelShape);
46 if (argMap.count(
"order")) {
47 auto order = argMap[
"order"].s();
48 if (order ==
"NCHW") {
49 c->setLayout(repr::Conv::NNLayout::NCHW);
50 }
else if (order ==
"NHWC") {
51 c->setLayout(repr::Conv::NNLayout::NHWC);
57 if (argMap.count(
"stride")) {
58 assert(argMap[
"stride"].has_i() &&
"Invalid stride argument");
59 int stride =
static_cast<int>(argMap[
"stride"].i());
60 c->setStrides({stride, stride});
63 if (argMap.count(
"pad")) {
64 assert(argMap[
"pad"].has_i() &&
"Invalid pad argument");
65 int pad =
static_cast<int>(argMap[
"pad"].i());
66 c->setPads({pad, pad, pad, pad});
69 if (argMap.count(
"dilation")) {
70 assert(argMap[
"dilation"].has_i() &&
"Invalid dilation argument");
71 int dilation =
static_cast<int>(argMap[
"dilation"].i());
72 c->setDilations({dilation, dilation});
78 if (op.type() ==
"Relu") {
79 auto relu = util::make_unique<repr::Relu>();
80 return std::move(relu);
83 return util::make_unique<repr::GenericOperator>(op.type());
90 repr::NNModule convertFromCaffe2Proto(
const caffe2::NetDef &net, std::unordered_map<std::string, repr::NNGraph::NodeRef>* blobMapOut) {
99 std::unordered_map<std::string, repr::NNGraph::NodeRef> blobMap;
106 cfg.createNode(util::make_unique<repr::BasicBlockType<repr::NNGraph>>());
108 for (
const auto &op : net.op()) {
109 auto opNode = dfg.createNode();
111 for (
const auto &input : op.input()) {
113 if (!blobMap.count(input)) {
114 auto tensor = util::make_unique<repr::Tensor>(input);
116 dfg.createNode(unique_dyn_cast<repr::NeuralNetData>(tensor));
119 auto tensorNode = blobMap[input];
120 dfg.createEdge(tensorNode, opNode);
124 for (
const auto &output : op.output()) {
125 auto tensor = util::make_unique<repr::Tensor>(output);
127 dfg.createNode(unique_dyn_cast<repr::NeuralNetData>(tensor));
128 dfg.createEdge(opNode, tensorNode);
129 blobMap[output] = tensorNode;
132 if (op.type() ==
"While") {
133 opNode->resetData(util::make_unique<repr::While>());
134 auto argMap = getArgumentsFromOperator(op);
135 std::string bodyNetSerialized = argMap[
"body"].s();
136 auto bodyNet = caffe2::NetDef();
137 bodyNet.ParseFromString(bodyNetSerialized);
139 std::unordered_map<std::string, repr::NNGraph::NodeRef> bodyBlobMap;
140 auto bodyNN = convertFromCaffe2Proto(bodyNet, &bodyBlobMap);
141 repr::NNGraph bodyGraph = std::move(bodyNN.dataFlow);
142 repr::NNCFGraph bodyCFGraph = std::move(bodyNN.controlFlow);
144 auto rev_sorted = algorithm::tarjans(&bodyGraph);
146 for (
auto& k : bodyBlobMap) {
148 if (blobMap.count(name)) {
149 auto oldNode = blobMap[name];
150 printf(
"Exit tensor %s is in the parent scope, inserting Phi node...\n", k.first.c_str());
151 auto phiNode = dfg.createNode(util::make_unique<repr::NNPhi>());
153 auto tensor = dyn_cast<repr::NeuralNetData>(blobMap[name]->data().get());
154 auto* clonedTensor = tensor->clone();
155 auto phiOut = dfg.createNode(std::unique_ptr<repr::NeuralNetData>(clonedTensor));
156 dfg.createEdge(phiNode, phiOut);
157 dfg.createEdge(oldNode, phiNode);
158 dfg.createEdge(bodyBlobMap[name], phiNode);
159 blobMap[name] = phiOut;
160 for (
auto& inEdge : opNode->getInEdges()) {
161 if (inEdge->tail() == oldNode) {
162 dfg.deleteEdge(inEdge);
163 dfg.createEdge(phiOut, opNode);
170 std::unordered_map<repr::NNGraph::NodeRef, repr::NNGraph::NodeRef> inNodeMap;
171 for (
auto& n : bodyGraph.getMutableNodes()) {
172 if (!isa<repr::NeuralNetData>(n->data())) {
continue; }
173 if (n->getInEdges().size() == 0) {
174 auto name = dyn_cast<repr::NeuralNetData>(n->data().get())->getName();
177 if (blobMap.count(name)) {
178 inNodeMap[n] = blobMap[name];
183 assert(rev_sorted.front().getNodes().size() == 1 &&
184 "More than one exit node.");
185 assert(rev_sorted.back().getNodes().size() == 1 &&
186 "More than one entry node.");
188 auto exit_tensor = *(rev_sorted.front().getNodes().begin());
189 assert(isa<repr::NeuralNetData>(exit_tensor->data()) &&
190 "Exit node is not a tensor.");
192 auto bodyNodes = bodyGraph.getMutableNodes();
193 auto bodyEdges = bodyGraph.getMutableEdges();
195 for (
auto node : bodyNodes) {
196 bodyGraph.swapNode(node, dfg);
199 for (
auto edge : bodyEdges) {
200 bodyGraph.swapEdge(edge, dfg);
204 for (
auto node : dfg.getMutableNodes()) {
205 if (inNodeMap.count(node)) {
206 dfg.replaceNode(node, inNodeMap[node]);
207 dfg.deleteNode(node);
211 for (
const auto& inEdge : opNode->getInEdges()) {
212 auto* inputData = dyn_cast<repr::NeuralNetData>(inEdge->tail()->data().get());
213 auto* exitData = dyn_cast<repr::NeuralNetData>(exit_tensor->data().get());
214 if (inputData->getName() == exitData->getName()) {
215 dfg.replaceNode(exit_tensor, inEdge->tail());
216 dfg.deleteNode(exit_tensor);
221 auto bodyCFNodes = bodyCFGraph.getMutableNodes();
222 auto bodyCFEdges = bodyCFGraph.getMutableEdges();
225 auto whileBasicBlock = util::make_unique<repr::BasicBlockType<repr::NNGraph>>();
226 for (
auto& inEdge : opNode->getInEdges()) {
227 auto node = inEdge->tail();
228 for (
auto& parentInEdge : node->getInEdges()) {
229 auto parentNode = parentInEdge->tail();
230 if (isa<repr::Phi>(parentNode->data().get())) {
231 whileBasicBlock->pushInstructionNode(parentNode);
235 whileBasicBlock->pushInstructionNode(opNode);
237 auto whileCFNode = cfg.createNode(std::move(whileBasicBlock));
238 cfg.createEdge(bbNode, whileCFNode, 0);
242 for (
auto cfNode : bodyCFNodes) {
243 bodyCFGraph.swapNode(cfNode, cfg);
246 if (cfNode->getOutEdges().size() == 0) {
247 cfg.createEdge(cfNode, whileCFNode, 0);
250 if (cfNode->getInEdges().size() == 0) {
251 cfg.createEdge(whileCFNode, cfNode, 1);
254 for (
auto cfEdge : bodyCFEdges) {
255 bodyCFGraph.swapEdge(cfEdge, cfg);
260 cfg.createNode(util::make_unique<repr::BasicBlockType<repr::NNGraph>>());
261 cfg.createEdge(whileCFNode, bbNode, -1);
263 opNode->resetData(convertOperator(op));
264 auto currentBasicBlock = bbNode->mutableData()->get();
265 currentBasicBlock->pushInstructionNode(opNode);
267 auto opRef = dyn_cast<repr::NeuralNetOperator>(opNode->data().get());
269 assert(opNode->data());
271 auto device_name = op.device_option().node_name();
272 if (device_name !=
"") {
273 auto device = util::make_unique<repr::DeviceAnnotation>(device_name);
274 opRef->setAnnotation(std::move(device));
276 opRef->setAnnotation(util::make_unique<repr::Annotation>());
279 opRef->getMutableAnnotation()->setSaved((
void *)&op);
282 repr::NNModule module;
283 module.dataFlow = std::move(dfg);
284 module.controlFlow = std::move(cfg);
286 *blobMapOut = blobMap;
291 caffe2::NetDef convertToCaffe2Proto(repr::NNModule &m) {
292 auto predictNet = caffe2::NetDef();
294 repr::nn::coalesceInsertedDataDependencies(&m);
298 for (
const auto &bbNode : m.controlFlow.getMutableNodes()) {
299 if (bbNode->getOutEdges().size() > 1) {
300 assert(0 &&
"Control flow not yet supported in Caffe2 converter.");
302 auto bb = bbNode->data().get();
303 for (
const auto &instrNode : bb->getInstructions()) {
304 auto *nnOp = dyn_cast<repr::NeuralNetOperator>(instrNode->data().get());
305 auto *annotation = nnOp->getAnnotation();
306 assert(annotation->getSaved() &&
307 "Generating Caffe2 operators from IR not yet supported.\n");
309 reinterpret_cast<caffe2::OperatorDef *
>(annotation->getSaved());
314 for (
const auto &inEdge : instrNode->getInEdges()) {
316 dyn_cast<repr::NeuralNetData>(inEdge->tail()->data().get());
317 *op->add_input() = tensorNode->getName();
319 for (
const auto &outEdge : instrNode->getOutEdges()) {
321 dyn_cast<repr::NeuralNetData>(outEdge->head()->data().get());
322 *op->add_output() = tensorNode->getName();
325 *predictNet.add_op() = *op;