| 492 | } |
| 493 | |
| 494 | static std::vector<SubModuleInfo> _createSubModuleInfo(std::shared_ptr<BufferStorage> bufferStorage, const std::set<int>& inputIndexes, const std::set<int>& outputIndexes, const std::set<int>& noComputeIndexes, std::shared_ptr<Schedule::ScheduleInfo> sharedConst, bool& success) { |
| 495 | std::vector<SubModuleInfo> submodule; |
| 496 | auto net = flatbuffers::GetRoot<Net>(bufferStorage->buffer()); |
| 497 | auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes); |
| 498 | success = true; |
| 499 | |
| 500 | // Separate the graph to serveral submodule |
| 501 | SubModuleInfo current; |
| 502 | for (int si=0; si<selectOps.size(); ++si) { |
| 503 | auto i = selectOps[si]; |
| 504 | auto op = net->oplists()->GetAs<Op>(i); |
| 505 | if (isBreakOp(op)) { |
| 506 | // TODO: Don't need split segment |
| 507 | if (current.opList.size() > 0) { |
| 508 | // Not empty |
| 509 | submodule.emplace_back(std::move(current)); |
| 510 | } |
| 511 | SubModuleInfo controlOp; |
| 512 | controlOp.opList = {i}; |
| 513 | controlOp.isBreak = true; |
| 514 | if (nullptr != op->inputIndexes()) { |
| 515 | controlOp.inputs.resize(op->inputIndexes()->size()); |
| 516 | ::memcpy(controlOp.inputs.data(), op->inputIndexes()->data(), controlOp.inputs.size() * sizeof(int)); |
| 517 | } |
| 518 | if (nullptr != op->outputIndexes()) { |
| 519 | controlOp.outputs.resize(op->outputIndexes()->size()); |
| 520 | ::memcpy(controlOp.outputs.data(), op->outputIndexes()->data(), controlOp.outputs.size() * sizeof(int)); |
| 521 | } |
| 522 | submodule.emplace_back(std::move(controlOp)); |
| 523 | continue; |
| 524 | } |
| 525 | current.opList.emplace_back(i); |
| 526 | } |
| 527 | if (!current.opList.empty()) { |
| 528 | submodule.emplace_back(std::move(current)); |
| 529 | } |
| 530 | submodule = _splitSubModuleForShapeConst(submodule, net, sharedConst); |
| 531 | for (int moduleIndex=0; moduleIndex < submodule.size(); ++moduleIndex) { |
| 532 | auto& m = submodule[moduleIndex]; |
| 533 | // Compute input / output |
| 534 | if (!m.isBreak) { |
| 535 | _computeTensorMask(m, net); |
| 536 | for (int i=0; i<m.tensorMask.size(); ++i) { |
| 537 | if (0 == m.tensorMask[i]) { |
| 538 | continue; |
| 539 | } |
| 540 | if (1 == m.tensorMask[i]) { |
| 541 | if (noComputeIndexes.find(i) != noComputeIndexes.end()) { |
| 542 | continue; |
| 543 | } |
| 544 | m.inputs.emplace_back(i); |
| 545 | continue; |
| 546 | } |
| 547 | if (2 == m.tensorMask[i]) { |
| 548 | m.outputs.emplace_back(i); |
| 549 | continue; |
| 550 | } |
| 551 | if (3 == m.tensorMask[i]) { |
no test coverage detected