Star694
464
代码
代码
Issues33
Pull Requests111
Actions
项目讨论
Wiki
分析
Star694
464
已合并
【PR】: 合入多流加权均衡算法 #3812
xionglan2创建于 7月1日
【PR】: 合入多流加权均衡算法 #3812
已合并
xionglan2创建于 7月1日
从xionglan2/gepr3699_covercommit合入到cann/gedevelop
共 13 个文件变更+2769-9
Mcompiler/CMakeLists.txt+1-0
@@ -120,6 +120,7 @@ set(COMPILER_SRC_LIST
120 "graph/build/dag/dag_graph.cc"120 "graph/build/dag/dag_graph.cc"
121 "graph/build/dag/dag_stream_allocator.cc"121 "graph/build/dag/dag_stream_allocator.cc"
122 "graph/build/dag/dag_stream_merger.cc"122 "graph/build/dag/dag_stream_merger.cc"
123+ "graph/build/dag/dag_weighted_stream_merger.cc"
123 "graph/build/dag/dag_profiling_parser.cc"124 "graph/build/dag/dag_profiling_parser.cc"
124 "graph/build/stream/dag_adapter.cc"125 "graph/build/stream/dag_adapter.cc"
125 "graph/build/stream/dag_stream_allocator_pass.cc"126 "graph/build/stream/dag_stream_allocator_pass.cc"
Mcompiler/graph/build/dag/dag_stream_allocator.cc+17-8
@@ -13,12 +13,12 @@
13#include "graph/build/dag/dag_log.h"13#include "graph/build/dag/dag_log.h"
14#include "graph/build/dag/dag_stream_divide.h"14#include "graph/build/dag/dag_stream_divide.h"
15#include "graph/build/dag/dag_stream_merger.h"15#include "graph/build/dag/dag_stream_merger.h"
16+#include "graph/build/dag/dag_weighted_stream_merger.h"
16#include <algorithm>17#include <algorithm>
17#include <map>18#include <map>
18#include <memory>19#include <memory>
19#include <queue>20#include <queue>
20#include <set>21#include <set>
21-#include <string>
22#include <utility>22#include <utility>
23#include <vector>23#include <vector>
24 24 
@@ -275,14 +275,23 @@ graphStatus ByPathCoverCore(DAGGraph &graph, StreamAllocConfig &config) {
275 275 
276 auto index_routes = ConvertRoutesToIndexRoutes(routes);276 auto index_routes = ConvertRoutesToIndexRoutes(routes);
277 277 
278- StreamMergeOptions options;
279- options.physical_stream_limit =
280- (config.max_stream_id >= 0) ? static_cast<int32_t>(config.max_stream_id + 1) : kDefaultMaxPhysicalStreams;
281- options.strategy = config.merge_strategy;
282- StreamMerger merger(options);
283 std::vector<int32_t> logical_to_physical;278 std::vector<int32_t> logical_to_physical;
284- MINIDAG_ASSERT_SUCCESS(merger.Merge(graph, index_routes, logical_to_physical),279+ const auto physical_stream_limit =
285- "StreamMerger failed, ByPathCover abort.");280+ (config.max_stream_id >= 0) ? static_cast<int32_t>(config.max_stream_id + 1) : kDefaultMaxPhysicalStreams;
281+ if (config.merge_strategy == StreamMergeStrategy::kWeightedLoadBalance) {
282+ WeightedStreamMergeOptions options;
283+ options.physical_stream_limit = physical_stream_limit;
284+ WeightedStreamMerger merger(options);
285+ MINIDAG_ASSERT_SUCCESS(merger.Merge(graph, index_routes, logical_to_physical),
286+ "WeightedStreamMerger failed, ByPathCover abort.");
287+ } else {
288+ StreamMergeOptions options;
289+ options.physical_stream_limit = physical_stream_limit;
290+ options.strategy = config.merge_strategy;
291+ StreamMerger merger(options);
292+ MINIDAG_ASSERT_SUCCESS(merger.Merge(graph, index_routes, logical_to_physical),
293+ "StreamMerger failed, ByPathCover abort.");
294+ }
286 295 
287 AssignStreamIds(routes, logical_to_physical, id_to_node, config);296 AssignStreamIds(routes, logical_to_physical, id_to_node, config);
288 MINIDAG_LOG_INFO("Logical stream num:%zu, merged physical stream num:%ld", routes.size(), config.required_streams);297 MINIDAG_LOG_INFO("Logical stream num:%zu, merged physical stream num:%ld", routes.size(), config.required_streams);
Acompiler/graph/build/dag/dag_weighted_stream_merger.cc+1146-0
@@ -0,0 +1,1146 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software; you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "graph/build/dag/dag_weighted_stream_merger.h"
12+ 
13+#include <algorithm>
14+#include <cmath>
15+#include <cstdint>
16+#include <deque>
17+#include <functional>
18+#include <limits>
19+#include <map>
20+#include <memory>
21+#include <queue>
22+#include <set>
23+#include <string>
24+#include <unordered_map>
25+#include <utility>
26+#include <vector>
27+ 
28+#include "graph/build/dag/dag_checker.h"
29+#include "graph/build/dag/dag_log.h"
30+ 
31+namespace minidag {
32+namespace {
33+static constexpr double kImproveEps = 1e-9;
34+ 
35+struct UnitProfile {
36+ int32_t unit_id = -1;
37+ std::vector<int32_t> node_indices;
38+ int32_t size = 0;
39+ int32_t earliest_level = 0;
40+ int32_t latest_level = 0;
41+ int32_t peak_level_load = 0;
42+ std::vector<int32_t> level_hist;
43+ int32_t interaction = 0;
44+};
45+ 
46+struct WeightedRuntimeInfo {
47+ double duration = 0.0;
48+ int32_t current_stream_id = -1;
49+ int32_t aiv_cores = 0;
50+ int32_t aic_cores = 0;
51+};
52+ 
53+struct LiteScore {
54+ double total = 0.0;
55+ double time_conflict = 0.0;
56+ int32_t event_local = 0;
57+ double time_load = 0.0;
58+ int32_t candidate_stream = 0;
59+};
60+ 
61+struct ResourceScheduleResult {
62+ bool feasible = false;
63+ double makespan = 0.0;
64+};
65+ 
66+struct StreamingState {
67+ std::vector<int32_t> assignment;
68+ std::vector<double> stream_total_durations;
69+ std::vector<std::vector<double>> stream_duration_hists;
70+ std::map<int32_t, std::map<int32_t, int32_t>> origin_stream_to_flow_counts;
71+ int32_t active_stream_count = 0;
72+};
73+ 
74+struct BestRepairMove {
75+ int32_t unit_id = -1;
76+ int32_t from_stream = -1;
77+ int32_t to_stream = -1;
78+ double improvement = 0.0;
79+};
80+ 
81+class WeightedStreamMergeSolver {
82+ public:
83+ WeightedStreamMergeSolver(const DAGGraph &dag, const std::vector<std::vector<int32_t>> &logical_stream_routes,
84+ const WeightedStreamMergeOptions &options)
85+ : dag_(dag), logical_stream_routes_(logical_stream_routes), options_(options) {}
86+ 
87+ graphStatus Solve(std::vector<int32_t> &logical_to_physical_stream) {
88+ MINIDAG_ASSERT_SUCCESS(BuildNodeOrder(), "Build weighted node order failed.");
89+ MINIDAG_ASSERT_SUCCESS(BuildEdges(), "Build weighted edges failed.");
90+ MINIDAG_ASSERT_SUCCESS(BuildLevels(), "Build weighted levels failed.");
91+ MINIDAG_ASSERT_SUCCESS(BuildUnits(), "Build weighted units failed.");
92+ if (unit_count_ == 0) {
93+ logical_to_physical_stream.clear();
94+ return graphStatus::SUCCESS;
95+ }
96+ MINIDAG_ASSERT_SUCCESS(LoadNodeCostProfiles(), "Load weighted node cost profiles failed.");
97+ BuildUnitDependencyGraph();
98+ BuildUnitsByEarliestLevel();
99+ auto state = NewStreamingState();
100+ for (int32_t level = 0; level <= max_level_; ++level) {
101+ const auto level_iter = units_by_earliest_level_.find(level);
102+ if (level_iter == units_by_earliest_level_.end()) {
103+ continue;
104+ }
105+ std::vector<int32_t> recent_units;
106+ for (const auto unit_id : level_iter->second) {
107+ const auto candidate_flows = CandidateFlows(unit_id, level, state, true);
108+ if (candidate_flows.empty()) {
109+ MINIDAG_LOG_ERROR("No candidate physical stream for weighted unit %d.", unit_id);
110+ return graphStatus::FAILED;
111+ }
112+ const auto best_stream = SelectBestStream(unit_id, level, candidate_flows, state);
113+ ApplyUnitToStream(unit_id, best_stream, state);
114+ recent_units.emplace_back(unit_id);
115+ }
116+ if ((options_.repair_moves > 0) && (!recent_units.empty())) {
117+ RepairRecentUnits(level, recent_units, state);
118+ }
119+ }
120+ 
121+ logical_to_physical_stream = CompactAssignment(state.assignment);
122+ if (logical_to_physical_stream.empty()) {
123+ return graphStatus::FAILED;
124+ }
125+ return graphStatus::SUCCESS;
126+ }
127+ 
128+ private:
129+ struct RunningNode {
130+ double finish_time = 0.0;
131+ int32_t topo_pos = 0;
132+ int32_t node_index = 0;
133+ };
134+ 
135+ struct RunningNodeGreater {
136+ bool operator()(const RunningNode &lhs, const RunningNode &rhs) const {
137+ const auto finish_cmp = CompareDouble(lhs.finish_time, rhs.finish_time);
138+ if (finish_cmp != 0) {
139+ return finish_cmp > 0;
140+ }
141+ if (lhs.topo_pos != rhs.topo_pos) {
142+ return lhs.topo_pos > rhs.topo_pos;
143+ }
144+ return lhs.node_index > rhs.node_index;
145+ }
146+ };
147+ 
148+ using RunningNodeHeap = std::priority_queue<RunningNode, std::vector<RunningNode>, RunningNodeGreater>;
149+ 
150+ struct SimulationContext {
151+ std::vector<int32_t> node_streams;
152+ std::map<int32_t, std::deque<int32_t>> stream_to_queue;
153+ std::vector<std::vector<int32_t>> filtered_preds;
154+ std::vector<char> finished;
155+ std::vector<char> running_in_stream;
156+ RunningNodeHeap running_heap;
157+ double current_time = 0.0;
158+ int32_t executed_count = 0;
159+ };
160+ 
161+ static int32_t CompareDouble(const double lhs, const double rhs) {
162+ const auto diff = lhs - rhs;
163+ if (std::fabs(diff) <= kImproveEps) {
164+ return 0;
165+ }
166+ return (diff < 0.0) ? -1 : 1;
167+ }
168+ 
169+ static bool LiteScoreLess(const LiteScore &lhs, const LiteScore &rhs) {
170+ const auto total_cmp = CompareDouble(lhs.total, rhs.total);
171+ if (total_cmp != 0) {
172+ return lhs.total < rhs.total;
173+ }
174+ const auto conflict_cmp = CompareDouble(lhs.time_conflict, rhs.time_conflict);
175+ if (conflict_cmp != 0) {
176+ return lhs.time_conflict < rhs.time_conflict;
177+ }
178+ if (lhs.event_local != rhs.event_local) {
179+ return lhs.event_local < rhs.event_local;
180+ }
181+ const auto load_cmp = CompareDouble(lhs.time_load, rhs.time_load);
182+ if (load_cmp != 0) {
183+ return lhs.time_load < rhs.time_load;
184+ }
185+ return lhs.candidate_stream < rhs.candidate_stream;
186+ }
187+ 
188+ graphStatus BuildNodeOrder() {
189+ topo_nodes_.clear();
190+ const auto all_nodes = dag_.GetAllNodes();
191+ topo_nodes_.reserve(all_nodes.size());
192+ for (size_t idx = 0UL; idx < all_nodes.size(); ++idx) {
193+ MINIDAG_ASSERT_NOTNULL(all_nodes[idx], "Node at index %zu is nullptr.", idx);
194+ topo_nodes_.emplace_back(all_nodes[idx]);
195+ }
196+ std::stable_sort(topo_nodes_.begin(), topo_nodes_.end(),
197+ [](const std::shared_ptr<DAGNode> &lhs, const std::shared_ptr<DAGNode> &rhs) {
198+ return lhs->GetTopoId() < rhs->GetTopoId();
199+ });
200+ node_name_to_index_.clear();
201+ for (size_t idx = 0UL; idx < topo_nodes_.size(); ++idx) {
202+ node_name_to_index_[topo_nodes_[idx]->GetName()] = static_cast<int32_t>(idx);
203+ }
204+ return graphStatus::SUCCESS;
205+ }
206+ 
207+ graphStatus BuildEdges() {
208+ preds_.assign(topo_nodes_.size(), {});
209+ succs_.assign(topo_nodes_.size(), {});
210+ edge_pairs_.clear();
211+ for (size_t node_index = 0UL; node_index < topo_nodes_.size(); ++node_index) {
212+ const auto &node = topo_nodes_[node_index];
213+ std::set<int32_t> unique_pred_indices;
214+ for (const auto &pred_node : node->GetInputNodes()) {
215+ const auto pred_iter = node_name_to_index_.find(pred_node->GetName());
216+ if (pred_iter == node_name_to_index_.end()) {
217+ continue;
218+ }
219+ unique_pred_indices.insert(pred_iter->second);
220+ }
221+ preds_[node_index].assign(unique_pred_indices.begin(), unique_pred_indices.end());
222+ for (const auto pred_index : preds_[node_index]) {
223+ succs_[pred_index].emplace_back(static_cast<int32_t>(node_index));
224+ edge_pairs_.emplace_back(pred_index, static_cast<int32_t>(node_index));
225+ }
226+ }
227+ return graphStatus::SUCCESS;
228+ }
229+ 
230+ graphStatus BuildLevels() {
231+ std::vector<int32_t> indegree(topo_nodes_.size(), 0);
232+ for (size_t idx = 0UL; idx < preds_.size(); ++idx) {
233+ indegree[idx] = static_cast<int32_t>(preds_[idx].size());
234+ }
235+ std::priority_queue<int32_t, std::vector<int32_t>, std::greater<int32_t>> ready_nodes;
236+ for (size_t idx = 0UL; idx < indegree.size(); ++idx) {
237+ if (indegree[idx] == 0) {
238+ ready_nodes.push(static_cast<int32_t>(idx));
239+ }
240+ }
241+ 
242+ topo_order_.clear();
243+ topo_order_.reserve(topo_nodes_.size());
244+ while (!ready_nodes.empty()) {
245+ const auto node_index = ready_nodes.top();
246+ ready_nodes.pop();
247+ topo_order_.emplace_back(node_index);
248+ for (const auto succ_index : succs_[node_index]) {
249+ --indegree[succ_index];
250+ if (indegree[succ_index] == 0) {
251+ ready_nodes.push(succ_index);
252+ }
253+ }
254+ }
255+ if (topo_order_.size() != topo_nodes_.size()) {
256+ MINIDAG_LOG_ERROR("Weighted merge only supports DAG graphs.");
257+ return graphStatus::FAILED;
258+ }
259+ 
260+ topo_position_.assign(topo_nodes_.size(), -1);
261+ levels_.assign(topo_nodes_.size(), 0);
262+ max_level_ = 0;
263+ for (size_t pos = 0UL; pos < topo_order_.size(); ++pos) {
264+ const auto node_index = topo_order_[pos];
265+ topo_position_[node_index] = static_cast<int32_t>(pos);
266+ int32_t level = 0;
267+ for (const auto pred_index : preds_[node_index]) {
268+ level = std::max(level, levels_[pred_index] + 1);
269+ }
270+ levels_[node_index] = level;
271+ max_level_ = std::max(max_level_, level);
272+ }
273+ return graphStatus::SUCCESS;
274+ }
275+ 
276+ graphStatus BuildUnits() {
277+ unit_count_ = static_cast<int32_t>(logical_stream_routes_.size());
278+ unit_profiles_.clear();
279+ unit_profiles_.reserve(logical_stream_routes_.size());
280+ node_to_unit_.assign(topo_nodes_.size(), -1);
281+ 
282+ for (size_t route_id = 0UL; route_id < logical_stream_routes_.size(); ++route_id) {
283+ MINIDAG_ASSERT_SUCCESS(BuildUnitProfile(route_id), "Build weighted unit profile %zu failed.", route_id);
284+ }
285+ 
286+ return ValidateAllNodesAssigned();
287+ }
288+ 
289+ graphStatus BuildUnitProfile(const size_t route_id) {
290+ std::vector<int32_t> node_indices;
291+ node_indices.reserve(logical_stream_routes_[route_id].size());
292+ for (const auto node_index : logical_stream_routes_[route_id]) {
293+ MINIDAG_ASSERT_SUCCESS(AddNodeToUnit(route_id, node_index, node_indices),
294+ "Add node %d to weighted unit %zu failed.", node_index, route_id);
295+ }
296+ if (node_indices.empty()) {
297+ MINIDAG_LOG_ERROR("Weighted logical stream %zu should not be empty.", route_id);
298+ return graphStatus::FAILED;
299+ }
300+ std::sort(node_indices.begin(), node_indices.end(),
301+ [this](const int32_t lhs, const int32_t rhs) { return topo_position_[lhs] < topo_position_[rhs]; });
302+ unit_profiles_.emplace_back(MakeUnitProfile(route_id, std::move(node_indices)));
303+ return graphStatus::SUCCESS;
304+ }
305+ 
306+ graphStatus AddNodeToUnit(const size_t route_id, const int32_t node_index, std::vector<int32_t> &node_indices) {
307+ if ((node_index < 0) || (node_index >= static_cast<int32_t>(topo_nodes_.size()))) {
308+ MINIDAG_LOG_ERROR("Node index %d is out of range [0, %zu).", node_index, topo_nodes_.size());
309+ return graphStatus::FAILED;
310+ }
311+ if (node_to_unit_[node_index] != -1) {
312+ MINIDAG_LOG_ERROR("Node %s appears in more than one logical stream.", topo_nodes_[node_index]->GetName().c_str());
313+ return graphStatus::FAILED;
314+ }
315+ node_to_unit_[node_index] = static_cast<int32_t>(route_id);
316+ node_indices.emplace_back(node_index);
317+ return graphStatus::SUCCESS;
318+ }
319+ 
320+ UnitProfile MakeUnitProfile(const size_t route_id, std::vector<int32_t> node_indices) const {
321+ std::vector<int32_t> level_hist(static_cast<size_t>(max_level_ + 1), 0);
322+ int32_t earliest_level = max_level_;
323+ int32_t latest_level = 0;
324+ for (const auto node_index : node_indices) {
325+ const auto level = levels_[node_index];
326+ ++level_hist[level];
327+ earliest_level = std::min(earliest_level, level);
328+ latest_level = std::max(latest_level, level);
329+ }
330+ int32_t peak_level_load = 0;
331+ for (const auto load : level_hist) {
332+ peak_level_load = std::max(peak_level_load, load);
333+ }
334+ return {static_cast<int32_t>(route_id),
335+ std::move(node_indices),
336+ static_cast<int32_t>(logical_stream_routes_[route_id].size()),
337+ earliest_level,
338+ latest_level,
339+ peak_level_load,
340+ std::move(level_hist),
341+ 0};
342+ }
343+ 
344+ graphStatus ValidateAllNodesAssigned() const {
345+ for (size_t node_index = 0UL; node_index < node_to_unit_.size(); ++node_index) {
346+ if (node_to_unit_[node_index] < 0) {
347+ MINIDAG_LOG_ERROR("Node %s is not assigned to any logical stream.", topo_nodes_[node_index]->GetName().c_str());
348+ return graphStatus::FAILED;
349+ }
350+ }
351+ return graphStatus::SUCCESS;
352+ }
353+ 
354+ graphStatus LoadNodeCostProfiles() {
355+ node_durations_.assign(topo_nodes_.size(), 0.0);
356+ node_origin_stream_hints_.assign(topo_nodes_.size(), 0);
357+ node_aiv_cores_.assign(topo_nodes_.size(), 0);
358+ node_aic_cores_.assign(topo_nodes_.size(), 0);
359+ node_bottom_ranks_.assign(topo_nodes_.size(), 0.0);
360+ 
361+ size_t valid_cost_count = 0UL;
362+ size_t missing_cost_count = 0UL;
363+ for (size_t node_index = 0UL; node_index < topo_nodes_.size(); ++node_index) {
364+ const auto &node = topo_nodes_[node_index];
365+ const auto &cost = node->GetCost();
366+ WeightedRuntimeInfo info;
367+ info.duration = (cost.execution_time >= 0.0f) ? static_cast<double>(cost.execution_time) : 0.0;
368+ info.current_stream_id = node_to_unit_[node_index];
369+ info.aiv_cores = SizeToInt32(cost.vec_block_num);
370+ info.aic_cores = SizeToInt32(cost.cube_block_num);
371+ 
372+ node_durations_[node_index] = info.duration;
373+ node_origin_stream_hints_[node_index] = info.current_stream_id;
374+ node_aiv_cores_[node_index] = info.aiv_cores;
375+ node_aic_cores_[node_index] = info.aic_cores;
376+ if (cost.execution_time >= 0.0f) {
377+ ++valid_cost_count;
378+ } else {
379+ ++missing_cost_count;
380+ }
381+ }
382+ if (missing_cost_count > 0UL) {
383+ MINIDAG_LOG_WARN("Weighted stream NodeCost summary: valid=%zu, missing=%zu.", valid_cost_count,
384+ missing_cost_count);
385+ } else {
386+ MINIDAG_LOG_INFO("Weighted stream NodeCost summary: valid=%zu, missing=%zu.", valid_cost_count,
387+ missing_cost_count);
388+ }
389+ 
390+ ComputeBottomRanks();
391+ BuildResourceUnitProfiles();
392+ return graphStatus::SUCCESS;
393+ }
394+ 
395+ static int32_t SizeToInt32(const size_t value) {
396+ const auto limit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
397+ return static_cast<int32_t>(std::min(value, limit));
398+ }
399+ 
400+ void ComputeBottomRanks() {
401+ for (auto iter = topo_order_.rbegin(); iter != topo_order_.rend(); ++iter) {
402+ const auto node_index = *iter;
403+ double succ_rank = 0.0;
404+ for (const auto succ_index : succs_[node_index]) {
405+ succ_rank = std::max(succ_rank, node_bottom_ranks_[succ_index]);
406+ }
407+ node_bottom_ranks_[node_index] = node_durations_[node_index] + succ_rank;
408+ }
409+ }
410+ 
411+ void BuildResourceUnitProfiles() {
412+ InitResourceUnitProfiles();
413+ for (const auto &unit : unit_profiles_) {
414+ BuildResourceUnitProfile(unit);
415+ }
416+ }
417+ 
418+ void InitResourceUnitProfiles() {
419+ const auto level_count = static_cast<size_t>(max_level_ + 1);
420+ unit_total_duration_.assign(static_cast<size_t>(unit_count_), 0.0);
421+ unit_total_aiv_time_.assign(static_cast<size_t>(unit_count_), 0.0);
422+ unit_total_aic_time_.assign(static_cast<size_t>(unit_count_), 0.0);
423+ unit_duration_hist_.assign(static_cast<size_t>(unit_count_), std::vector<double>(level_count, 0.0));
424+ unit_peak_aiv_hist_.assign(static_cast<size_t>(unit_count_), std::vector<int32_t>(level_count, 0));
425+ unit_peak_aic_hist_.assign(static_cast<size_t>(unit_count_), std::vector<int32_t>(level_count, 0));
426+ unit_peak_aiv_.assign(static_cast<size_t>(unit_count_), 0);
427+ unit_peak_aic_.assign(static_cast<size_t>(unit_count_), 0);
428+ unit_rank_.assign(static_cast<size_t>(unit_count_), 0.0);
429+ unit_origin_stream_hint_.assign(static_cast<size_t>(unit_count_), -1);
430+ unit_active_levels_.assign(static_cast<size_t>(unit_count_), {});
431+ }
432+ 
433+ void BuildResourceUnitProfile(const UnitProfile &unit) {
434+ std::map<int32_t, int32_t> origin_counter;
435+ const auto unit_id = unit.unit_id;
436+ for (const auto node_index : unit.node_indices) {
437+ AccumulateUnitResourceNode(unit_id, node_index, origin_counter);
438+ }
439+ unit_peak_aiv_[unit_id] =
440+ *std::max_element(unit_peak_aiv_hist_[unit_id].begin(), unit_peak_aiv_hist_[unit_id].end());
441+ unit_peak_aic_[unit_id] =
442+ *std::max_element(unit_peak_aic_hist_[unit_id].begin(), unit_peak_aic_hist_[unit_id].end());
443+ unit_origin_stream_hint_[unit_id] = SelectOriginStreamHint(origin_counter);
444+ BuildUnitActiveLevels(unit_id);
445+ }
446+ 
447+ void AccumulateUnitResourceNode(const int32_t unit_id, const int32_t node_index,
448+ std::map<int32_t, int32_t> &origin_counter) {
449+ const auto level = levels_[node_index];
450+ const auto duration = node_durations_[node_index];
451+ const auto aiv = node_aiv_cores_[node_index];
452+ const auto aic = node_aic_cores_[node_index];
453+ ++origin_counter[node_origin_stream_hints_[node_index]];
454+ unit_total_duration_[unit_id] += duration;
455+ unit_total_aiv_time_[unit_id] += duration * static_cast<double>(aiv);
456+ unit_total_aic_time_[unit_id] += duration * static_cast<double>(aic);
457+ unit_duration_hist_[unit_id][level] += duration;
458+ unit_peak_aiv_hist_[unit_id][level] = std::max(unit_peak_aiv_hist_[unit_id][level], aiv);
459+ unit_peak_aic_hist_[unit_id][level] = std::max(unit_peak_aic_hist_[unit_id][level], aic);
460+ unit_rank_[unit_id] = std::max(unit_rank_[unit_id], node_bottom_ranks_[node_index]);
461+ }
462+ 
463+ int32_t SelectOriginStreamHint(const std::map<int32_t, int32_t> &origin_counter) const {
464+ int32_t best_count = -1;
465+ int32_t best_stream = -1;
466+ for (const auto &origin_and_count : origin_counter) {
467+ if ((origin_and_count.second > best_count) ||
468+ ((origin_and_count.second == best_count) && (origin_and_count.first < best_stream))) {
469+ best_count = origin_and_count.second;
470+ best_stream = origin_and_count.first;
471+ }
472+ }
473+ return best_stream;
474+ }
475+ 
476+ void BuildUnitActiveLevels(const int32_t unit_id) {
477+ for (int32_t level = 0; level <= max_level_; ++level) {
478+ if ((unit_duration_hist_[unit_id][level] > kImproveEps) || (unit_peak_aiv_hist_[unit_id][level] > 0) ||
479+ (unit_peak_aic_hist_[unit_id][level] > 0)) {
480+ unit_active_levels_[unit_id].emplace_back(level);
481+ }
482+ }
483+ }
484+ 
485+ void BuildUnitDependencyGraph() {
486+ unit_pred_edges_.assign(static_cast<size_t>(unit_count_), {});
487+ unit_succ_edges_.assign(static_cast<size_t>(unit_count_), {});
488+ std::vector<int32_t> interactions(static_cast<size_t>(unit_count_), 0);
489+ for (const auto &edge_pair : edge_pairs_) {
490+ const auto src_unit = node_to_unit_[edge_pair.first];
491+ const auto dst_unit = node_to_unit_[edge_pair.second];
492+ if (src_unit == dst_unit) {
493+ continue;
494+ }
495+ ++unit_succ_edges_[src_unit][dst_unit];
496+ ++unit_pred_edges_[dst_unit][src_unit];
497+ ++interactions[src_unit];
498+ ++interactions[dst_unit];
499+ }
500+ for (auto &unit : unit_profiles_) {
501+ unit.interaction = interactions[unit.unit_id];
502+ }
503+ }
504+ 
505+ void BuildUnitsByEarliestLevel() {
506+ units_by_earliest_level_.clear();
507+ for (const auto &unit : unit_profiles_) {
508+ units_by_earliest_level_[unit.earliest_level].emplace_back(unit.unit_id);
509+ }
510+ for (auto &level_units : units_by_earliest_level_) {
511+ auto &unit_ids = level_units.second;
512+ std::sort(unit_ids.begin(), unit_ids.end(), [this](const int32_t lhs, const int32_t rhs) {
513+ if (CompareDouble(unit_total_duration_[lhs], unit_total_duration_[rhs]) != 0) {
514+ return unit_total_duration_[lhs] > unit_total_duration_[rhs];
515+ }
516+ const auto lhs_resource_time = unit_total_aiv_time_[lhs] + unit_total_aic_time_[lhs];
517+ const auto rhs_resource_time = unit_total_aiv_time_[rhs] + unit_total_aic_time_[rhs];
518+ if (CompareDouble(lhs_resource_time, rhs_resource_time) != 0) {
519+ return lhs_resource_time > rhs_resource_time;
520+ }
521+ const auto lhs_peak = std::max(unit_peak_aiv_[lhs], unit_peak_aic_[lhs]);
522+ const auto rhs_peak = std::max(unit_peak_aiv_[rhs], unit_peak_aic_[rhs]);
523+ if (lhs_peak != rhs_peak) {
524+ return lhs_peak > rhs_peak;
525+ }
526+ if (unit_profiles_[lhs].interaction != unit_profiles_[rhs].interaction) {
527+ return unit_profiles_[lhs].interaction > unit_profiles_[rhs].interaction;
528+ }
529+ if (unit_profiles_[lhs].earliest_level != unit_profiles_[rhs].earliest_level) {
530+ return unit_profiles_[lhs].earliest_level < unit_profiles_[rhs].earliest_level;
531+ }
532+ return lhs < rhs;
533+ });
534+ }
535+ }
536+ 
537+ StreamingState NewStreamingState() const {
538+ const auto level_count = static_cast<size_t>(max_level_ + 1);
539+ StreamingState state;
540+ state.assignment.assign(static_cast<size_t>(unit_count_), -1);
541+ state.stream_total_durations.assign(static_cast<size_t>(options_.physical_stream_limit), 0.0);
542+ state.stream_duration_hists.assign(static_cast<size_t>(options_.physical_stream_limit),
543+ std::vector<double>(level_count, 0.0));
544+ state.active_stream_count = 0;
545+ return state;
546+ }
547+ 
548+ int32_t WindowEnd(const int32_t start_level) const {
549+ return std::min(max_level_, start_level + options_.window_width - 1);
550+ }
551+ 
552+ double WindowMass(const std::vector<double> &hist, const int32_t start_level) const {
553+ double mass = 0.0;
554+ const auto end_level = WindowEnd(start_level);
555+ for (int32_t level = start_level; level <= end_level; ++level) {
556+ mass += hist[level];
557+ }
558+ return mass;
559+ }
560+ 
561+ double StreamWindowDurationLoad(const StreamingState &state, const int32_t physical_stream,
562+ const int32_t start_level) const {
563+ return WindowMass(state.stream_duration_hists[physical_stream], start_level);
564+ }
565+ 
566+ std::vector<int32_t> WindowActiveLevels(const int32_t unit_id, const int32_t start_level) const {
567+ const auto end_level = WindowEnd(start_level);
568+ std::vector<int32_t> levels;
569+ for (const auto level : unit_active_levels_[unit_id]) {
570+ if ((level >= start_level) && (level <= end_level)) {
571+ levels.emplace_back(level);
572+ }
573+ }
574+ return levels.empty() ? unit_active_levels_[unit_id] : levels;
575+ }
576+ 
577+ std::map<int32_t, int32_t> AdjacentAssignedFlows(const int32_t unit_id,
578+ const std::vector<int32_t> &assignment) const {
579+ std::map<int32_t, int32_t> flow_weights;
580+ for (const auto &pred_flow : unit_pred_edges_[unit_id]) {
581+ const auto flow = assignment[pred_flow.first];
582+ if (flow >= 0) {
583+ flow_weights[flow] += pred_flow.second;
584+ }
585+ }
586+ for (const auto &succ_flow : unit_succ_edges_[unit_id]) {
587+ const auto flow = assignment[succ_flow.first];
588+ if (flow >= 0) {
589+ flow_weights[flow] += succ_flow.second;
590+ }
591+ }
592+ return flow_weights;
593+ }
594+ 
595+ std::vector<int32_t> OriginHintFlows(const int32_t unit_id, const StreamingState &state) const {
596+ std::vector<int32_t> flows;
597+ const auto hint_stream_id = unit_origin_stream_hint_[unit_id];
598+ if (hint_stream_id < 0) {
599+ return flows;
600+ }
601+ const auto hint_iter = state.origin_stream_to_flow_counts.find(hint_stream_id);
602+ if (hint_iter == state.origin_stream_to_flow_counts.end()) {
603+ return flows;
604+ }
605+ std::vector<std::pair<int32_t, int32_t>> flow_counts(hint_iter->second.begin(), hint_iter->second.end());
606+ std::sort(
607+ flow_counts.begin(), flow_counts.end(),
608+ [&state](const std::pair<int32_t, int32_t> &lhs, const std::pair<int32_t, int32_t> &rhs) {
609+ if (lhs.second != rhs.second) {
610+ return lhs.second > rhs.second;
611+ }
612+ if (CompareDouble(state.stream_total_durations[lhs.first], state.stream_total_durations[rhs.first]) != 0) {
613+ return state.stream_total_durations[lhs.first] < state.stream_total_durations[rhs.first];
614+ }
615+ return lhs.first < rhs.first;
616+ });
617+ for (const auto &flow_and_count : flow_counts) {
618+ flows.emplace_back(flow_and_count.first);
619+ }
620+ return flows;
621+ }
622+ 
623+ std::vector<int32_t> CandidateFlows(const int32_t unit_id, const int32_t level, const StreamingState &state,
624+ const bool include_new_flow) const {
625+ auto candidates = AdjacentCandidateFlows(unit_id, state);
626+ const auto hint_flows = OriginHintFlows(unit_id, state);
627+ candidates.insert(candidates.end(), hint_flows.begin(), hint_flows.end());
628+ 
629+ const auto light_flows = LightCandidateFlows(level, state);
630+ candidates.insert(candidates.end(), light_flows.begin(), light_flows.end());
631+ if (include_new_flow && (state.active_stream_count < options_.physical_stream_limit)) {
632+ candidates.emplace_back(state.active_stream_count);
633+ }
634+ 
635+ const auto unique_candidates = UniqueLimitedCandidates(candidates);
636+ if (!unique_candidates.empty()) {
637+ return unique_candidates;
638+ }
639+ return FallbackCandidateFlows(state, include_new_flow);
640+ }
641+ 
642+ std::vector<int32_t> AdjacentCandidateFlows(const int32_t unit_id, const StreamingState &state) const {
643+ const auto adjacent_flow_weights = AdjacentAssignedFlows(unit_id, state.assignment);
644+ std::vector<std::pair<int32_t, int32_t>> weighted_flows(adjacent_flow_weights.begin(), adjacent_flow_weights.end());
645+ std::sort(
646+ weighted_flows.begin(), weighted_flows.end(),
647+ [&state](const std::pair<int32_t, int32_t> &lhs, const std::pair<int32_t, int32_t> &rhs) {
648+ if (lhs.second != rhs.second) {
649+ return lhs.second > rhs.second;
650+ }
651+ if (CompareDouble(state.stream_total_durations[lhs.first], state.stream_total_durations[rhs.first]) != 0) {
652+ return state.stream_total_durations[lhs.first] < state.stream_total_durations[rhs.first];
653+ }
654+ return lhs.first < rhs.first;
655+ });
656+ std::vector<int32_t> candidates;
657+ for (const auto &flow_and_weight : weighted_flows) {
658+ candidates.emplace_back(flow_and_weight.first);
659+ }
660+ return candidates;
661+ }
662+ 
663+ std::vector<int32_t> LightCandidateFlows(const int32_t level, const StreamingState &state) const {
664+ std::vector<int32_t> light_flows;
665+ for (int32_t stream = 0; stream < state.active_stream_count; ++stream) {
666+ light_flows.emplace_back(stream);
667+ }
668+ std::sort(light_flows.begin(), light_flows.end(), [this, &state, level](const int32_t lhs, const int32_t rhs) {
669+ const auto lhs_load = StreamWindowDurationLoad(state, lhs, level);
670+ const auto rhs_load = StreamWindowDurationLoad(state, rhs, level);
671+ if (CompareDouble(lhs_load, rhs_load) != 0) {
672+ return lhs_load < rhs_load;
673+ }
674+ if (CompareDouble(state.stream_total_durations[lhs], state.stream_total_durations[rhs]) != 0) {
675+ return state.stream_total_durations[lhs] < state.stream_total_durations[rhs];
676+ }
677+ return lhs < rhs;
678+ });
679+ std::vector<int32_t> candidates;
680+ for (int32_t idx = 0; (idx < options_.light_stream_limit) && (idx < static_cast<int32_t>(light_flows.size()));
681+ ++idx) {
682+ candidates.emplace_back(light_flows[idx]);
683+ }
684+ return candidates;
685+ }
686+ 
687+ std::vector<int32_t> UniqueLimitedCandidates(const std::vector<int32_t> &candidates) const {
688+ std::vector<int32_t> unique_candidates;
689+ std::set<int32_t> seen;
690+ for (const auto candidate : candidates) {
691+ if (seen.insert(candidate).second) {
692+ unique_candidates.emplace_back(candidate);
693+ }
694+ if (static_cast<int32_t>(unique_candidates.size()) >= options_.candidate_limit) {
695+ break;
696+ }
697+ }
698+ return unique_candidates;
699+ }
700+ 
701+ std::vector<int32_t> FallbackCandidateFlows(const StreamingState &state, const bool include_new_flow) const {
702+ if (include_new_flow && (state.active_stream_count < options_.physical_stream_limit)) {
703+ return {state.active_stream_count};
704+ }
705+ std::vector<int32_t> fallback_candidates;
706+ for (int32_t stream = 0; (stream < state.active_stream_count) && (stream < options_.candidate_limit); ++stream) {
707+ fallback_candidates.emplace_back(stream);
708+ }
709+ return fallback_candidates;
710+ }
711+ 
712+ LiteScore EvaluateLiteScore(const int32_t unit_id, const int32_t candidate_stream, const int32_t level,
713+ const StreamingState &state) const {
714+ LiteScore score;
715+ const auto opens_new_stream = (candidate_stream == state.active_stream_count);
716+ for (const auto &pred_edge : unit_pred_edges_[unit_id]) {
717+ const auto neighbor_stream = state.assignment[pred_edge.first];
718+ if ((neighbor_stream >= 0) && (neighbor_stream != candidate_stream)) {
719+ score.event_local += pred_edge.second;
720+ }
721+ }
722+ for (const auto &succ_edge : unit_succ_edges_[unit_id]) {
723+ const auto neighbor_stream = state.assignment[succ_edge.first];
724+ if ((neighbor_stream >= 0) && (neighbor_stream != candidate_stream)) {
725+ score.event_local += succ_edge.second;
726+ }
727+ }
728+ for (const auto current_level : WindowActiveLevels(unit_id, level)) {
729+ score.time_conflict +=
730+ unit_duration_hist_[unit_id][current_level] * state.stream_duration_hists[candidate_stream][current_level];
731+ }
732+ score.time_load = state.stream_total_durations[candidate_stream] + unit_total_duration_[unit_id];
733+ score.candidate_stream = candidate_stream;
734+ score.total = options_.event_local_weight * static_cast<double>(score.event_local) +
735+ options_.time_conflict_weight * score.time_conflict + options_.time_load_weight * score.time_load +
736+ options_.new_flow_penalty_weight * static_cast<double>(opens_new_stream ? 1 : 0);
737+ return score;
738+ }
739+ 
740+ bool ApplyUnitToStream(const int32_t unit_id, const int32_t candidate_stream, StreamingState &state) const {
741+ const auto opens_new_stream = (candidate_stream == state.active_stream_count);
742+ if (opens_new_stream) {
743+ ++state.active_stream_count;
744+ }
745+ state.assignment[unit_id] = candidate_stream;
746+ state.stream_total_durations[candidate_stream] += unit_total_duration_[unit_id];
747+ for (const auto level : unit_active_levels_[unit_id]) {
748+ state.stream_duration_hists[candidate_stream][level] += unit_duration_hist_[unit_id][level];
749+ }
750+ const auto hint_stream_id = unit_origin_stream_hint_[unit_id];
751+ if (hint_stream_id >= 0) {
752+ ++state.origin_stream_to_flow_counts[hint_stream_id][candidate_stream];
753+ }
754+ return opens_new_stream;
755+ }
756+ 
757+ int32_t RemoveUnitFromStream(const int32_t unit_id, StreamingState &state) const {
758+ const auto physical_stream = state.assignment[unit_id];
759+ if (physical_stream < 0) {
760+ MINIDAG_LOG_ERROR("Try to remove unassigned weighted unit %d.", unit_id);
761+ return -1;
762+ }
763+ state.assignment[unit_id] = -1;
764+ state.stream_total_durations[physical_stream] -= unit_total_duration_[unit_id];
765+ for (const auto level : unit_active_levels_[unit_id]) {
766+ state.stream_duration_hists[physical_stream][level] -= unit_duration_hist_[unit_id][level];
767+ }
768+ const auto hint_stream_id = unit_origin_stream_hint_[unit_id];
769+ if (hint_stream_id >= 0) {
770+ auto hint_iter = state.origin_stream_to_flow_counts.find(hint_stream_id);
771+ if (hint_iter != state.origin_stream_to_flow_counts.end()) {
772+ auto &flow_counter = hint_iter->second;
773+ auto flow_iter = flow_counter.find(physical_stream);
774+ if (flow_iter != flow_counter.end()) {
775+ --flow_iter->second;
776+ if (flow_iter->second <= 0) {
777+ flow_counter.erase(flow_iter);
778+ }
779+ }
780+ if (flow_counter.empty()) {
781+ state.origin_stream_to_flow_counts.erase(hint_iter);
782+ }
783+ }
784+ }
785+ return physical_stream;
786+ }
787+ 
788+ ResourceScheduleResult SimulateAssignment(const std::vector<int32_t> &assignment,
789+ const bool include_unassigned) const {
790+ return BuildAndRunSimulation(assignment, include_unassigned);
791+ }
792+ 
793+ ResourceScheduleResult BuildAndRunSimulation(const std::vector<int32_t> &assignment,
794+ const bool include_unassigned) const {
795+ std::vector<int32_t> node_streams(topo_nodes_.size(), -1);
796+ std::vector<char> active_mask(topo_nodes_.size(), 0);
797+ const auto active_nodes = BuildSimulationNodeStreams(assignment, include_unassigned, node_streams, active_mask);
798+ if (active_nodes < 0) {
799+ return {false, 0.0};
800+ }
801+ if (active_nodes == 0) {
802+ return {true, 0.0};
803+ }
804+ 
805+ auto context = NewSimulationContext(node_streams, active_mask);
806+ return RunSimulation(active_nodes, context);
807+ }
808+ 
809+ SimulationContext NewSimulationContext(const std::vector<int32_t> &node_streams,
810+ const std::vector<char> &active_mask) const {
811+ SimulationContext context;
812+ context.node_streams = node_streams;
813+ context.stream_to_queue = BuildSimulationQueues(node_streams, active_mask);
814+ context.filtered_preds = BuildSimulationFilteredPreds(active_mask);
815+ context.finished.assign(topo_nodes_.size(), 0);
816+ context.running_in_stream.assign(static_cast<size_t>(options_.physical_stream_limit), 0);
817+ return context;
818+ }
819+ 
820+ ResourceScheduleResult RunSimulation(const int32_t active_nodes, SimulationContext &context) const {
821+ while (context.executed_count < active_nodes) {
822+ if (!RunSimulationStep(context)) {
823+ return {false, context.current_time};
824+ }
825+ }
826+ return {true, context.current_time};
827+ }
828+ 
829+ bool RunSimulationStep(SimulationContext &context) const {
830+ const auto ready_heads = SimulationReadyHeads(context.stream_to_queue, context.filtered_preds, context.finished,
831+ context.running_in_stream);
832+ if (!StartSimulationReadyHead(ready_heads, context)) {
833+ return false;
834+ }
835+ if (context.running_heap.empty()) {
836+ return false;
837+ }
838+ FinishNextSimulationNodes(context);
839+ return true;
840+ }
841+ 
842+ int32_t BuildSimulationNodeStreams(const std::vector<int32_t> &assignment, const bool include_unassigned,
843+ std::vector<int32_t> &node_streams, std::vector<char> &active_mask) const {
844+ int32_t active_nodes = 0;
845+ for (size_t node_index = 0UL; node_index < topo_nodes_.size(); ++node_index) {
846+ const auto unit_id = node_to_unit_[node_index];
847+ const auto physical_stream = assignment[unit_id];
848+ if (physical_stream < 0) {
849+ if (include_unassigned) {
850+ MINIDAG_LOG_ERROR("There is still an unassigned weighted unit in final simulation.");
851+ return -1;
852+ }
853+ continue;
854+ }
855+ node_streams[node_index] = physical_stream;
856+ active_mask[node_index] = 1;
857+ ++active_nodes;
858+ }
859+ return active_nodes;
860+ }
861+ 
862+ std::map<int32_t, std::deque<int32_t>> BuildSimulationQueues(const std::vector<int32_t> &node_streams,
863+ const std::vector<char> &active_mask) const {
864+ std::map<int32_t, std::deque<int32_t>> stream_to_queue;
865+ for (const auto node_index : topo_order_) {
866+ if (active_mask[node_index] == 0) {
867+ continue;
868+ }
869+ stream_to_queue[node_streams[node_index]].emplace_back(node_index);
870+ }
871+ return stream_to_queue;
872+ }
873+ 
874+ std::vector<std::vector<int32_t>> BuildSimulationFilteredPreds(const std::vector<char> &active_mask) const {
875+ std::vector<std::vector<int32_t>> filtered_preds(topo_nodes_.size());
876+ for (size_t node_index = 0UL; node_index < topo_nodes_.size(); ++node_index) {
877+ if (active_mask[node_index] == 0) {
878+ continue;
879+ }
880+ for (const auto pred_index : preds_[node_index]) {
881+ if (active_mask[pred_index] != 0) {
882+ filtered_preds[node_index].emplace_back(pred_index);
883+ }
884+ }
885+ }
886+ return filtered_preds;
887+ }
888+ 
889+ std::vector<int32_t> SimulationReadyHeads(const std::map<int32_t, std::deque<int32_t>> &stream_to_queue,
890+ const std::vector<std::vector<int32_t>> &filtered_preds,
891+ const std::vector<char> &finished,
892+ const std::vector<char> &running_in_stream) const {
893+ std::vector<int32_t> ready_heads;
894+ for (const auto &stream_and_queue : stream_to_queue) {
895+ const auto stream = stream_and_queue.first;
896+ const auto &queue = stream_and_queue.second;
897+ if (queue.empty() || (running_in_stream[stream] != 0)) {
898+ continue;
899+ }
900+ if (AreSimulationPredsFinished(queue.front(), filtered_preds, finished)) {
901+ ready_heads.emplace_back(queue.front());
902+ }
903+ }
904+ SortSimulationReadyHeads(ready_heads);
905+ return ready_heads;
906+ }
907+ 
908+ bool AreSimulationPredsFinished(const int32_t node_index, const std::vector<std::vector<int32_t>> &filtered_preds,
909+ const std::vector<char> &finished) const {
910+ for (const auto pred_index : filtered_preds[node_index]) {
911+ if (finished[pred_index] == 0) {
912+ return false;
913+ }
914+ }
915+ return true;
916+ }
917+ 
918+ void SortSimulationReadyHeads(std::vector<int32_t> &ready_heads) const {
919+ std::sort(ready_heads.begin(), ready_heads.end(), [this](const int32_t lhs, const int32_t rhs) {
920+ if (CompareDouble(node_bottom_ranks_[lhs], node_bottom_ranks_[rhs]) != 0) {
921+ return node_bottom_ranks_[lhs] > node_bottom_ranks_[rhs];
922+ }
923+ if (CompareDouble(node_durations_[lhs], node_durations_[rhs]) != 0) {
924+ return node_durations_[lhs] > node_durations_[rhs];
925+ }
926+ if (topo_position_[lhs] != topo_position_[rhs]) {
927+ return topo_position_[lhs] < topo_position_[rhs];
928+ }
929+ return lhs < rhs;
930+ });
931+ }
932+ 
933+ bool StartSimulationReadyHead(const std::vector<int32_t> &ready_heads, SimulationContext &context) const {
934+ if (ready_heads.empty()) {
935+ return true;
936+ }
937+ const auto node_index = ready_heads.front();
938+ const auto stream = context.node_streams[node_index];
939+ auto &queue = context.stream_to_queue[stream];
940+ if (queue.empty() || (queue.front() != node_index)) {
941+ MINIDAG_LOG_ERROR("Weighted stream simulation queue state is inconsistent.");
942+ return false;
943+ }
944+ queue.pop_front();
945+ context.running_in_stream[stream] = 1;
946+ context.running_heap.push(
947+ {context.current_time + node_durations_[node_index], topo_position_[node_index], node_index});
948+ return true;
949+ }
950+ 
951+ void FinishNextSimulationNodes(SimulationContext &context) const {
952+ const auto next_finish = context.running_heap.top().finish_time;
953+ context.current_time = next_finish;
954+ while (!context.running_heap.empty() && (CompareDouble(context.running_heap.top().finish_time, next_finish) == 0)) {
955+ const auto finished_node = context.running_heap.top();
956+ context.running_heap.pop();
957+ context.finished[finished_node.node_index] = 1;
958+ context.running_in_stream[context.node_streams[finished_node.node_index]] = 0;
959+ ++context.executed_count;
960+ }
961+ }
962+ 
963+ int32_t SelectBestStream(const int32_t unit_id, const int32_t level, const std::vector<int32_t> &candidate_flows,
964+ StreamingState &state) const {
965+ std::vector<LiteScore> scored_candidates;
966+ scored_candidates.reserve(candidate_flows.size());
967+ for (const auto candidate_stream : candidate_flows) {
968+ scored_candidates.emplace_back(EvaluateLiteScore(unit_id, candidate_stream, level, state));
969+ }
970+ std::sort(scored_candidates.begin(), scored_candidates.end(), LiteScoreLess);
971+ if (options_.resim_candidate_limit <= 0) {
972+ return scored_candidates.front().candidate_stream;
973+ }
974+ 
975+ bool has_best_resim = false;
976+ double best_resim_span = 0.0;
977+ LiteScore best_resim_score;
978+ const auto resim_limit = std::min(static_cast<int32_t>(scored_candidates.size()), options_.resim_candidate_limit);
979+ for (int32_t idx = 0; idx < resim_limit; ++idx) {
980+ const auto &candidate = scored_candidates[idx];
981+ const auto previous_active_stream_count = state.active_stream_count;
982+ (void)ApplyUnitToStream(unit_id, candidate.candidate_stream, state);
983+ const auto schedule = SimulateAssignment(state.assignment, false);
984+ (void)RemoveUnitFromStream(unit_id, state);
985+ state.active_stream_count = previous_active_stream_count;
986+ if (!schedule.feasible) {
987+ continue;
988+ }
989+ if ((!has_best_resim) || (CompareDouble(schedule.makespan, best_resim_span) < 0) ||
990+ ((CompareDouble(schedule.makespan, best_resim_span) == 0) && LiteScoreLess(candidate, best_resim_score))) {
991+ has_best_resim = true;
992+ best_resim_span = schedule.makespan;
993+ best_resim_score = candidate;
994+ }
995+ }
996+ return has_best_resim ? best_resim_score.candidate_stream : scored_candidates.front().candidate_stream;
997+ }
998+ 
999+ void RepairRecentUnits(const int32_t level, const std::vector<int32_t> &recent_unit_ids,
1000+ StreamingState &state) const {
1001+ int32_t move_count = 0;
1002+ while (move_count < options_.repair_moves) {
1003+ BestRepairMove best_move;
1004+ for (const auto unit_id : recent_unit_ids) {
1005+ const auto previous_active_stream_count = state.active_stream_count;
1006+ const auto current_stream = RemoveUnitFromStream(unit_id, state);
1007+ if (current_stream < 0) {
1008+ return;
1009+ }
1010+ const auto current_score = EvaluateLiteScore(unit_id, current_stream, level, state);
1011+ auto candidate_flows = CandidateFlows(unit_id, level, state, false);
1012+ if (std::find(candidate_flows.begin(), candidate_flows.end(), current_stream) == candidate_flows.end()) {
1013+ candidate_flows.emplace_back(current_stream);
1014+ }
1015+ const auto best_candidate = SelectBestStream(unit_id, level, candidate_flows, state);
1016+ const auto best_score = EvaluateLiteScore(unit_id, best_candidate, level, state);
1017+ (void)ApplyUnitToStream(unit_id, current_stream, state);
1018+ state.active_stream_count = previous_active_stream_count;
1019+ const auto improvement = current_score.total - best_score.total;
1020+ if ((best_candidate != current_stream) && (improvement > kImproveEps) &&
1021+ (improvement > best_move.improvement)) {
1022+ best_move = {unit_id, current_stream, best_candidate, improvement};
1023+ }
1024+ }
1025+ if (best_move.unit_id < 0) {
1026+ return;
1027+ }
1028+ (void)RemoveUnitFromStream(best_move.unit_id, state);
1029+ (void)ApplyUnitToStream(best_move.unit_id, best_move.to_stream, state);
1030+ ++move_count;
1031+ }
1032+ }
1033+ 
1034+ std::vector<int32_t> CompactAssignment(const std::vector<int32_t> &assignment) const {
1035+ std::set<int32_t> used_streams;
1036+ for (const auto stream : assignment) {
1037+ if (stream < 0) {
1038+ MINIDAG_LOG_ERROR("There is still an unassigned weighted unit when compacting assignment.");
1039+ return {};
1040+ }
1041+ used_streams.insert(stream);
1042+ }
1043+ if (used_streams.empty()) {
1044+ MINIDAG_LOG_ERROR("Weighted merge should produce at least one physical stream.");
1045+ return {};
1046+ }
1047+ 
1048+ std::map<int32_t, int32_t> remap;
1049+ int32_t next_stream = 0;
1050+ for (const auto stream : used_streams) {
1051+ remap.emplace(stream, next_stream++);
1052+ }
1053+ std::vector<int32_t> compact_assignment;
1054+ compact_assignment.reserve(assignment.size());
1055+ for (const auto stream : assignment) {
1056+ compact_assignment.emplace_back(remap.at(stream));
1057+ }
1058+ return compact_assignment;
1059+ }
1060+ 
1061+ const DAGGraph &dag_;
1062+ const std::vector<std::vector<int32_t>> &logical_stream_routes_;
1063+ const WeightedStreamMergeOptions &options_;
1064+ 
1065+ std::vector<std::shared_ptr<DAGNode>> topo_nodes_;
1066+ std::unordered_map<std::string, int32_t> node_name_to_index_;
1067+ std::vector<std::vector<int32_t>> preds_;
1068+ std::vector<std::vector<int32_t>> succs_;
1069+ std::vector<int32_t> topo_order_;
1070+ std::vector<int32_t> topo_position_;
1071+ std::vector<int32_t> levels_;
1072+ int32_t max_level_ = 0;
1073+ std::vector<std::pair<int32_t, int32_t>> edge_pairs_;
1074+ 
1075+ std::vector<UnitProfile> unit_profiles_;
1076+ std::vector<int32_t> node_to_unit_;
1077+ int32_t unit_count_ = 0;
1078+ std::vector<std::map<int32_t, int32_t>> unit_pred_edges_;
1079+ std::vector<std::map<int32_t, int32_t>> unit_succ_edges_;
1080+ std::map<int32_t, std::vector<int32_t>> units_by_earliest_level_;
1081+ 
1082+ std::vector<double> node_durations_;
1083+ std::vector<int32_t> node_origin_stream_hints_;
1084+ std::vector<int32_t> node_aiv_cores_;
1085+ std::vector<int32_t> node_aic_cores_;
1086+ std::vector<double> node_bottom_ranks_;
1087+ 
1088+ std::vector<double> unit_total_duration_;
1089+ std::vector<double> unit_total_aiv_time_;
1090+ std::vector<double> unit_total_aic_time_;
1091+ std::vector<std::vector<double>> unit_duration_hist_;
1092+ std::vector<std::vector<int32_t>> unit_peak_aiv_hist_;
1093+ std::vector<std::vector<int32_t>> unit_peak_aic_hist_;
1094+ std::vector<int32_t> unit_peak_aiv_;
1095+ std::vector<int32_t> unit_peak_aic_;
1096+ std::vector<double> unit_rank_;
1097+ std::vector<int32_t> unit_origin_stream_hint_;
1098+ std::vector<std::vector<int32_t>> unit_active_levels_;
1099+};
1100+ 
1101+graphStatus ValidateWeightedStreamMergeOptions(const WeightedStreamMergeOptions &options) {
1102+ if (options.physical_stream_limit <= 0) {
1103+ MINIDAG_LOG_ERROR("Weighted stream merge physical stream limit must be greater than 0.");
1104+ return graphStatus::FAILED;
1105+ }
1106+ if (options.window_width <= 0) {
1107+ MINIDAG_LOG_ERROR("Weighted stream merge window width must be greater than 0.");
1108+ return graphStatus::FAILED;
1109+ }
1110+ if (options.candidate_limit <= 0) {
1111+ MINIDAG_LOG_ERROR("Weighted stream merge candidate limit must be greater than 0.");
1112+ return graphStatus::FAILED;
1113+ }
1114+ if (options.light_stream_limit <= 0) {
1115+ MINIDAG_LOG_ERROR("Weighted stream merge light_stream_limit must be greater than 0.");
1116+ return graphStatus::FAILED;
1117+ }
1118+ if (options.repair_moves < 0) {
1119+ MINIDAG_LOG_ERROR("Weighted stream merge repair_moves must be greater than or equal to 0.");
1120+ return graphStatus::FAILED;
1121+ }
1122+ if (options.resim_candidate_limit < 0) {
1123+ MINIDAG_LOG_ERROR("Weighted stream merge resim_candidate_limit must be greater than or equal to 0.");
1124+ return graphStatus::FAILED;
1125+ }
1126+ return graphStatus::SUCCESS;
1127+}
1128+} // namespace
1129+ 
1130+WeightedStreamMerger::WeightedStreamMerger(const WeightedStreamMergeOptions &options) : options_(options) {}
1131+ 
1132+graphStatus WeightedStreamMerger::Merge(const DAGGraph &dag,
1133+ const std::vector<std::vector<int32_t>> &logical_stream_routes,
1134+ std::vector<int32_t> &logical_to_physical_stream) const {
1135+ logical_to_physical_stream.clear();
1136+ MINIDAG_ASSERT_SUCCESS(ValidateWeightedStreamMergeOptions(options_),
1137+ "Validate weighted stream merge options failed.");
1138+ if (logical_stream_routes.empty()) {
1139+ return graphStatus::SUCCESS;
1140+ }
1141+ 
1142+ WeightedStreamMergeSolver solver(dag, logical_stream_routes, options_);
1143+ return solver.Solve(logical_to_physical_stream);
1144+}
1145+ 
1146+} // namespace minidag
Acompiler/graph/build/dag/dag_weighted_stream_merger.h+56-0
@@ -0,0 +1,56 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software; you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef GE_GRAPH_BUILD_DAG_DAG_WEIGHTED_STREAM_MERGER_H_
12+#define GE_GRAPH_BUILD_DAG_DAG_WEIGHTED_STREAM_MERGER_H_
13+ 
14+#include <cstdint>
15+#include <vector>
16+ 
17+#include "graph/build/dag/dag_graph.h"
18+#include "graph/build/dag/dag_types.h"
19+ 
20+namespace minidag {
21+ 
22+struct WeightedStreamMergeOptions {
23+ int32_t physical_stream_limit = 8;
24+ int32_t window_width = 6;
25+ int32_t candidate_limit = 8;
26+ int32_t light_stream_limit = 3;
27+ int32_t repair_moves = 0;
28+ int32_t resim_candidate_limit = 3;
29+ 
30+ double ressimspan_weight = 8.0;
31+ double stage_var_weight = 3.0;
32+ double event_weight = 3.0;
33+ double time_imbalance_weight = 2.0;
34+ double stream_count_weight = 2.0;
35+ 
36+ double event_local_weight = 3.0;
37+ double time_conflict_weight = 4.0;
38+ double time_load_weight = 2.0;
39+ double new_flow_penalty_weight = 2.0;
40+};
41+ 
42+class WeightedStreamMerger {
43+ public:
44+ explicit WeightedStreamMerger(const WeightedStreamMergeOptions &options = WeightedStreamMergeOptions());
45+ ~WeightedStreamMerger() = default;
46+ 
47+ graphStatus Merge(const DAGGraph &dag, const std::vector<std::vector<int32_t>> &logical_stream_routes,
48+ std::vector<int32_t> &logical_to_physical_stream) const;
49+ 
50+ private:
51+ WeightedStreamMergeOptions options_;
52+};
53+ 
54+} // namespace minidag
55+ 
56+#endif // GE_GRAPH_BUILD_DAG_DAG_WEIGHTED_STREAM_MERGER_H_
Mcompiler/graph/build/stream/dag_stream_allocator_pass.cc+2-0
@@ -65,6 +65,8 @@ bool ParseStreamConfig(const std::string &multi_stream_mode, int64_t &out_max_st
65 out_strategy = minidag::StreamMergeStrategy::kMainStream;65 out_strategy = minidag::StreamMergeStrategy::kMainStream;
66 } else if (algo == "LoadBalance") {66 } else if (algo == "LoadBalance") {
67 out_strategy = minidag::StreamMergeStrategy::kLoadBalance;67 out_strategy = minidag::StreamMergeStrategy::kLoadBalance;
68+ } else if (algo == "WeightedLoadBalance") {
69+ out_strategy = minidag::StreamMergeStrategy::kWeightedLoadBalance;
68 } else {70 } else {
69 const auto invalid_strategy = static_cast<minidag::StreamMergeStrategy>(-1);71 const auto invalid_strategy = static_cast<minidag::StreamMergeStrategy>(-1);
70 const auto *strategy_name = GetStrategyName(invalid_strategy);72 const auto *strategy_name = GetStrategyName(invalid_strategy);
Mtests/ge/st/testcase/CMakeLists.txt+1-0
@@ -106,6 +106,7 @@ add_executable(graph_engine_test
106 "${AIR_CODE_DIR}/tests/ge/st/testcase/test_om2.cc"106 "${AIR_CODE_DIR}/tests/ge/st/testcase/test_om2.cc"
107 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_adapter_integration_test.cc"107 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_adapter_integration_test.cc"
108 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_stream_allocator_pass_test.cc"108 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_stream_allocator_pass_test.cc"
109+ "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/build/dag/dag_weighted_stream_merger_public_st_test.cc"
109 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/ir/named_io_node_builder_test.cc"110 "${AIR_CODE_DIR}/tests/ge/st/testcase/graph/ir/named_io_node_builder_test.cc"
110)111)
111 112 
Mtests/ge/st/testcase/graph/build/dag/dag_stream_allocator_pass_test.cc+134-0
@@ -9,6 +9,11 @@
9 */9 */
10 10 
11#include <gtest/gtest.h>11#include <gtest/gtest.h>
12+#include <cstdio>
13+#include <cstdlib>
14+#include <fstream>
15+#include <string>
16+#include <vector>
12#include <ge_running_env/ge_running_env_faker.h>17#include <ge_running_env/ge_running_env_faker.h>
13#include <ge_running_env/fake_op.h>18#include <ge_running_env/fake_op.h>
14#include <common/share_graph.h>19#include <common/share_graph.h>
@@ -34,6 +39,34 @@ graphStatus CallFromGEGraph(const ConstGraphPtr &ge_graph, std::shared_ptr<minid
34 bool has_profiled_node_cost = false;39 bool has_profiled_node_cost = false;
35 return DAGAdapter::FromGEGraph(ge_graph, dag, has_profiled_node_cost);40 return DAGAdapter::FromGEGraph(ge_graph, dag, has_profiled_node_cost);
36}41}
42+ 
43+struct GraphOptionGuard {
44+ ~GraphOptionGuard() {
45+ ge::GetThreadLocalContext().SetGraphOption({});
46+ }
47+};
48+ 
49+struct ProfilingPathGuard {
50+ explicit ProfilingPathGuard(const char *path) : path_(path) {}
51+ ~ProfilingPathGuard() {
52+ unsetenv("MINIDAG_PROFILING_PATH");
53+ std::remove(path_);
54+ ge::GetThreadLocalContext().SetGraphOption({});
55+ }
56+ const char *path_;
57+};
58+ 
59+bool WriteProfilingCsv(const char *profiling_path, const std::vector<std::string> &rows) {
60+ std::ofstream file(profiling_path);
61+ if (!file.is_open()) {
62+ return false;
63+ }
64+ file << "Op Name,Task Type,Task Duration(us),Block Num,Mix Block Num\n";
65+ for (const auto &row : rows) {
66+ file << row << "\n";
67+ }
68+ return true;
69+}
37} // namespace70} // namespace
38 71 
39class MiniDAGStreamPassTest : public testing::Test {72class MiniDAGStreamPassTest : public testing::Test {
@@ -579,4 +612,105 @@ TEST_F(MiniDAGStreamPassTest, RunPass_InvalidUnknownAlgoName) {
579 ge::GetThreadLocalContext().SetGraphOption({});612 ge::GetThreadLocalContext().SetGraphOption({});
580}613}
581 614 
615+/**
616+ * 场景: profiling 命中时，MiniDAG Stream Pass 使用 WeightedLoadBalance 路径
617+ */
618+TEST_F(MiniDAGStreamPassTest, RunPass_ProfileHitUsesWeightedLoadBalance) {
619+ const char *profiling_path = "/tmp/test_minidag_stream_pass_weighted_hit.csv";
620+ ProfilingPathGuard guard(profiling_path);
621+ ASSERT_TRUE(WriteProfilingCsv(profiling_path, {"add1,AI_CORE,100.0,8,0"}));
622+ setenv("MINIDAG_PROFILING_PATH", profiling_path, 1);
623+ 
624+ std::map<std::string, std::string> options;
625+ options["ge.autoMultistreamParallelMode"] = "LoadBalance:8";
626+ ge::GetThreadLocalContext().SetGraphOption(options);
627+ 
628+ auto compute_graph = gert::ShareGraph::BuildTwoAddNodeKnownShapeGraph();
629+ ASSERT_NE(compute_graph, nullptr);
630+ 
631+ auto graph = GraphUtilsEx::CreateGraphPtrFromComputeGraph(compute_graph);
632+ ASSERT_NE(graph, nullptr);
633+ 
634+ std::shared_ptr<minidag::DAGGraph> dag;
635+ bool has_profiled_node_cost = false;
636+ ASSERT_EQ(DAGAdapter::FromGEGraph(graph, dag, has_profiled_node_cost), ge::GRAPH_SUCCESS);
637+ EXPECT_TRUE(has_profiled_node_cost);
638+ 
639+ ge::StreamPassContext context(0);
640+ auto ret = RunMiniDAGStreamPass(graph, context);
641+ EXPECT_EQ(ret, ge::SUCCESS);
642+ EXPECT_GT(context.GetCurrMaxStreamId(), 0);
643+}
644+ 
645+/**
646+ * 场景: 直接配置 WeightedLoadBalance，覆盖加权均衡策略解析分支
647+ */
648+TEST_F(MiniDAGStreamPassTest, RunPass_WithWeightedLoadBalanceMode) {
649+ GraphOptionGuard guard;
650+ std::map<std::string, std::string> options;
651+ options["ge.autoMultistreamParallelMode"] = "WeightedLoadBalance:8";
652+ ge::GetThreadLocalContext().SetGraphOption(options);
653+ 
654+ auto compute_graph = gert::ShareGraph::BuildTwoAddNodeKnownShapeGraph();
655+ ASSERT_NE(compute_graph, nullptr);
656+ 
657+ auto graph = GraphUtilsEx::CreateGraphPtrFromComputeGraph(compute_graph);
658+ ASSERT_NE(graph, nullptr);
659+ 
660+ ge::StreamPassContext context(0);
661+ auto ret = RunMiniDAGStreamPass(graph, context);
662+ EXPECT_EQ(ret, ge::SUCCESS);
663+ EXPECT_GT(context.GetCurrMaxStreamId(), 0);
664+}
665+ 
666+/**
667+ * 场景: 直接配置 WeightedLoadBalance，复杂图端到端覆盖加权均衡路径
668+ */
669+TEST_F(MiniDAGStreamPassTest, RunPass_WithWeightedLoadBalanceModeOnComplexGraph) {
670+ GraphOptionGuard guard;
671+ std::map<std::string, std::string> options;
672+ options["ge.autoMultistreamParallelMode"] = "WeightedLoadBalance:4";
673+ ge::GetThreadLocalContext().SetGraphOption(options);
674+ 
675+ auto compute_graph = gert::ShareGraph::BuildStaticAbsReluExpAddNodeGraph();
676+ ASSERT_NE(compute_graph, nullptr);
677+ auto graph = GraphUtilsEx::CreateGraphPtrFromComputeGraph(compute_graph);
678+ ASSERT_NE(graph, nullptr);
679+ 
680+ ge::StreamPassContext context(0);
681+ auto ret = RunMiniDAGStreamPass(graph, context);
682+ EXPECT_EQ(ret, ge::SUCCESS);
683+ EXPECT_GT(context.GetCurrMaxStreamId(), 0);
684+}
685+ 
686+/**
687+ * 场景: 多节点 profiling 命中时，复杂图自动使用 WeightedLoadBalance 路径
688+ */
689+TEST_F(MiniDAGStreamPassTest, RunPass_ProfileMultiNodeHitUsesWeightedLoadBalance) {
690+ const char *profiling_path = "/tmp/test_minidag_stream_pass_weighted_multi_hit.csv";
691+ ProfilingPathGuard guard(profiling_path);
692+ ASSERT_TRUE(WriteProfilingCsv(profiling_path, {"abs1,AI_CORE,120.0,8,0", "exp,MIX_AIC,60.0,6,2",
693+ "relu,AI_VECTOR_CORE,30.0,4,0", "add,MIX_AIV,150.0,3,7"}));
694+ setenv("MINIDAG_PROFILING_PATH", profiling_path, 1);
695+ 
696+ std::map<std::string, std::string> options;
697+ options["ge.autoMultistreamParallelMode"] = "LoadBalance:4";
698+ ge::GetThreadLocalContext().SetGraphOption(options);
699+ 
700+ auto compute_graph = gert::ShareGraph::BuildStaticAbsReluExpAddNodeGraph();
701+ ASSERT_NE(compute_graph, nullptr);
702+ auto graph = GraphUtilsEx::CreateGraphPtrFromComputeGraph(compute_graph);
703+ ASSERT_NE(graph, nullptr);
704+ 
705+ std::shared_ptr<minidag::DAGGraph> dag;
706+ bool has_profiled_node_cost = false;
707+ ASSERT_EQ(DAGAdapter::FromGEGraph(graph, dag, has_profiled_node_cost), ge::GRAPH_SUCCESS);
708+ EXPECT_TRUE(has_profiled_node_cost);
709+ 
710+ ge::StreamPassContext context(0);
711+ auto ret = RunMiniDAGStreamPass(graph, context);
712+ EXPECT_EQ(ret, ge::SUCCESS);
713+ EXPECT_GT(context.GetCurrMaxStreamId(), 0);
714+}
715+ 
582} // namespace ge716} // namespace ge
Atests/ge/st/testcase/graph/build/dag/dag_weighted_stream_merger_public_st_test.cc+236-0
@@ -0,0 +1,236 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <gtest/gtest.h>
12+ 
13+#include <limits>
14+#include <memory>
15+#include <string>
16+#include <vector>
17+ 
18+#include "graph/build/dag/dag_graph.h"
19+#include "graph/build/dag/dag_weighted_stream_merger.h"
20+ 
21+namespace minidag {
22+namespace test {
23+namespace {
24+void SetNodeCost(const std::shared_ptr<DAGNode> &node, const float duration, const size_t cube_num,
25+ const size_t vec_num) {
26+ NodeCost cost;
27+ cost.execution_time = duration;
28+ cost.cube_block_num = cube_num;
29+ cost.vec_block_num = vec_num;
30+ node->SetCost(cost);
31+}
32+ 
33+std::shared_ptr<DAGNode> AddCostNode(DAGGraph &dag, const std::string &name, const int64_t topo_id,
34+ const float duration) {
35+ auto node = dag.AddNode(name, "Dummy");
36+ if (node == nullptr) {
37+ return nullptr;
38+ }
39+ node->SetTopoId(topo_id);
40+ SetNodeCost(node, duration, static_cast<size_t>((topo_id % 3) + 1), static_cast<size_t>((topo_id % 2) + 1));
41+ return node;
42+}
43+ 
44+void BuildTwoNodeGraph(DAGGraph &dag) {
45+ auto n0 = AddCostNode(dag, "n0", 0, 5.0F);
46+ auto n1 = AddCostNode(dag, "n1", 1, 7.0F);
47+ ASSERT_NE(n0, nullptr);
48+ ASSERT_NE(n1, nullptr);
49+ ASSERT_EQ(dag.AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
50+}
51+ 
52+void BuildCycleGraph(DAGGraph &dag) {
53+ auto n0 = AddCostNode(dag, "n0", 0, 5.0F);
54+ auto n1 = AddCostNode(dag, "n1", 1, 7.0F);
55+ ASSERT_NE(n0, nullptr);
56+ ASSERT_NE(n1, nullptr);
57+ ASSERT_EQ(dag.AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
58+ ASSERT_EQ(dag.AddEdge(n1, 0, n0, 0), graphStatus::SUCCESS);
59+}
60+ 
61+void BuildCrossStreamGraph(DAGGraph &dag) {
62+ auto n0 = AddCostNode(dag, "n0", 0, 30.0F);
63+ auto n1 = AddCostNode(dag, "n1", 1, 10.0F);
64+ auto n2 = AddCostNode(dag, "n2", 2, 20.0F);
65+ auto n3 = AddCostNode(dag, "n3", 3, 15.0F);
66+ auto n4 = AddCostNode(dag, "n4", 4, 8.0F);
67+ auto n5 = AddCostNode(dag, "n5", 5, 12.0F);
68+ ASSERT_NE(n0, nullptr);
69+ ASSERT_NE(n1, nullptr);
70+ ASSERT_NE(n2, nullptr);
71+ ASSERT_NE(n3, nullptr);
72+ ASSERT_NE(n4, nullptr);
73+ ASSERT_NE(n5, nullptr);
74+ ASSERT_EQ(dag.AddEdge(n0, 0, n3, 0), graphStatus::SUCCESS);
75+ ASSERT_EQ(dag.AddEdge(n1, 0, n3, 0), graphStatus::SUCCESS);
76+ ASSERT_EQ(dag.AddEdge(n2, 0, n4, 0), graphStatus::SUCCESS);
77+ ASSERT_EQ(dag.AddEdge(n3, 0, n5, 0), graphStatus::SUCCESS);
78+ ASSERT_EQ(dag.AddEdge(n4, 0, n5, 0), graphStatus::SUCCESS);
79+}
80+ 
81+void BuildInternalEdgeGraph(DAGGraph &dag) {
82+ auto n0 = AddCostNode(dag, "n0", 0, 9.0F);
83+ auto n1 = AddCostNode(dag, "n1", 1, 3.0F);
84+ auto n2 = AddCostNode(dag, "n2", 2, 5.0F);
85+ auto n3 = AddCostNode(dag, "n3", 3, 0.0F);
86+ ASSERT_NE(n0, nullptr);
87+ ASSERT_NE(n1, nullptr);
88+ ASSERT_NE(n2, nullptr);
89+ ASSERT_NE(n3, nullptr);
90+ SetNodeCost(n3, 0.0F, 0U, 0U);
91+ ASSERT_EQ(dag.AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
92+ ASSERT_EQ(dag.AddEdge(n1, 0, n2, 0), graphStatus::SUCCESS);
93+ ASSERT_EQ(dag.AddEdge(n2, 0, n3, 0), graphStatus::SUCCESS);
94+}
95+ 
96+void ExpectValidMapping(const std::vector<int32_t> &mapping, const size_t expected_size, const int32_t stream_limit) {
97+ ASSERT_EQ(mapping.size(), expected_size);
98+ for (const auto stream : mapping) {
99+ EXPECT_GE(stream, 0);
100+ EXPECT_LT(stream, stream_limit);
101+ }
102+}
103+} // namespace
104+ 
105+TEST(WeightedStreamMergerPublicStTest, Merge_CoversInvalidOptionsOnProductionObject) {
106+ DAGGraph dag("public_invalid_options");
107+ BuildTwoNodeGraph(dag);
108+ const std::vector<std::vector<int32_t>> routes = {{0}, {1}};
109+ std::vector<int32_t> mapping;
110+ 
111+ WeightedStreamMergeOptions options;
112+ options.physical_stream_limit = 0;
113+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
114+ options.physical_stream_limit = 2;
115+ options.window_width = 0;
116+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
117+ options.window_width = 1;
118+ options.candidate_limit = 0;
119+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
120+ options.candidate_limit = 1;
121+ options.light_stream_limit = 0;
122+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
123+ options.light_stream_limit = 1;
124+ options.repair_moves = -1;
125+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
126+ options.repair_moves = 0;
127+ options.resim_candidate_limit = -1;
128+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, routes, mapping), graphStatus::FAILED);
129+}
130+ 
131+TEST(WeightedStreamMergerPublicStTest, Merge_CoversRouteFailuresOnProductionObject) {
132+ DAGGraph dag("public_route_failures");
133+ BuildTwoNodeGraph(dag);
134+ std::vector<int32_t> mapping;
135+ WeightedStreamMergeOptions options;
136+ 
137+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, {}, mapping), graphStatus::SUCCESS);
138+ EXPECT_TRUE(mapping.empty());
139+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, {{0}}, mapping), graphStatus::FAILED);
140+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, {{0}, {0}, {1}}, mapping), graphStatus::FAILED);
141+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, {{0}, {2}}, mapping), graphStatus::FAILED);
142+ EXPECT_EQ(WeightedStreamMerger(options).Merge(dag, {{}, {1}}, mapping), graphStatus::FAILED);
143+}
144+ 
145+TEST(WeightedStreamMergerPublicStTest, Merge_CoversCycleFailureOnProductionObject) {
146+ DAGGraph dag("public_cycle");
147+ BuildCycleGraph(dag);
148+ std::vector<int32_t> mapping;
149+ 
150+ EXPECT_EQ(WeightedStreamMerger().Merge(dag, {{0}, {1}}, mapping), graphStatus::FAILED);
151+}
152+ 
153+TEST(WeightedStreamMergerPublicStTest, Merge_CoversMissingCostAndExtremeCoresOnProductionObject) {
154+ DAGGraph dag("public_mixed_cost");
155+ auto n0 = AddCostNode(dag, "n0", 0, 5.0F);
156+ auto n1 = AddCostNode(dag, "n1", 1, -1.0F);
157+ auto n2 = AddCostNode(dag, "n2", 2, 7.0F);
158+ auto n3 = dag.AddNode("n3", "Dummy");
159+ ASSERT_NE(n0, nullptr);
160+ ASSERT_NE(n1, nullptr);
161+ ASSERT_NE(n2, nullptr);
162+ ASSERT_NE(n3, nullptr);
163+ n3->SetTopoId(3);
164+ SetNodeCost(n2, 7.0F, std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
165+ ASSERT_EQ(dag.AddEdge(n0, 0, n2, 0), graphStatus::SUCCESS);
166+ ASSERT_EQ(dag.AddEdge(n1, 0, n2, 0), graphStatus::SUCCESS);
167+ ASSERT_EQ(dag.AddEdge(n2, 0, n3, 0), graphStatus::SUCCESS);
168+ 
169+ std::vector<int32_t> mapping;
170+ WeightedStreamMergeOptions options;
171+ options.physical_stream_limit = 3;
172+ options.repair_moves = 1;
173+ ASSERT_EQ(WeightedStreamMerger(options).Merge(dag, {{1, 0}, {2}, {3}}, mapping), graphStatus::SUCCESS);
174+ ExpectValidMapping(mapping, 3U, options.physical_stream_limit);
175+}
176+ 
177+TEST(WeightedStreamMergerPublicStTest, Merge_CoversCrossStreamDependencyOnProductionObject) {
178+ DAGGraph dag("public_cross_stream");
179+ BuildCrossStreamGraph(dag);
180+ std::vector<int32_t> mapping;
181+ WeightedStreamMergeOptions options;
182+ options.physical_stream_limit = 3;
183+ options.candidate_limit = 4;
184+ options.light_stream_limit = 3;
185+ options.repair_moves = 3;
186+ options.resim_candidate_limit = 3;
187+ 
188+ ASSERT_EQ(WeightedStreamMerger(options).Merge(dag, {{0}, {1}, {2}, {3}, {4}, {5}}, mapping), graphStatus::SUCCESS);
189+ ExpectValidMapping(mapping, 6U, options.physical_stream_limit);
190+}
191+ 
192+TEST(WeightedStreamMergerPublicStTest, Merge_CoversInternalEdgeAndFastScorePathOnProductionObject) {
193+ DAGGraph dag("public_internal_edge");
194+ BuildInternalEdgeGraph(dag);
195+ std::vector<int32_t> mapping;
196+ WeightedStreamMergeOptions options;
197+ options.physical_stream_limit = 2;
198+ options.resim_candidate_limit = 0;
199+ 
200+ ASSERT_EQ(WeightedStreamMerger(options).Merge(dag, {{1, 0}, {2}, {3}}, mapping), graphStatus::SUCCESS);
201+ ExpectValidMapping(mapping, 3U, options.physical_stream_limit);
202+}
203+ 
204+TEST(WeightedStreamMergerPublicStTest, Merge_CoversSinglePhysicalStreamOnProductionObject) {
205+ DAGGraph dag("public_single_stream");
206+ for (int32_t idx = 0; idx < 6; ++idx) {
207+ ASSERT_NE(AddCostNode(dag, "n" + std::to_string(idx), idx, 3.0F), nullptr);
208+ }
209+ std::vector<int32_t> mapping;
210+ WeightedStreamMergeOptions options;
211+ options.physical_stream_limit = 1;
212+ 
213+ ASSERT_EQ(WeightedStreamMerger(options).Merge(dag, {{0}, {1}, {2}, {3}, {4}, {5}}, mapping), graphStatus::SUCCESS);
214+ ASSERT_EQ(mapping.size(), 6U);
215+ for (const auto stream : mapping) {
216+ EXPECT_EQ(stream, 0);
217+ }
218+}
219+ 
220+TEST(WeightedStreamMergerPublicStTest, Merge_CoversExternalInputNodeOnProductionObject) {
221+ DAGGraph dag("public_external_pred");
222+ auto n0 = AddCostNode(dag, "n0", 0, 4.0F);
223+ auto n1 = AddCostNode(dag, "n1", 1, 6.0F);
224+ auto external = std::make_shared<DAGNode>("external", "Dummy");
225+ ASSERT_NE(n0, nullptr);
226+ ASSERT_NE(n1, nullptr);
227+ ASSERT_EQ(dag.AddEdge(external, 0, n0, 0), graphStatus::SUCCESS);
228+ ASSERT_EQ(dag.AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
229+ 
230+ std::vector<int32_t> mapping;
231+ ASSERT_EQ(WeightedStreamMerger().Merge(dag, {{0}, {1}}, mapping), graphStatus::SUCCESS);
232+ ExpectValidMapping(mapping, 2U, WeightedStreamMergeOptions().physical_stream_limit);
233+}
234+ 
235+} // namespace test
236+} // namespace minidag
Mtests/ge/ut/ge/CMakeLists.txt+2-0
@@ -499,6 +499,8 @@ set(MULTI_PARTS_TEST_FILES
499 "graph/build/dag/dag_stream_allocator_test.cc"499 "graph/build/dag/dag_stream_allocator_test.cc"
500 "graph/build/dag/dag_stream_divide_test.cc"500 "graph/build/dag/dag_stream_divide_test.cc"
501 "graph/build/dag/dag_stream_merger_test.cc"501 "graph/build/dag/dag_stream_merger_test.cc"
502+ "graph/build/dag/dag_weighted_stream_merger_test.cc"
503+ "graph/build/dag/dag_weighted_stream_merger_whitebox_test.cc"
502 "graph/build/dag/dag_profiling_parser_test.cc"504 "graph/build/dag/dag_profiling_parser_test.cc"
503 "graph/execute/graph_execute_unittest.cc"505 "graph/execute/graph_execute_unittest.cc"
504 "graph/execute/model_executor_unittest.cc"506 "graph/execute/model_executor_unittest.cc"
Mtests/ge/ut/ge/graph/build/dag/dag_stream_allocator_pass_ut_test.cc+36-0
@@ -459,6 +459,24 @@ TEST(DagStreamAllocatorPassTest, RunPass_WithAutoMultistreamMode_MainStream) {
459 ge::GetThreadLocalContext().SetGraphOption({});459 ge::GetThreadLocalContext().SetGraphOption({});
460}460}
461 461 
462+/**
463+ * 场景 C2-1: 设置 ge.autoMultistreamParallelMode="WeightedLoadBalance:6" - 解析冒号格式+WeightedLoadBalance策略
464+ */
465+TEST(DagStreamAllocatorPassTest, RunPass_WithAutoMultistreamMode_WeightedLoadBalance) {
466+ std::map<std::string, std::string> options;
467+ options["ge.autoMultistreamParallelMode"] = "WeightedLoadBalance:6";
468+ ge::GetThreadLocalContext().SetGraphOption(options);
469+ 
470+ auto graph = BuildGraphWithControlEdge();
471+ ASSERT_NE(graph, nullptr);
472+ 
473+ ge::StreamPassContext context(0);
474+ auto ret = RunMiniDAGStreamPass(graph, context);
475+ EXPECT_EQ(ret, ge::SUCCESS);
476+ 
477+ ge::GetThreadLocalContext().SetGraphOption({});
478+}
479+ 
462/**480/**
463 * 场景 C3: 设置 ge.autoMultistreamParallelMode="LoadBalance:invalid" - 无效max_stream值，返回FAILED481 * 场景 C3: 设置 ge.autoMultistreamParallelMode="LoadBalance:invalid" - 无效max_stream值，返回FAILED
464 */482 */
@@ -495,6 +513,24 @@ TEST(DagStreamAllocatorPassTest, RunPass_WithAutoMultistreamMode_LegacyFormat) {
495 ge::GetThreadLocalContext().SetGraphOption({});513 ge::GetThreadLocalContext().SetGraphOption({});
496}514}
497 515 
516+/**
517+ * 场景 C4-1: 设置 ge.autoMultistreamParallelMode="WeightedLoadBalance" - 无冒号格式，返回FAILED
518+ */
519+TEST(DagStreamAllocatorPassTest, RunPass_WithAutoMultistreamMode_WeightedLoadBalanceLegacyFormat) {
520+ std::map<std::string, std::string> options;
521+ options["ge.autoMultistreamParallelMode"] = "WeightedLoadBalance";
522+ ge::GetThreadLocalContext().SetGraphOption(options);
523+ 
524+ auto graph = BuildGraphWithControlEdge();
525+ ASSERT_NE(graph, nullptr);
526+ 
527+ ge::StreamPassContext context(0);
528+ auto ret = RunMiniDAGStreamPass(graph, context);
529+ EXPECT_EQ(ret, ge::FAILED);
530+ 
531+ ge::GetThreadLocalContext().SetGraphOption({});
532+}
533+ 
498/**534/**
499 * 场景 C5: 设置 ge.autoMultistreamParallelMode="LoadBalance:0" - max_val <= 0，返回FAILED535 * 场景 C5: 设置 ge.autoMultistreamParallelMode="LoadBalance:0" - max_val <= 0，返回FAILED
500 */536 */
Mtests/ge/ut/ge/graph/build/dag/dag_stream_allocator_test.cc+36-1
@@ -9,12 +9,12 @@
9 */9 */
10 10 
11#include <gtest/gtest.h>11#include <gtest/gtest.h>
12+ 
12#include "graph/build/dag/dag_graph.h"13#include "graph/build/dag/dag_graph.h"
13#include "graph/build/dag/dag_stream_allocator.h"14#include "graph/build/dag/dag_stream_allocator.h"
14 15 
15namespace minidag {16namespace minidag {
16namespace test {17namespace test {
17- 
18class DagStreamAllocatorTest : public testing::Test {18class DagStreamAllocatorTest : public testing::Test {
19 protected:19 protected:
20 void SetUp() override {20 void SetUp() override {
@@ -160,5 +160,40 @@ TEST_F(DagStreamAllocatorTest, ByPathCover_EdgeToNodeOutsideGraph_AbortWithoutAs
160 EXPECT_EQ(n1->GetStreamId(), INVALID_STREAM_ID);160 EXPECT_EQ(n1->GetStreamId(), INVALID_STREAM_ID);
161}161}
162 162 
163+TEST_F(DagStreamAllocatorTest, ByPathCover_WeightedLoadBalanceStrategy_UsesNodeCostWithoutJson) {
164+ auto n1 = graph_->AddNode("n1", "Op1");
165+ auto n2 = graph_->AddNode("n2", "Op2");
166+ auto n3 = graph_->AddNode("n3", "Op3");
167+ auto n4 = graph_->AddNode("n4", "Op4");
168+ ASSERT_EQ(graph_->AddEdge(n1, 0, n2, 0), graphStatus::SUCCESS);
169+ ASSERT_EQ(graph_->AddEdge(n1, 0, n3, 0), graphStatus::SUCCESS);
170+ ASSERT_EQ(graph_->AddEdge(n2, 0, n4, 0), graphStatus::SUCCESS);
171+ ASSERT_EQ(graph_->AddEdge(n3, 0, n4, 0), graphStatus::SUCCESS);
172+ 
173+ NodeCost cost;
174+ cost.execution_time = 10.0f;
175+ cost.cube_block_num = 4;
176+ cost.vec_block_num = 8;
177+ n1->SetCost(cost);
178+ n2->SetCost(cost);
179+ n3->SetCost(cost);
180+ n4->SetCost(cost);
181+ 
182+ StreamAllocConfig config{1, 0, 3};
183+ config.merge_strategy = StreamMergeStrategy::kWeightedLoadBalance;
184+ DagStreamAllocator::ByPathCover(*graph_, config);
185+ 
186+ EXPECT_GE(config.required_streams, 1);
187+ EXPECT_LE(config.required_streams, 2);
188+ EXPECT_GE(n1->GetStreamId(), 3);
189+ EXPECT_GE(n2->GetStreamId(), 3);
190+ EXPECT_GE(n3->GetStreamId(), 3);
191+ EXPECT_GE(n4->GetStreamId(), 3);
192+ EXPECT_LE(n1->GetStreamId(), 4);
193+ EXPECT_LE(n2->GetStreamId(), 4);
194+ EXPECT_LE(n3->GetStreamId(), 4);
195+ EXPECT_LE(n4->GetStreamId(), 4);
196+}
197+ 
163} // namespace test198} // namespace test
164} // namespace minidag199} // namespace minidag
Atests/ge/ut/ge/graph/build/dag/dag_weighted_stream_merger_test.cc+398-0
@@ -0,0 +1,398 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <gtest/gtest.h>
12+ 
13+#include <limits>
14+#include <memory>
15+#include <string>
16+#include <vector>
17+ 
18+#include "graph/build/dag/dag_graph.h"
19+#include "graph/build/dag/dag_weighted_stream_merger.h"
20+ 
21+namespace minidag {
22+namespace test {
23+ 
24+class WeightedStreamMergerTest : public testing::Test {
25+ protected:
26+ void SetUp() override {
27+ dag_ = std::make_shared<DAGGraph>("weighted_stream_merger_test");
28+ options_.physical_stream_limit = 2;
29+ options_.candidate_limit = 4;
30+ options_.light_stream_limit = 2;
31+ options_.resim_candidate_limit = 2;
32+ }
33+ 
34+ void SetCost(const std::shared_ptr<DAGNode> &node, const float duration, const size_t cube_num,
35+ const size_t vec_num) {
36+ NodeCost cost;
37+ cost.execution_time = duration;
38+ cost.cube_block_num = cube_num;
39+ cost.vec_block_num = vec_num;
40+ node->SetCost(cost);
41+ }
42+ 
43+ std::shared_ptr<DAGNode> AddNode(const std::string &name, const float duration, const size_t cube_num,
44+ const size_t vec_num) {
45+ auto node = dag_->AddNode(name, "Dummy");
46+ node->SetTopoId(static_cast<int64_t>(dag_->GetNodeCount()) - 1);
47+ SetCost(node, duration, cube_num, vec_num);
48+ return node;
49+ }
50+ 
51+ void BuildDiamondGraph() {
52+ auto n0 = AddNode("n0", 10.0F, 1U, 2U);
53+ auto n1 = AddNode("n1", 60.0F, 4U, 1U);
54+ auto n2 = AddNode("n2", 20.0F, 1U, 4U);
55+ auto n3 = AddNode("n3", 10.0F, 2U, 2U);
56+ ASSERT_EQ(dag_->AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
57+ ASSERT_EQ(dag_->AddEdge(n0, 0, n2, 0), graphStatus::SUCCESS);
58+ ASSERT_EQ(dag_->AddEdge(n1, 0, n3, 0), graphStatus::SUCCESS);
59+ ASSERT_EQ(dag_->AddEdge(n2, 0, n3, 0), graphStatus::SUCCESS);
60+ }
61+ 
62+ void BuildIndependentGraph(const int32_t count, const float duration) {
63+ for (int32_t idx = 0; idx < count; ++idx) {
64+ auto node = AddNode("n" + std::to_string(idx), duration, static_cast<size_t>((idx % 3) + 1),
65+ static_cast<size_t>((idx % 2) + 1));
66+ node->SetStreamId(idx % 2);
67+ }
68+ }
69+ 
70+ void BuildTwoStageParallelGraph() {
71+ auto n0 = AddNode("n0", 5.0F, 1U, 1U);
72+ auto n1 = AddNode("n1", 5.0F, 1U, 1U);
73+ auto n2 = AddNode("n2", 5.0F, 1U, 1U);
74+ auto n3 = AddNode("n3", 5.0F, 1U, 1U);
75+ auto n4 = AddNode("n4", 7.0F, 2U, 1U);
76+ auto n5 = AddNode("n5", 7.0F, 1U, 2U);
77+ n0->SetStreamId(0);
78+ n1->SetStreamId(1);
79+ n2->SetStreamId(0);
80+ n3->SetStreamId(1);
81+ n4->SetStreamId(0);
82+ n5->SetStreamId(1);
83+ ASSERT_EQ(dag_->AddEdge(n0, 0, n4, 0), graphStatus::SUCCESS);
84+ ASSERT_EQ(dag_->AddEdge(n1, 0, n4, 0), graphStatus::SUCCESS);
85+ ASSERT_EQ(dag_->AddEdge(n2, 0, n5, 0), graphStatus::SUCCESS);
86+ ASSERT_EQ(dag_->AddEdge(n3, 0, n5, 0), graphStatus::SUCCESS);
87+ }
88+ 
89+ void BuildCrossStreamDependencyGraph() {
90+ auto n0 = AddNode("n0", 30.0F, 4U, 1U);
91+ auto n1 = AddNode("n1", 10.0F, 1U, 4U);
92+ auto n2 = AddNode("n2", 20.0F, 2U, 2U);
93+ auto n3 = AddNode("n3", 15.0F, 3U, 1U);
94+ auto n4 = AddNode("n4", 8.0F, 1U, 3U);
95+ auto n5 = AddNode("n5", 12.0F, 2U, 2U);
96+ ASSERT_EQ(dag_->AddEdge(n0, 0, n3, 0), graphStatus::SUCCESS);
97+ ASSERT_EQ(dag_->AddEdge(n1, 0, n3, 0), graphStatus::SUCCESS);
98+ ASSERT_EQ(dag_->AddEdge(n2, 0, n4, 0), graphStatus::SUCCESS);
99+ ASSERT_EQ(dag_->AddEdge(n3, 0, n5, 0), graphStatus::SUCCESS);
100+ ASSERT_EQ(dag_->AddEdge(n4, 0, n5, 0), graphStatus::SUCCESS);
101+ }
102+ 
103+ void BuildCycleGraph() {
104+ auto n0 = AddNode("n0", 10.0F, 1U, 1U);
105+ auto n1 = AddNode("n1", 10.0F, 1U, 1U);
106+ ASSERT_EQ(dag_->AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
107+ ASSERT_EQ(dag_->AddEdge(n1, 0, n0, 0), graphStatus::SUCCESS);
108+ }
109+ 
110+ void BuildMissingAndExtremeCostGraph() {
111+ auto n0 = AddNode("n0", 5.0F, 1U, 1U);
112+ auto n1 = AddNode("n1", -1.0F, 2U, 2U);
113+ auto n2 = AddNode("n2", 7.0F, std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
114+ auto n3 = dag_->AddNode("n3", "Dummy");
115+ ASSERT_NE(n3, nullptr);
116+ n3->SetTopoId(3);
117+ ASSERT_EQ(dag_->AddEdge(n0, 0, n2, 0), graphStatus::SUCCESS);
118+ ASSERT_EQ(dag_->AddEdge(n1, 0, n2, 0), graphStatus::SUCCESS);
119+ ASSERT_EQ(dag_->AddEdge(n2, 0, n3, 0), graphStatus::SUCCESS);
120+ }
121+ 
122+ void ExpectValidMapping(const std::vector<int32_t> &mapping, const size_t expected_size) const {
123+ ASSERT_EQ(mapping.size(), expected_size);
124+ for (const auto stream_id : mapping) {
125+ EXPECT_GE(stream_id, 0);
126+ EXPECT_LT(stream_id, options_.physical_stream_limit);
127+ }
128+ }
129+ 
130+ std::shared_ptr<DAGGraph> dag_;
131+ WeightedStreamMergeOptions options_;
132+};
133+ 
134+TEST_F(WeightedStreamMergerTest, Merge_EmptyRoutes_ReturnsSuccess) {
135+ std::vector<std::vector<int32_t>> logical_routes;
136+ std::vector<int32_t> logical_to_physical;
137+ 
138+ WeightedStreamMerger merger(options_);
139+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
140+ 
141+ EXPECT_EQ(status, graphStatus::SUCCESS);
142+ EXPECT_TRUE(logical_to_physical.empty());
143+}
144+ 
145+TEST_F(WeightedStreamMergerTest, Merge_DiamondGraph_RespectsPhysicalStreamLimit) {
146+ BuildDiamondGraph();
147+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}};
148+ std::vector<int32_t> logical_to_physical;
149+ 
150+ WeightedStreamMerger merger(options_);
151+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
152+ 
153+ ASSERT_EQ(status, graphStatus::SUCCESS);
154+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
155+ for (const auto stream_id : logical_to_physical) {
156+ EXPECT_GE(stream_id, 0);
157+ EXPECT_LT(stream_id, options_.physical_stream_limit);
158+ }
159+}
160+ 
161+TEST_F(WeightedStreamMergerTest, Merge_GroupedLogicalRoutes_ReturnsRouteMapping) {
162+ BuildDiamondGraph();
163+ std::vector<std::vector<int32_t>> logical_routes = {{0, 1}, {2, 3}};
164+ std::vector<int32_t> logical_to_physical;
165+ 
166+ WeightedStreamMerger merger(options_);
167+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
168+ 
169+ ASSERT_EQ(status, graphStatus::SUCCESS);
170+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
171+ for (const auto stream_id : logical_to_physical) {
172+ EXPECT_GE(stream_id, 0);
173+ EXPECT_LT(stream_id, options_.physical_stream_limit);
174+ }
175+}
176+ 
177+TEST_F(WeightedStreamMergerTest, Merge_MissingNodeInRoutes_ReturnsFailed) {
178+ BuildDiamondGraph();
179+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}};
180+ std::vector<int32_t> logical_to_physical;
181+ 
182+ WeightedStreamMerger merger(options_);
183+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
184+ 
185+ EXPECT_EQ(status, graphStatus::FAILED);
186+}
187+ 
188+TEST_F(WeightedStreamMergerTest, Merge_DuplicatedNodeInRoutes_ReturnsFailed) {
189+ BuildDiamondGraph();
190+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {1}, {2}, {3}};
191+ std::vector<int32_t> logical_to_physical;
192+ 
193+ WeightedStreamMerger merger(options_);
194+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
195+ 
196+ EXPECT_EQ(status, graphStatus::FAILED);
197+}
198+ 
199+TEST_F(WeightedStreamMergerTest, Merge_InvalidOptions_ReturnsFailed) {
200+ BuildDiamondGraph();
201+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}};
202+ std::vector<int32_t> logical_to_physical;
203+ 
204+ options_.physical_stream_limit = 0;
205+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
206+ EXPECT_TRUE(logical_to_physical.empty());
207+ 
208+ options_.physical_stream_limit = 2;
209+ options_.window_width = 0;
210+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
211+ 
212+ options_.window_width = 1;
213+ options_.candidate_limit = 0;
214+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
215+ 
216+ options_.candidate_limit = 1;
217+ options_.light_stream_limit = 0;
218+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
219+ 
220+ options_.light_stream_limit = 1;
221+ options_.repair_moves = -1;
222+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
223+ 
224+ options_.repair_moves = 0;
225+ options_.resim_candidate_limit = -1;
226+ EXPECT_EQ(WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical), graphStatus::FAILED);
227+}
228+ 
229+TEST_F(WeightedStreamMergerTest, Merge_WithOriginHintsAndRepair_ReturnsValidMapping) {
230+ BuildIndependentGraph(8, 12.0F);
231+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}};
232+ std::vector<int32_t> logical_to_physical;
233+ options_.physical_stream_limit = 3;
234+ options_.candidate_limit = 3;
235+ options_.light_stream_limit = 2;
236+ options_.repair_moves = 2;
237+ options_.resim_candidate_limit = 3;
238+ 
239+ WeightedStreamMerger merger(options_);
240+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
241+ 
242+ ASSERT_EQ(status, graphStatus::SUCCESS);
243+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
244+ for (const auto stream_id : logical_to_physical) {
245+ EXPECT_GE(stream_id, 0);
246+ EXPECT_LT(stream_id, options_.physical_stream_limit);
247+ }
248+}
249+ 
250+TEST_F(WeightedStreamMergerTest, Merge_WithLimitedCandidates_ReturnsValidMapping) {
251+ BuildIndependentGraph(6, 8.0F);
252+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}, {4}, {5}};
253+ std::vector<int32_t> logical_to_physical;
254+ options_.physical_stream_limit = 2;
255+ options_.candidate_limit = 1;
256+ options_.light_stream_limit = 1;
257+ options_.resim_candidate_limit = 0;
258+ 
259+ WeightedStreamMerger merger(options_);
260+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
261+ 
262+ ASSERT_EQ(status, graphStatus::SUCCESS);
263+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
264+ for (const auto stream_id : logical_to_physical) {
265+ EXPECT_GE(stream_id, 0);
266+ EXPECT_LT(stream_id, options_.physical_stream_limit);
267+ }
268+}
269+ 
270+TEST_F(WeightedStreamMergerTest, Merge_TwoStageParallelGraph_TriggersSimulationTieBreaks) {
271+ BuildTwoStageParallelGraph();
272+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}, {4}, {5}};
273+ std::vector<int32_t> logical_to_physical;
274+ options_.physical_stream_limit = 3;
275+ options_.candidate_limit = 4;
276+ options_.light_stream_limit = 3;
277+ options_.resim_candidate_limit = 3;
278+ 
279+ WeightedStreamMerger merger(options_);
280+ auto status = merger.Merge(*dag_, logical_routes, logical_to_physical);
281+ 
282+ ASSERT_EQ(status, graphStatus::SUCCESS);
283+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
284+ for (const auto stream_id : logical_to_physical) {
285+ EXPECT_GE(stream_id, 0);
286+ EXPECT_LT(stream_id, options_.physical_stream_limit);
287+ }
288+}
289+ 
290+TEST_F(WeightedStreamMergerTest, Merge_EmptyRoutesWithNodes_ReturnsSuccess) {
291+ BuildIndependentGraph(3, 1.0F);
292+ std::vector<int32_t> logical_to_physical;
293+ 
294+ auto status = WeightedStreamMerger(options_).Merge(*dag_, {}, logical_to_physical);
295+ 
296+ EXPECT_EQ(status, graphStatus::SUCCESS);
297+ EXPECT_TRUE(logical_to_physical.empty());
298+}
299+ 
300+TEST_F(WeightedStreamMergerTest, Merge_EmptyLogicalRoute_ReturnsFailed) {
301+ BuildDiamondGraph();
302+ std::vector<int32_t> logical_to_physical;
303+ 
304+ auto status = WeightedStreamMerger(options_).Merge(*dag_, {{0}, {}}, logical_to_physical);
305+ 
306+ EXPECT_EQ(status, graphStatus::FAILED);
307+}
308+ 
309+TEST_F(WeightedStreamMergerTest, Merge_NodeIndexOutOfRange_ReturnsFailed) {
310+ BuildDiamondGraph();
311+ std::vector<int32_t> logical_to_physical;
312+ 
313+ auto status = WeightedStreamMerger(options_).Merge(*dag_, {{0}, {1}, {2}, {4}}, logical_to_physical);
314+ 
315+ EXPECT_EQ(status, graphStatus::FAILED);
316+}
317+ 
318+TEST_F(WeightedStreamMergerTest, Merge_CyclicGraph_ReturnsFailed) {
319+ BuildCycleGraph();
320+ std::vector<int32_t> logical_to_physical;
321+ 
322+ auto status = WeightedStreamMerger(options_).Merge(*dag_, {{0}, {1}}, logical_to_physical);
323+ 
324+ EXPECT_EQ(status, graphStatus::FAILED);
325+}
326+ 
327+TEST_F(WeightedStreamMergerTest, Merge_SinglePhysicalStream_MapsAllToZero) {
328+ BuildIndependentGraph(6, 3.0F);
329+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}, {4}, {5}};
330+ std::vector<int32_t> logical_to_physical;
331+ options_.physical_stream_limit = 1;
332+ 
333+ auto status = WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical);
334+ 
335+ ASSERT_EQ(status, graphStatus::SUCCESS);
336+ ASSERT_EQ(logical_to_physical.size(), logical_routes.size());
337+ for (const auto stream_id : logical_to_physical) {
338+ EXPECT_EQ(stream_id, 0);
339+ }
340+}
341+ 
342+TEST_F(WeightedStreamMergerTest, Merge_CrossStreamDependency_ReturnsValidMapping) {
343+ BuildCrossStreamDependencyGraph();
344+ std::vector<std::vector<int32_t>> logical_routes = {{0}, {1}, {2}, {3}, {4}, {5}};
345+ std::vector<int32_t> logical_to_physical;
346+ options_.physical_stream_limit = 3;
347+ options_.repair_moves = 3;
348+ 
349+ auto status = WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical);
350+ 
351+ ASSERT_EQ(status, graphStatus::SUCCESS);
352+ ExpectValidMapping(logical_to_physical, logical_routes.size());
353+}
354+ 
355+TEST_F(WeightedStreamMergerTest, Merge_MissingAndExtremeCost_ReturnsValidMapping) {
356+ BuildMissingAndExtremeCostGraph();
357+ std::vector<std::vector<int32_t>> logical_routes = {{0, 1}, {2}, {3}};
358+ std::vector<int32_t> logical_to_physical;
359+ options_.physical_stream_limit = 3;
360+ 
361+ auto status = WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical);
362+ 
363+ ASSERT_EQ(status, graphStatus::SUCCESS);
364+ ExpectValidMapping(logical_to_physical, logical_routes.size());
365+}
366+ 
367+TEST_F(WeightedStreamMergerTest, Merge_EdgeFromExternalNode_IgnoresOutsidePred) {
368+ auto n0 = AddNode("n0", 4.0F, 1U, 1U);
369+ auto n1 = AddNode("n1", 6.0F, 2U, 1U);
370+ auto external = std::make_shared<DAGNode>("external", "Dummy");
371+ ASSERT_EQ(dag_->AddEdge(external, 0, n0, 0), graphStatus::SUCCESS);
372+ ASSERT_EQ(dag_->AddEdge(n0, 0, n1, 0), graphStatus::SUCCESS);
373+ std::vector<int32_t> logical_to_physical;
374+ 
375+ auto status = WeightedStreamMerger(options_).Merge(*dag_, {{0}, {1}}, logical_to_physical);
376+ 
377+ ASSERT_EQ(status, graphStatus::SUCCESS);
378+ ExpectValidMapping(logical_to_physical, 2U);
379+}
380+ 
381+TEST_F(WeightedStreamMergerTest, Merge_LargeScaleIndependentGraph_ReturnsValidMapping) {
382+ BuildIndependentGraph(20, 2.0F);
383+ std::vector<std::vector<int32_t>> logical_routes;
384+ for (int32_t idx = 0; idx < 20; ++idx) {
385+ logical_routes.push_back({idx});
386+ }
387+ std::vector<int32_t> logical_to_physical;
388+ options_.physical_stream_limit = 4;
389+ options_.repair_moves = 4;
390+ 
391+ auto status = WeightedStreamMerger(options_).Merge(*dag_, logical_routes, logical_to_physical);
392+ 
393+ ASSERT_EQ(status, graphStatus::SUCCESS);
394+ ExpectValidMapping(logical_to_physical, logical_routes.size());
395+}
396+ 
397+} // namespace test
398+} // namespace minidag
Atests/ge/ut/ge/graph/build/dag/dag_weighted_stream_merger_whitebox_test.cc+704-0
@@ -0,0 +1,704 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <gtest/gtest.h>
12+ 
13+#include <algorithm>
14+#include <cmath>
15+#include <cstdint>
16+#include <deque>
17+#include <functional>
18+#include <limits>
19+#include <map>
20+#include <memory>
21+#include <queue>
22+#include <set>
23+#include <string>
24+#include <unordered_map>
25+#include <utility>
26+#include <vector>
27+ 
28+#include "graph/build/dag/dag_edge.h"
29+#include "graph/build/dag/dag_graph.h"
30+#include "graph/build/dag/dag_log.h"
31+#include "graph/build/dag/dag_node.h"
32+#include "graph/build/dag/dag_types.h"
33+ 
34+#define private public
35+#define WeightedStreamMerger WeightedStreamMergerWhiteBox
36+#include "graph/build/dag/dag_weighted_stream_merger.cc"
37+#undef WeightedStreamMerger
38+#undef private
39+ 
40+namespace minidag {
41+namespace test {
42+namespace {
43+void PrepareTwoUnitSolver(WeightedStreamMergeSolver &solver) {
44+ solver.max_level_ = 2;
45+ solver.unit_count_ = 2;
46+ solver.unit_total_duration_ = {10.0, 1.0};
47+ solver.unit_total_aiv_time_ = {20.0, 1.0};
48+ solver.unit_total_aic_time_ = {10.0, 1.0};
49+ solver.unit_duration_hist_ = {{5.0, 0.0, 5.0}, {0.0, 1.0, 0.0}};
50+ solver.unit_peak_aiv_hist_ = {{2, 0, 1}, {0, 1, 0}};
51+ solver.unit_peak_aic_hist_ = {{1, 0, 3}, {0, 1, 0}};
52+ solver.unit_peak_aiv_ = {2, 1};
53+ solver.unit_peak_aic_ = {3, 1};
54+ solver.unit_origin_stream_hint_ = {-1, -1};
55+ solver.unit_active_levels_ = {{0, 2}, {1}};
56+ solver.unit_pred_edges_ = {{}, {{0, 2}}};
57+ solver.unit_succ_edges_ = {{{1, 3}}, {}};
58+}
59+} // namespace
60+ 
61+TEST(WeightedStreamMergerWhiteBoxTest, RunningNodeGreater_TieBreaksByFinishTopoAndNode) {
62+ WeightedStreamMergeSolver::RunningNodeGreater greater;
63+ 
64+ EXPECT_TRUE(greater({2.0, 0, 0}, {1.0, 0, 0}));
65+ EXPECT_FALSE(greater({1.0, 0, 0}, {2.0, 0, 0}));
66+ EXPECT_TRUE(greater({1.0, 2, 0}, {1.0, 1, 0}));
67+ EXPECT_TRUE(greater({1.0, 1, 2}, {1.0, 1, 1}));
68+}
69+ 
70+TEST(WeightedStreamMergerWhiteBoxTest, RunningNodeHeap_OrdersByFinishTopoAndNode) {
71+ WeightedStreamMergeSolver::RunningNodeHeap heap;
72+ 
73+ heap.push({3.0, 0, 0});
74+ heap.push({1.0, 2, 0});
75+ heap.push({1.0, 1, 3});
76+ heap.push({1.0, 1, 2});
77+ 
78+ EXPECT_EQ(heap.top().node_index, 2);
79+ heap.pop();
80+ EXPECT_EQ(heap.top().node_index, 3);
81+ heap.pop();
82+ EXPECT_EQ(heap.top().node_index, 0);
83+}
84+ 
85+TEST(WeightedStreamMergerWhiteBoxTest, LiteScoreLess_TieBreaksByConflictEventLoadAndStream) {
86+ LiteScore lhs;
87+ LiteScore rhs;
88+ lhs.total = 1.0;
89+ rhs.total = 1.0;
90+ lhs.time_conflict = 1.0;
91+ rhs.time_conflict = 2.0;
92+ EXPECT_TRUE(WeightedStreamMergeSolver::LiteScoreLess(lhs, rhs));
93+ 
94+ rhs.time_conflict = lhs.time_conflict;
95+ lhs.event_local = 1;
96+ rhs.event_local = 2;
97+ EXPECT_TRUE(WeightedStreamMergeSolver::LiteScoreLess(lhs, rhs));
98+ 
99+ rhs.event_local = lhs.event_local;
100+ lhs.time_load = 3.0;
101+ rhs.time_load = 4.0;
102+ EXPECT_TRUE(WeightedStreamMergeSolver::LiteScoreLess(lhs, rhs));
103+ 
104+ rhs.time_load = lhs.time_load;
105+ lhs.candidate_stream = 0;
106+ rhs.candidate_stream = 1;
107+ EXPECT_TRUE(WeightedStreamMergeSolver::LiteScoreLess(lhs, rhs));
108+}
109+ 
110+TEST(WeightedStreamMergerWhiteBoxTest, OriginHintFlows_SortsByCountLoadAndStreamId) {
111+ DAGGraph dag("whitebox_origin_hint");
112+ std::vector<std::vector<int32_t>> routes;
113+ WeightedStreamMergeOptions options;
114+ WeightedStreamMergeSolver solver(dag, routes, options);
115+ solver.unit_origin_stream_hint_ = {7};
116+ 
117+ StreamingState state;
118+ state.stream_total_durations = {5.0, 4.0, 2.0, 2.0};
119+ state.origin_stream_to_flow_counts[7][0] = 3;
120+ state.origin_stream_to_flow_counts[7][1] = 1;
121+ state.origin_stream_to_flow_counts[7][2] = 1;
122+ state.origin_stream_to_flow_counts[7][3] = 1;
123+ 
124+ const auto flows = solver.OriginHintFlows(0, state);
125+ 
126+ ASSERT_EQ(flows.size(), 4U);
127+ EXPECT_EQ(flows[0], 0);
128+ EXPECT_EQ(flows[1], 2);
129+ EXPECT_EQ(flows[2], 3);
130+ EXPECT_EQ(flows[3], 1);
131+}
132+ 
133+TEST(WeightedStreamMergerWhiteBoxTest, CandidateFlows_AddsOriginHintsBeforeLightFlows) {
134+ DAGGraph dag("whitebox_origin_hint_candidate");
135+ std::vector<std::vector<int32_t>> routes;
136+ WeightedStreamMergeOptions options;
137+ options.physical_stream_limit = 4;
138+ options.candidate_limit = 4;
139+ options.light_stream_limit = 2;
140+ WeightedStreamMergeSolver solver(dag, routes, options);
141+ PrepareTwoUnitSolver(solver);
142+ solver.unit_origin_stream_hint_ = {9, -1};
143+ solver.unit_pred_edges_ = {{}, {}};
144+ solver.unit_succ_edges_ = {{}, {}};
145+ 
146+ StreamingState state;
147+ state.assignment = {-1, -1};
148+ state.stream_total_durations = {8.0, 4.0, 2.0, 0.0};
149+ state.stream_duration_hists = {{3.0, 0.0, 0.0}, {2.0, 0.0, 0.0}, {1.0, 0.0, 0.0}, {0.0, 0.0, 0.0}};
150+ state.origin_stream_to_flow_counts[9][0] = 1;
151+ state.origin_stream_to_flow_counts[9][1] = 2;
152+ state.origin_stream_to_flow_counts[9][2] = 1;
153+ state.active_stream_count = 3;
154+ 
155+ EXPECT_EQ(solver.CandidateFlows(0, 0, state, true), std::vector<int32_t>({1, 2, 0, 3}));
156+}
157+ 
158+TEST(WeightedStreamMergerWhiteBoxTest, OriginHintFlows_ReturnsEmptyWithoutHintOrCounter) {
159+ DAGGraph dag("whitebox_origin_hint_empty");
160+ std::vector<std::vector<int32_t>> routes;
161+ WeightedStreamMergeOptions options;
162+ WeightedStreamMergeSolver solver(dag, routes, options);
163+ StreamingState state;
164+ 
165+ solver.unit_origin_stream_hint_ = {-1};
166+ EXPECT_TRUE(solver.OriginHintFlows(0, state).empty());
167+ 
168+ solver.unit_origin_stream_hint_ = {3};
169+ EXPECT_TRUE(solver.OriginHintFlows(0, state).empty());
170+}
171+ 
172+TEST(WeightedStreamMergerWhiteBoxTest, AdjacentAndLightCandidateFlows_SortAndLimitCandidates) {
173+ DAGGraph dag("whitebox_candidates");
174+ std::vector<std::vector<int32_t>> routes;
175+ WeightedStreamMergeOptions options;
176+ options.physical_stream_limit = 4;
177+ options.candidate_limit = 3;
178+ options.light_stream_limit = 2;
179+ WeightedStreamMergeSolver solver(dag, routes, options);
180+ PrepareTwoUnitSolver(solver);
181+ 
182+ StreamingState state;
183+ state.assignment = {0, 2};
184+ state.stream_total_durations = {8.0, 1.0, 3.0, 0.0};
185+ state.stream_duration_hists = {{4.0, 4.0, 0.0}, {1.0, 0.0, 0.0}, {2.0, 0.0, 1.0}, {0.0, 0.0, 0.0}};
186+ state.active_stream_count = 3;
187+ 
188+ const std::map<int32_t, int32_t> expected_unit0_flows = {{2, 3}};
189+ const std::map<int32_t, int32_t> expected_unit1_flows = {{0, 2}};
190+ EXPECT_EQ(solver.AdjacentAssignedFlows(0, state.assignment), expected_unit0_flows);
191+ EXPECT_EQ(solver.AdjacentAssignedFlows(1, state.assignment), expected_unit1_flows);
192+ EXPECT_EQ(solver.AdjacentCandidateFlows(0, state), std::vector<int32_t>({2}));
193+ EXPECT_EQ(solver.LightCandidateFlows(0, state), std::vector<int32_t>({1, 2}));
194+ EXPECT_EQ(solver.UniqueLimitedCandidates({2, 1, 2, 0}), std::vector<int32_t>({2, 1, 0}));
195+ EXPECT_EQ(solver.CandidateFlows(0, 0, state, true), std::vector<int32_t>({2, 1, 3}));
196+}
197+ 
198+TEST(WeightedStreamMergerWhiteBoxTest, WindowHelpers_ReturnExpectedLevelsAndLoad) {
199+ DAGGraph dag("whitebox_window");
200+ std::vector<std::vector<int32_t>> routes;
201+ WeightedStreamMergeOptions options;
202+ options.window_width = 2;
203+ options.physical_stream_limit = 2;
204+ WeightedStreamMergeSolver solver(dag, routes, options);
205+ PrepareTwoUnitSolver(solver);
206+ 
207+ StreamingState state;
208+ state.stream_duration_hists = {{1.0, 2.0, 4.0}, {0.0, 3.0, 5.0}};
209+ 
210+ EXPECT_EQ(solver.WindowEnd(0), 1);
211+ EXPECT_EQ(solver.WindowEnd(2), 2);
212+ EXPECT_DOUBLE_EQ(solver.WindowMass(state.stream_duration_hists[0], 0), 3.0);
213+ EXPECT_DOUBLE_EQ(solver.StreamWindowDurationLoad(state, 1, 1), 8.0);
214+ EXPECT_EQ(solver.WindowActiveLevels(0, 1), std::vector<int32_t>({2}));
215+ EXPECT_EQ(solver.WindowActiveLevels(1, 2), std::vector<int32_t>({1}));
216+}
217+ 
218+TEST(WeightedStreamMergerWhiteBoxTest, FallbackCandidateFlows_ReturnsNewOrExistingStreams) {
219+ DAGGraph dag("whitebox_fallback");
220+ std::vector<std::vector<int32_t>> routes;
221+ WeightedStreamMergeOptions options;
222+ options.physical_stream_limit = 3;
223+ options.candidate_limit = 2;
224+ WeightedStreamMergeSolver solver(dag, routes, options);
225+ 
226+ StreamingState state;
227+ state.active_stream_count = 1;
228+ EXPECT_EQ(solver.FallbackCandidateFlows(state, true), std::vector<int32_t>({1}));
229+ 
230+ state.active_stream_count = 3;
231+ EXPECT_EQ(solver.FallbackCandidateFlows(state, false), std::vector<int32_t>({0, 1}));
232+}
233+ 
234+TEST(WeightedStreamMergerWhiteBoxTest, CandidateFlows_FallsBackWhenNoRegularCandidateExists) {
235+ DAGGraph dag("whitebox_candidate_fallback");
236+ std::vector<std::vector<int32_t>> routes;
237+ WeightedStreamMergeOptions options;
238+ options.physical_stream_limit = 2;
239+ options.candidate_limit = 2;
240+ options.light_stream_limit = 1;
241+ WeightedStreamMergeSolver solver(dag, routes, options);
242+ PrepareTwoUnitSolver(solver);
243+ solver.unit_origin_stream_hint_ = {-1, -1};
244+ solver.unit_pred_edges_ = {{}, {}};
245+ solver.unit_succ_edges_ = {{}, {}};
246+ 
247+ StreamingState state;
248+ state.assignment = {-1, -1};
249+ state.stream_total_durations = {0.0, 0.0};
250+ state.stream_duration_hists = {{0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}};
251+ state.active_stream_count = 0;
252+ 
253+ EXPECT_EQ(solver.CandidateFlows(0, 0, state, true), std::vector<int32_t>({0}));
254+ 
255+ state.active_stream_count = 2;
256+ EXPECT_EQ(solver.CandidateFlows(0, 0, state, false), std::vector<int32_t>({0}));
257+}
258+ 
259+TEST(WeightedStreamMergerWhiteBoxTest, ScoreApplyAndRemove_UpdateStateAndCounters) {
260+ DAGGraph dag("whitebox_score");
261+ std::vector<std::vector<int32_t>> routes;
262+ WeightedStreamMergeOptions options;
263+ options.physical_stream_limit = 3;
264+ options.window_width = 3;
265+ WeightedStreamMergeSolver solver(dag, routes, options);
266+ PrepareTwoUnitSolver(solver);
267+ solver.unit_origin_stream_hint_ = {7, 9};
268+ 
269+ StreamingState state;
270+ state.assignment = {-1, 1};
271+ state.stream_total_durations = {3.0, 2.0, 0.0};
272+ state.stream_duration_hists = {{1.0, 1.0, 1.0}, {0.0, 2.0, 0.0}, {0.0, 0.0, 0.0}};
273+ state.active_stream_count = 2;
274+ 
275+ const auto score = solver.EvaluateLiteScore(0, 2, 0, state);
276+ EXPECT_EQ(score.event_local, 3);
277+ EXPECT_DOUBLE_EQ(score.time_conflict, 0.0);
278+ EXPECT_DOUBLE_EQ(score.time_load, 10.0);
279+ EXPECT_EQ(score.candidate_stream, 2);
280+ EXPECT_GT(score.total, score.time_load);
281+ 
282+ EXPECT_TRUE(solver.ApplyUnitToStream(0, 2, state));
283+ EXPECT_EQ(state.active_stream_count, 3);
284+ EXPECT_EQ(state.assignment[0], 2);
285+ EXPECT_DOUBLE_EQ(state.stream_total_durations[2], 10.0);
286+ EXPECT_EQ(state.origin_stream_to_flow_counts[7][2], 1);
287+ 
288+ EXPECT_EQ(solver.RemoveUnitFromStream(0, state), 2);
289+ EXPECT_EQ(state.assignment[0], -1);
290+ EXPECT_DOUBLE_EQ(state.stream_total_durations[2], 0.0);
291+ EXPECT_EQ(state.origin_stream_to_flow_counts.count(7), 0U);
292+}
293+ 
294+TEST(WeightedStreamMergerWhiteBoxTest, RepairRecentUnits_MovesUnitWhenScoreImproves) {
295+ DAGGraph dag("whitebox_repair");
296+ std::vector<std::vector<int32_t>> routes;
297+ WeightedStreamMergeOptions options;
298+ options.physical_stream_limit = 2;
299+ options.candidate_limit = 1;
300+ options.light_stream_limit = 2;
301+ options.repair_moves = 1;
302+ options.resim_candidate_limit = 0;
303+ WeightedStreamMergeSolver solver(dag, routes, options);
304+ 
305+ solver.max_level_ = 0;
306+ solver.unit_active_levels_ = {{0}, {0}};
307+ solver.unit_duration_hist_ = {{10.0}, {1.0}};
308+ solver.unit_total_duration_ = {10.0, 1.0};
309+ solver.unit_origin_stream_hint_ = {-1, -1};
310+ solver.unit_pred_edges_ = {{}, {}};
311+ solver.unit_succ_edges_ = {{}, {}};
312+ 
313+ StreamingState state;
314+ state.assignment = {0, 0};
315+ state.stream_total_durations = {11.0, 0.0};
316+ state.stream_duration_hists = {{11.0}, {0.0}};
317+ state.active_stream_count = 2;
318+ 
319+ solver.RepairRecentUnits(0, {0}, state);
320+ 
321+ EXPECT_EQ(state.assignment[0], 1);
322+ EXPECT_EQ(state.assignment[1], 0);
323+}
324+ 
325+TEST(WeightedStreamMergerWhiteBoxTest, RepairRecentUnits_ReturnsWhenRecentUnitIsUnassigned) {
326+ DAGGraph dag("whitebox_repair_unassigned");
327+ std::vector<std::vector<int32_t>> routes;
328+ WeightedStreamMergeOptions options;
329+ options.physical_stream_limit = 1;
330+ options.candidate_limit = 1;
331+ options.light_stream_limit = 1;
332+ options.repair_moves = 1;
333+ WeightedStreamMergeSolver solver(dag, routes, options);
334+ 
335+ solver.max_level_ = 0;
336+ solver.unit_active_levels_ = {{0}};
337+ solver.unit_duration_hist_ = {{1.0}};
338+ solver.unit_total_duration_ = {1.0};
339+ solver.unit_origin_stream_hint_ = {-1};
340+ solver.unit_pred_edges_ = {{}};
341+ solver.unit_succ_edges_ = {{}};
342+ 
343+ StreamingState state;
344+ state.assignment = {-1};
345+ state.stream_total_durations = {0.0};
346+ state.stream_duration_hists = {{0.0}};
347+ state.active_stream_count = 1;
348+ 
349+ solver.RepairRecentUnits(0, {0}, state);
350+ 
351+ EXPECT_EQ(state.assignment[0], -1);
352+}
353+ 
354+TEST(WeightedStreamMergerWhiteBoxTest, RepairRecentUnits_ReturnsWhenNoMoveImproves) {
355+ DAGGraph dag("whitebox_repair_no_move");
356+ std::vector<std::vector<int32_t>> routes;
357+ WeightedStreamMergeOptions options;
358+ options.physical_stream_limit = 2;
359+ options.candidate_limit = 2;
360+ options.light_stream_limit = 2;
361+ options.repair_moves = 1;
362+ options.resim_candidate_limit = 0;
363+ WeightedStreamMergeSolver solver(dag, routes, options);
364+ 
365+ solver.max_level_ = 0;
366+ solver.unit_active_levels_ = {{0}, {0}};
367+ solver.unit_duration_hist_ = {{10.0}, {1.0}};
368+ solver.unit_total_duration_ = {10.0, 1.0};
369+ solver.unit_origin_stream_hint_ = {-1, -1};
370+ solver.unit_pred_edges_ = {{}, {}};
371+ solver.unit_succ_edges_ = {{}, {}};
372+ 
373+ StreamingState state;
374+ state.assignment = {0, 1};
375+ state.stream_total_durations = {10.0, 1.0};
376+ state.stream_duration_hists = {{10.0}, {1.0}};
377+ state.active_stream_count = 2;
378+ 
379+ solver.RepairRecentUnits(0, {0}, state);
380+ 
381+ EXPECT_EQ(state.assignment[0], 0);
382+ EXPECT_EQ(state.assignment[1], 1);
383+}
384+ 
385+TEST(WeightedStreamMergerWhiteBoxTest, SimulationHelpers_HandleSuccessFailureAndTieBreaks) {
386+ DAGGraph dag("whitebox_simulation");
387+ std::vector<std::vector<int32_t>> routes;
388+ WeightedStreamMergeOptions options;
389+ options.physical_stream_limit = 2;
390+ WeightedStreamMergeSolver solver(dag, routes, options);
391+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy"), std::make_shared<DAGNode>("n1", "Dummy"),
392+ std::make_shared<DAGNode>("n2", "Dummy")};
393+ solver.topo_order_ = {0, 1, 2};
394+ solver.topo_position_ = {0, 1, 2};
395+ solver.node_to_unit_ = {0, 1, 2};
396+ solver.node_durations_ = {5.0, 3.0, 3.0};
397+ solver.node_bottom_ranks_ = {8.0, 3.0, 3.0};
398+ solver.preds_ = {{}, {0}, {}};
399+ 
400+ std::vector<int32_t> node_streams(3, -1);
401+ std::vector<char> active_mask(3, 0);
402+ EXPECT_EQ(solver.BuildSimulationNodeStreams({0, -1, 1}, false, node_streams, active_mask), 2);
403+ EXPECT_EQ(node_streams, std::vector<int32_t>({0, -1, 1}));
404+ EXPECT_EQ(active_mask, std::vector<char>({1, 0, 1}));
405+ 
406+ node_streams.assign(3, -1);
407+ active_mask.assign(3, 0);
408+ EXPECT_EQ(solver.BuildSimulationNodeStreams({0, -1, 1}, true, node_streams, active_mask), -1);
409+ 
410+ auto queues = solver.BuildSimulationQueues({0, 1, 0}, {1, 1, 0});
411+ ASSERT_EQ(queues.size(), 2U);
412+ EXPECT_EQ(queues[0], std::deque<int32_t>({0}));
413+ EXPECT_EQ(queues[1], std::deque<int32_t>({1}));
414+ 
415+ auto filtered_preds = solver.BuildSimulationFilteredPreds({1, 1, 0});
416+ EXPECT_EQ(filtered_preds[1], std::vector<int32_t>({0}));
417+ EXPECT_TRUE(filtered_preds[2].empty());
418+ 
419+ EXPECT_FALSE(solver.AreSimulationPredsFinished(1, filtered_preds, {0, 0, 0}));
420+ EXPECT_TRUE(solver.AreSimulationPredsFinished(1, filtered_preds, {1, 0, 0}));
421+ 
422+ std::vector<int32_t> ready_heads = {2, 1};
423+ solver.SortSimulationReadyHeads(ready_heads);
424+ EXPECT_EQ(ready_heads, std::vector<int32_t>({1, 2}));
425+ 
426+ auto result = solver.BuildAndRunSimulation({0, 1, 0}, true);
427+ EXPECT_TRUE(result.feasible);
428+ EXPECT_DOUBLE_EQ(result.makespan, 11.0);
429+ 
430+ result = solver.BuildAndRunSimulation({0, -1, 0}, true);
431+ EXPECT_FALSE(result.feasible);
432+}
433+ 
434+TEST(WeightedStreamMergerWhiteBoxTest, ResourceProfilesAndUnitOrder_UpdateDerivedFields) {
435+ DAGGraph dag("whitebox_resource_profiles");
436+ std::vector<std::vector<int32_t>> routes;
437+ WeightedStreamMergeOptions options;
438+ options.physical_stream_limit = 3;
439+ WeightedStreamMergeSolver solver(dag, routes, options);
440+ solver.max_level_ = 2;
441+ solver.unit_count_ = 3;
442+ solver.levels_ = {0, 1, 1, 2};
443+ solver.node_durations_ = {2.0, 5.0, 3.0, 7.0};
444+ solver.node_aiv_cores_ = {1, 2, 3, 1};
445+ solver.node_aic_cores_ = {2, 1, 1, 4};
446+ solver.node_origin_stream_hints_ = {4, 4, 2, 2};
447+ solver.node_bottom_ranks_ = {10.0, 8.0, 6.0, 7.0};
448+ solver.unit_profiles_ = {
449+ {0, {0, 1}, 2, 0, 1, 1, {1, 1, 0}, 0}, {1, {2}, 1, 1, 1, 1, {0, 1, 0}, 0}, {2, {3}, 1, 1, 1, 1, {0, 1, 0}, 0}};
450+ 
451+ solver.InitResourceUnitProfiles();
452+ solver.BuildResourceUnitProfile(solver.unit_profiles_[0]);
453+ 
454+ EXPECT_DOUBLE_EQ(solver.unit_total_duration_[0], 7.0);
455+ EXPECT_DOUBLE_EQ(solver.unit_total_aiv_time_[0], 12.0);
456+ EXPECT_DOUBLE_EQ(solver.unit_total_aic_time_[0], 9.0);
457+ EXPECT_EQ(solver.unit_peak_aiv_[0], 2);
458+ EXPECT_EQ(solver.unit_peak_aic_[0], 2);
459+ EXPECT_EQ(solver.unit_origin_stream_hint_[0], 4);
460+ EXPECT_EQ(solver.unit_active_levels_[0], std::vector<int32_t>({0, 1}));
461+ 
462+ EXPECT_EQ(solver.SelectOriginStreamHint({{5, 1}, {3, 1}}), 3);
463+ 
464+ solver.unit_total_duration_ = {10.0, 10.0, 10.0};
465+ solver.unit_total_aiv_time_ = {1.0, 2.0, 2.0};
466+ solver.unit_total_aic_time_ = {1.0, 1.0, 1.0};
467+ solver.unit_peak_aiv_ = {1, 1, 1};
468+ solver.unit_peak_aic_ = {1, 1, 1};
469+ solver.unit_profiles_[0].earliest_level = 1;
470+ solver.unit_profiles_[0].interaction = 0;
471+ solver.unit_profiles_[1].interaction = 2;
472+ solver.unit_profiles_[2].interaction = 1;
473+ solver.unit_profiles_[2].earliest_level = 0;
474+ solver.BuildUnitsByEarliestLevel();
475+ 
476+ ASSERT_EQ(solver.units_by_earliest_level_[1].size(), 2U);
477+ EXPECT_EQ(solver.units_by_earliest_level_[1][0], 1);
478+ EXPECT_EQ(solver.units_by_earliest_level_[1][1], 0);
479+ EXPECT_EQ(solver.units_by_earliest_level_[0], std::vector<int32_t>({2}));
480+}
481+ 
482+TEST(WeightedStreamMergerWhiteBoxTest, Solve_EmptyDagClearsMapping) {
483+ DAGGraph dag("whitebox_empty_solve");
484+ std::vector<std::vector<int32_t>> routes;
485+ WeightedStreamMergeOptions options;
486+ WeightedStreamMergeSolver solver(dag, routes, options);
487+ std::vector<int32_t> mapping = {1, 2};
488+ 
489+ EXPECT_EQ(solver.Solve(mapping), graphStatus::SUCCESS);
490+ EXPECT_TRUE(mapping.empty());
491+}
492+ 
493+TEST(WeightedStreamMergerWhiteBoxTest, UnitProfile_SortsRouteNodesByTopoPosition) {
494+ DAGGraph dag("whitebox_route_sort");
495+ std::vector<std::vector<int32_t>> routes = {{1, 0}};
496+ WeightedStreamMergeOptions options;
497+ WeightedStreamMergeSolver solver(dag, routes, options);
498+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy"), std::make_shared<DAGNode>("n1", "Dummy")};
499+ solver.topo_position_ = {0, 1};
500+ solver.levels_ = {0, 1};
501+ solver.max_level_ = 1;
502+ solver.node_to_unit_.assign(2U, -1);
503+ 
504+ ASSERT_EQ(solver.BuildUnitProfile(0), graphStatus::SUCCESS);
505+ EXPECT_EQ(solver.unit_profiles_[0].node_indices, std::vector<int32_t>({0, 1}));
506+ EXPECT_EQ(solver.unit_profiles_[0].earliest_level, 0);
507+ EXPECT_EQ(solver.unit_profiles_[0].latest_level, 1);
508+}
509+ 
510+TEST(WeightedStreamMergerWhiteBoxTest, SizeToInt32_ClampsLargeValue) {
511+ const auto max_int = std::numeric_limits<int32_t>::max();
512+ 
513+ EXPECT_EQ(WeightedStreamMergeSolver::SizeToInt32(16U), 16);
514+ EXPECT_EQ(WeightedStreamMergeSolver::SizeToInt32(std::numeric_limits<size_t>::max()), max_int);
515+}
516+ 
517+TEST(WeightedStreamMergerWhiteBoxTest, SimulationStep_ReturnsFalseForBlockedAndInconsistentState) {
518+ DAGGraph dag("whitebox_simulation_step");
519+ std::vector<std::vector<int32_t>> routes;
520+ WeightedStreamMergeOptions options;
521+ options.physical_stream_limit = 1;
522+ WeightedStreamMergeSolver solver(dag, routes, options);
523+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy")};
524+ solver.topo_order_ = {0};
525+ solver.topo_position_ = {0};
526+ solver.node_durations_ = {1.0};
527+ solver.node_bottom_ranks_ = {1.0};
528+ 
529+ WeightedStreamMergeSolver::SimulationContext context;
530+ context.node_streams = {0};
531+ context.stream_to_queue[0] = {0};
532+ context.filtered_preds = {{}};
533+ context.finished = {0};
534+ context.running_in_stream = {1};
535+ EXPECT_FALSE(solver.RunSimulationStep(context));
536+ 
537+ EXPECT_TRUE(solver.StartSimulationReadyHead({}, context));
538+ EXPECT_TRUE(context.running_heap.empty());
539+ 
540+ context.running_in_stream = {0};
541+ context.stream_to_queue[0].clear();
542+ EXPECT_FALSE(solver.StartSimulationReadyHead({0}, context));
543+}
544+ 
545+TEST(WeightedStreamMergerWhiteBoxTest, FinishNextSimulationNodes_FinishesSameTimeNodes) {
546+ DAGGraph dag("whitebox_finish_same_time");
547+ std::vector<std::vector<int32_t>> routes;
548+ WeightedStreamMergeOptions options;
549+ options.physical_stream_limit = 2;
550+ WeightedStreamMergeSolver solver(dag, routes, options);
551+ 
552+ WeightedStreamMergeSolver::SimulationContext context;
553+ context.node_streams = {0, 1};
554+ context.finished = {0, 0};
555+ context.running_in_stream = {1, 1};
556+ context.running_heap.push({2.0, 0, 0});
557+ context.running_heap.push({2.0, 1, 1});
558+ 
559+ solver.FinishNextSimulationNodes(context);
560+ 
561+ EXPECT_EQ(context.executed_count, 2);
562+ EXPECT_EQ(context.current_time, 2.0);
563+ EXPECT_EQ(context.finished, std::vector<char>({1, 1}));
564+ EXPECT_EQ(context.running_in_stream, std::vector<char>({0, 0}));
565+}
566+ 
567+TEST(WeightedStreamMergerWhiteBoxTest, SelectBestStream_UsesSimulationResultTieBreaks) {
568+ DAGGraph dag("whitebox_select_best");
569+ std::vector<std::vector<int32_t>> routes;
570+ WeightedStreamMergeOptions options;
571+ options.physical_stream_limit = 2;
572+ options.resim_candidate_limit = 2;
573+ options.new_flow_penalty_weight = 0.0;
574+ WeightedStreamMergeSolver solver(dag, routes, options);
575+ PrepareTwoUnitSolver(solver);
576+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy"), std::make_shared<DAGNode>("n1", "Dummy")};
577+ solver.topo_order_ = {0, 1};
578+ solver.topo_position_ = {0, 1};
579+ solver.node_to_unit_ = {0, 1};
580+ solver.node_durations_ = {10.0, 1.0};
581+ solver.node_bottom_ranks_ = {11.0, 1.0};
582+ solver.preds_ = {{}, {0}};
583+ 
584+ StreamingState state;
585+ state.assignment = {-1, 0};
586+ state.stream_total_durations = {1.0, 0.0};
587+ state.stream_duration_hists = {{0.0, 1.0, 0.0}, {0.0, 0.0, 0.0}};
588+ state.active_stream_count = 1;
589+ 
590+ EXPECT_EQ(solver.SelectBestStream(0, 0, {0, 1}, state), 0);
591+}
592+ 
593+TEST(WeightedStreamMergerWhiteBoxTest, SelectBestStream_FallsBackWhenAllResimulationsAreInfeasible) {
594+ DAGGraph dag("whitebox_select_best_infeasible");
595+ std::vector<std::vector<int32_t>> routes;
596+ WeightedStreamMergeOptions options;
597+ options.physical_stream_limit = 2;
598+ options.resim_candidate_limit = 2;
599+ options.new_flow_penalty_weight = 0.0;
600+ WeightedStreamMergeSolver solver(dag, routes, options);
601+ PrepareTwoUnitSolver(solver);
602+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy"), std::make_shared<DAGNode>("n1", "Dummy")};
603+ solver.topo_order_ = {0, 1};
604+ solver.topo_position_ = {0, 1};
605+ solver.node_to_unit_ = {0, 1};
606+ solver.node_durations_ = {10.0, 1.0};
607+ solver.node_bottom_ranks_ = {11.0, 1.0};
608+ solver.preds_ = {{1}, {0}};
609+ 
610+ StreamingState state;
611+ state.assignment = {-1, 0};
612+ state.stream_total_durations = {1.0, 0.0};
613+ state.stream_duration_hists = {{0.0, 1.0, 0.0}, {0.0, 0.0, 0.0}};
614+ state.active_stream_count = 1;
615+ 
616+ EXPECT_EQ(solver.SelectBestStream(0, 0, {0, 1}, state), 0);
617+}
618+ 
619+TEST(WeightedStreamMergerWhiteBoxTest, BuildUnitDependencyGraph_IgnoresSameUnitEdges) {
620+ DAGGraph dag("whitebox_unit_dependency");
621+ std::vector<std::vector<int32_t>> routes;
622+ WeightedStreamMergeOptions options;
623+ WeightedStreamMergeSolver solver(dag, routes, options);
624+ solver.unit_count_ = 2;
625+ solver.node_to_unit_ = {0, 0, 1};
626+ solver.edge_pairs_ = {{0, 1}, {1, 2}, {2, 0}};
627+ solver.unit_profiles_ = {{0, {0, 1}, 2, 0, 1, 1, {1, 1}, 0}, {1, {2}, 1, 1, 1, 1, {0, 1}, 0}};
628+ 
629+ solver.BuildUnitDependencyGraph();
630+ 
631+ EXPECT_EQ(solver.unit_succ_edges_[0].count(0), 0U);
632+ EXPECT_EQ(solver.unit_succ_edges_[0][1], 1);
633+ EXPECT_EQ(solver.unit_pred_edges_[1][0], 1);
634+ EXPECT_EQ(solver.unit_succ_edges_[1][0], 1);
635+ EXPECT_EQ(solver.unit_pred_edges_[0][1], 1);
636+ EXPECT_EQ(solver.unit_profiles_[0].interaction, 2);
637+ EXPECT_EQ(solver.unit_profiles_[1].interaction, 2);
638+}
639+ 
640+TEST(WeightedStreamMergerWhiteBoxTest, SimulationReadyHeads_FiltersAndSortsQueues) {
641+ DAGGraph dag("whitebox_ready_heads");
642+ std::vector<std::vector<int32_t>> routes;
643+ WeightedStreamMergeOptions options;
644+ WeightedStreamMergeSolver solver(dag, routes, options);
645+ solver.topo_position_ = {0, 1, 2, 3, 4};
646+ solver.node_bottom_ranks_ = {0.0, 1.0, 1.0, 4.0, 4.0};
647+ solver.node_durations_ = {0.0, 1.0, 1.0, 2.0, 3.0};
648+ 
649+ const std::map<int32_t, std::deque<int32_t>> queues = {{0, {}}, {1, {1}}, {2, {2}}, {3, {3}}, {4, {4}}};
650+ const std::vector<std::vector<int32_t>> filtered_preds = {{}, {0}, {}, {}, {}};
651+ const std::vector<char> finished = {0, 0, 0, 0, 0};
652+ const std::vector<char> running = {0, 0, 1, 0, 0};
653+ 
654+ EXPECT_EQ(solver.SimulationReadyHeads(queues, filtered_preds, finished, running), std::vector<int32_t>({4, 3}));
655+}
656+ 
657+TEST(WeightedStreamMergerWhiteBoxTest, BuildAndRunSimulation_ReturnsZeroWhenNoNodeActive) {
658+ DAGGraph dag("whitebox_empty_simulation");
659+ std::vector<std::vector<int32_t>> routes;
660+ WeightedStreamMergeOptions options;
661+ WeightedStreamMergeSolver solver(dag, routes, options);
662+ solver.topo_nodes_ = {std::make_shared<DAGNode>("n0", "Dummy"), std::make_shared<DAGNode>("n1", "Dummy")};
663+ solver.node_to_unit_ = {0, 1};
664+ 
665+ const auto result = solver.BuildAndRunSimulation({-1, -1}, false);
666+ 
667+ EXPECT_TRUE(result.feasible);
668+ EXPECT_DOUBLE_EQ(result.makespan, 0.0);
669+}
670+ 
671+TEST(WeightedStreamMergerWhiteBoxTest, RemoveUnitFromStream_HandlesMissingOriginCounter) {
672+ DAGGraph dag("whitebox_remove_missing_origin");
673+ std::vector<std::vector<int32_t>> routes;
674+ WeightedStreamMergeOptions options;
675+ WeightedStreamMergeSolver solver(dag, routes, options);
676+ solver.unit_total_duration_ = {2.0};
677+ solver.unit_duration_hist_ = {{2.0}};
678+ solver.unit_active_levels_ = {{0}};
679+ solver.unit_origin_stream_hint_ = {8};
680+ 
681+ StreamingState state;
682+ state.assignment = {0};
683+ state.stream_total_durations = {2.0};
684+ state.stream_duration_hists = {{2.0}};
685+ state.origin_stream_to_flow_counts[8][1] = 1;
686+ 
687+ EXPECT_EQ(solver.RemoveUnitFromStream(0, state), 0);
688+ EXPECT_EQ(state.origin_stream_to_flow_counts[8][1], 1);
689+ EXPECT_EQ(state.assignment[0], -1);
690+}
691+ 
692+TEST(WeightedStreamMergerWhiteBoxTest, CompactAssignment_CompactsAndRejectsInvalidAssignments) {
693+ DAGGraph dag("whitebox_compact");
694+ std::vector<std::vector<int32_t>> routes;
695+ WeightedStreamMergeOptions options;
696+ WeightedStreamMergeSolver solver(dag, routes, options);
697+ 
698+ EXPECT_EQ(solver.CompactAssignment({2, 5, 2}), std::vector<int32_t>({0, 1, 0}));
699+ EXPECT_TRUE(solver.CompactAssignment({}).empty());
700+ EXPECT_TRUE(solver.CompactAssignment({0, -1}).empty());
701+}
702+ 
703+} // namespace test
704+} // namespace minidag