From 972c1e6ebe2f616181f4a8f7a5d90c500f9ac518 Mon Sep 17 00:00:00 2001
From: Aaron Councilman <aaronjc4@illinois.edu>
Date: Sat, 22 Feb 2025 18:42:41 -0600
Subject: [PATCH] Propagate ParallelFork in coalesce and interchange
---
hercules_opt/src/fork_transforms.rs | 8 ++++++++
1 file changed, 8 insertions(+)
diff --git a/hercules_opt/src/fork_transforms.rs b/hercules_opt/src/fork_transforms.rs
index ec111e69..9fceb892 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)
});
--
GitLab