From 3ba235b888b84448d7c10305aa1f4fa6ce2e732d Mon Sep 17 00:00:00 2001
From: rarbore2 <rarbore2@illinois.edu>
Date: Mon, 18 Nov 2024 22:26:03 -0600
Subject: [PATCH] Fix inlining and editor handling of gravestone nodes
---
hercules_opt/src/editor.rs | 12 +++++++++++-
hercules_opt/src/inline.rs | 9 ++++++++-
2 files changed, 19 insertions(+), 2 deletions(-)
diff --git a/hercules_opt/src/editor.rs b/hercules_opt/src/editor.rs
index b830af98..95fd1669 100644
--- a/hercules_opt/src/editor.rs
+++ b/hercules_opt/src/editor.rs
@@ -634,7 +634,17 @@ pub fn repair_plan(plan: &mut Plan, new_function: &Function, edits: &[Edit]) {
new_partitions[data_id.idx()] = new_partitions[bbs[data_id.idx()].idx()];
}
- // Step 6: wrap everything up.
+ // Step 6: create a solitary gravestone partition. This will get removed
+ // when gravestone nodes are removed.
+ let gravestone_partition = PartitionID::new(plan.num_partitions);
+ plan.num_partitions += 1;
+ for (idx, node) in new_function.nodes.iter().enumerate() {
+ if idx > 0 && node.is_start() {
+ new_partitions[idx] = Some(gravestone_partition);
+ }
+ }
+
+ // Step 7: wrap everything up.
plan.partitions = new_partitions.into_iter().map(|id| id.unwrap()).collect();
plan.partition_devices
.resize(plan.num_partitions, Device::CPU);
diff --git a/hercules_opt/src/inline.rs b/hercules_opt/src/inline.rs
index ecc8027f..8e43c1c4 100644
--- a/hercules_opt/src/inline.rs
+++ b/hercules_opt/src/inline.rs
@@ -197,7 +197,14 @@ fn inline_func(
// Stich the control use of the original call node's region with
// the predecessor control of the inlined function's return.
- edit = edit.replace_all_uses(control, old_id_to_new_id(called_return_pred))?;
+ edit = edit.replace_all_uses(
+ control,
+ if called_return_pred == NodeID::new(0) {
+ call_pred
+ } else {
+ old_id_to_new_id(called_return_pred)
+ },
+ )?;
// Stitch uses of parameter nodes in the inlined function to the IDs
// of arguments provided to the call node.
--
GitLab