Traverses the subgraph from output_indices to input_indices and returns the set of ops that are visited.
| 671 | // Traverses the subgraph from output_indices to input_indices and returns the |
| 672 | // set of ops that are visited. |
| 673 | StatusOr<absl::flat_hash_set<const tflite::OperatorT*>> PruneSubgraph( |
| 674 | const tflite::SubGraphT& subgraph, ArrayRef<int32_t> input_indices, |
| 675 | ArrayRef<int32_t> output_indices) { |
| 676 | // Create a map from tensor index to defining op. |
| 677 | absl::flat_hash_map<int32_t, const tflite::OperatorT*> defining_op; |
| 678 | for (const auto& op : subgraph.operators) { |
| 679 | for (int32_t output : op->outputs) { |
| 680 | if (!llvm::is_contained(input_indices, output)) { |
| 681 | defining_op[output] = op.get(); |
| 682 | } |
| 683 | } |
| 684 | } |
| 685 | |
| 686 | std::vector<const tflite::OperatorT*> queue; |
| 687 | for (int32_t output : output_indices) { |
| 688 | if (auto& op = defining_op[output]) { |
| 689 | queue.push_back(op); |
| 690 | } else { |
| 691 | return errors::InvalidArgument("Output tensor doesn't have defining op"); |
| 692 | } |
| 693 | } |
| 694 | |
| 695 | // Traverse the graph towards inputs. |
| 696 | absl::flat_hash_set<const tflite::OperatorT*> visited; |
| 697 | while (!queue.empty()) { |
| 698 | const tflite::OperatorT* op = queue.back(); |
| 699 | queue.pop_back(); |
| 700 | if (!visited.insert(op).second) { |
| 701 | // The node has already been visited. |
| 702 | continue; |
| 703 | } |
| 704 | |
| 705 | for (int32_t input : op->inputs) { |
| 706 | // Input tensor may not have a defining op in case it is a subgraph input |
| 707 | // or a constant tensor. |
| 708 | if (auto& op = defining_op[input]) { |
| 709 | queue.push_back(op); |
| 710 | } |
| 711 | } |
| 712 | } |
| 713 | |
| 714 | return visited; |
| 715 | } |
| 716 | |
| 717 | // Build a FuncOp from a tflite SubGraph |
| 718 | // The op_names are a mapping from indexes into the TFLite operators array to |
no test coverage detected