diff --git a/hercules_opt/src/fork_transforms.rs b/hercules_opt/src/fork_transforms.rs index ec111e69d3f69270beffc0ca92ad3be67e6ff765..9fceb89232b957939adf172493e2868c58c9e033 100644 --- a/hercules_opt/src/fork_transforms.rs +++ b/hercules_opt/src/fork_transforms.rs @@ -719,6 +719,11 @@ pub fn fork_coalesce_helper( }; let new_fork = edit.add_node(new_fork); + if edit.get_schedule(outer_fork).contains(&Schedule::ParallelFork) + && edit.get_schedule(inner_fork).contains(&Schedule::ParallelFork) { + edit = edit.add_schedule(new_fork, Schedule::ParallelFork)?; + } + edit = edit.replace_all_uses(inner_fork, new_fork)?; edit = edit.replace_all_uses(outer_fork, new_fork)?; edit = edit.replace_all_uses(outer_join, inner_join)?; @@ -1271,6 +1276,9 @@ fn fork_interchange( edit = edit.delete_node(old_id)?; } let new_fork = edit.add_node(new_fork); + if edit.get_schedule(fork).contains(&Schedule::ParallelFork) { + edit = edit.add_schedule(new_fork, Schedule::ParallelFork)?; + } edit = edit.replace_all_uses(fork, new_fork)?; edit.delete_node(fork) });