From c7aa94687b733c112aca66ca22cc09c035879d8a Mon Sep 17 00:00:00 2001
From: Russel Arbore <russel.jma@gmail.com>
Date: Wed, 19 Feb 2025 21:35:50 -0600
Subject: [PATCH] Fix inlining id functions
---
hercules_opt/src/inline.rs | 6 +++++-
juno_samples/fork_join_tests/src/gpu.sch | 1 +
2 files changed, 6 insertions(+), 1 deletion(-)
diff --git a/hercules_opt/src/inline.rs b/hercules_opt/src/inline.rs
index f01b2366..6e308274 100644
--- a/hercules_opt/src/inline.rs
+++ b/hercules_opt/src/inline.rs
@@ -210,7 +210,11 @@ fn inline_func(
}
// Finally, delete the call node.
- edit = edit.replace_all_uses(id, old_id_to_new_id(called_return_data))?;
+ if let Node::Parameter { index } = called_func.nodes[called_return_data.idx()] {
+ edit = edit.replace_all_uses(id, args[index])?;
+ } else {
+ edit = edit.replace_all_uses(id, old_id_to_new_id(called_return_data))?;
+ }
edit = edit.delete_node(control)?;
edit = edit.delete_node(id)?;
diff --git a/juno_samples/fork_join_tests/src/gpu.sch b/juno_samples/fork_join_tests/src/gpu.sch
index 91bd6c79..364673cd 100644
--- a/juno_samples/fork_join_tests/src/gpu.sch
+++ b/juno_samples/fork_join_tests/src/gpu.sch
@@ -52,6 +52,7 @@ slf(auto.test2);
infer-schedules(auto.test2);
fork-interchange[0, 1](auto.test2);
+inline(test6);
fork-tile[32, 0, false, true](test6@loop);
let out = fork-split(test6@loop);
let out = auto-outline(test6);
--
GitLab