| 798 | } |
| 799 | |
| 800 | void TaskGraph::ForEachGpuDeviceNodes( |
| 801 | const std::function<void(const HashSet<TaskNode*>& dev_nodes)>& Handler) const { |
| 802 | HashMap<std::pair<int64_t, int64_t>, HashSet<TaskNode*>> global_dev_phy_id2nodes; |
| 803 | ForEachNode([&](TaskNode* task_node) { |
| 804 | if (task_node->device_type() == DeviceType::kCPU) { return; } |
| 805 | int64_t dev_phy_id = task_node->stream_id().device_id().device_index(); |
| 806 | global_dev_phy_id2nodes[{task_node->machine_id(), dev_phy_id}].emplace(task_node); |
| 807 | }); |
| 808 | for (const auto& pair : global_dev_phy_id2nodes) { Handler(pair.second); } |
| 809 | } |
| 810 | |
| 811 | void TaskGraph::EnableInplaceMemSharing( |
| 812 | const std::function<bool(const std::string&, const std::string&)>& |
no test coverage detected