diff --git a/Cargo.lock b/Cargo.lock
index c438e846d87f12fd7eea2c27d73a8fc7e043836d..32dc6a0ee754c243e4d1b7f45892190cf54c1d7c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1347,6 +1347,16 @@ dependencies = [
"with_builtin_macros",
]
+[[package]]
+name = "juno_multi_return"
+version = "0.1.0"
+dependencies = [
+ "async-std",
+ "hercules_rt",
+ "juno_build",
+ "with_builtin_macros",
+]
+
[[package]]
name = "juno_patterns"
version = "0.1.0"
diff --git a/Cargo.toml b/Cargo.toml
index 42d2813535ca9c00facef65212269555d0173f27..01f8cc138b7a2dfc2f3a6cc2882558961fd6b63f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ members = [
"juno_samples/matmul",
"juno_samples/median_window",
"juno_samples/multi_device",
+ "juno_samples/multi_return",
"juno_samples/patterns",
"juno_samples/products",
"juno_samples/rodinia/backprop",
diff --git a/juno_samples/multi_return/Cargo.toml b/juno_samples/multi_return/Cargo.toml
new file mode 100644
index 0000000000000000000000000000000000000000..0fb3de94cfff22a8720d084ad2bd7e43c87a4cd4
--- /dev/null
+++ b/juno_samples/multi_return/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "juno_multi_return"
+version = "0.1.0"
+authors = ["Aaron Councilman <aaronjc4@illinois.edu>"]
+edition = "2021"
+
+[[bin]]
+name = "juno_multi_return"
+path = "src/main.rs"
+
+[features]
+cuda = ["juno_build/cuda", "hercules_rt/cuda"]
+
+[build-dependencies]
+juno_build = { path = "../../juno_build" }
+
+[dependencies]
+juno_build = { path = "../../juno_build" }
+hercules_rt = { path = "../../hercules_rt" }
+with_builtin_macros = "0.1.0"
+async-std = "*"
diff --git a/juno_samples/multi_return/build.rs b/juno_samples/multi_return/build.rs
new file mode 100644
index 0000000000000000000000000000000000000000..3a8f9b1c503f8cc5fbca0baf06ff8edbd300bf9e
--- /dev/null
+++ b/juno_samples/multi_return/build.rs
@@ -0,0 +1,15 @@
+use juno_build::JunoCompiler;
+
+fn main() {
+ JunoCompiler::new()
+ .file_in_src("multi_return.jn")
+ .unwrap()
+ .schedule_in_src(if cfg!(feature = "cuda") {
+ "gpu.sch"
+ } else {
+ "cpu.sch"
+ })
+ .unwrap()
+ .build()
+ .unwrap();
+}
diff --git a/juno_samples/multi_return/src/cpu.sch b/juno_samples/multi_return/src/cpu.sch
new file mode 100644
index 0000000000000000000000000000000000000000..03fb2585cff07dff23d26bbd66296611e5555f46
--- /dev/null
+++ b/juno_samples/multi_return/src/cpu.sch
@@ -0,0 +1,31 @@
+gvn(*);
+phi-elim(*);
+dce(*);
+
+ip-sroa(*);
+sroa(*);
+dce(*);
+
+forkify(*);
+fork-guard-elim(*);
+gvn(*);
+dce(*);
+
+inline(*);
+delete-uncalled(*);
+
+let out = auto-outline(*);
+cpu(out.rolling_sum_prod);
+
+fork-fusion(out.rolling_sum_prod);
+gvn(*);
+dce(*);
+
+float-collections(*);
+
+unforkify(*);
+gvn(*);
+ccp(*);
+dce(*);
+
+gcm(*);
diff --git a/juno_samples/multi_return/src/gpu.sch b/juno_samples/multi_return/src/gpu.sch
new file mode 100644
index 0000000000000000000000000000000000000000..e733551d31c4bb6cfc76668a1b03e5be89851ec5
--- /dev/null
+++ b/juno_samples/multi_return/src/gpu.sch
@@ -0,0 +1,26 @@
+gvn(*);
+phi-elim(*);
+dce(*);
+
+ip-sroa(*);
+sroa(*);
+dce(*);
+
+forkify(*);
+fork-guard-elim(*);
+gvn(*);
+dce(*);
+
+inline(*);
+delete-uncalled(*);
+
+let out = auto-outline(*);
+gpu(out.rolling_sum_prod);
+
+fork-fusion(out.rolling_sum_prod);
+gvn(*);
+dce(*);
+
+float-collections(*);
+
+gcm(*);
diff --git a/juno_samples/multi_return/src/main.rs b/juno_samples/multi_return/src/main.rs
new file mode 100644
index 0000000000000000000000000000000000000000..63479dbab1b3f4aaede401566ea00f3917cfa3b9
--- /dev/null
+++ b/juno_samples/multi_return/src/main.rs
@@ -0,0 +1,22 @@
+#![feature(concat_idents)]
+
+juno_build::juno!("median");
+
+use hercules_rt::{runner, HerculesImmBox, HerculesImmBoxTo};
+
+fn main() {
+ let m = vec![
+ 86, 72, 14, 5, 55, 25, 98, 89, 3, 66, 44, 81, 27, 3, 40, 18, 4, 57, 93, 34, 70, 50, 50, 18,
+ 34,
+ ];
+ let m = HerculesImmBox::from(m.as_slice());
+
+ let mut r = runner!(median_window);
+ let res = async_std::task::block_on(async { r.run(m.to()).await });
+ assert_eq!(res, 57);
+}
+
+#[test]
+fn test_median_window() {
+ main()
+}
diff --git a/juno_samples/multi_return/src/multi_return.jn b/juno_samples/multi_return/src/multi_return.jn
new file mode 100644
index 0000000000000000000000000000000000000000..a49df91c2179d76242965f088cda93c261ebec28
--- /dev/null
+++ b/juno_samples/multi_return/src/multi_return.jn
@@ -0,0 +1,32 @@
+fn rolling_sum<t: number, n: usize>(x: t[n]) -> t, t[n + 1] {
+ let rolling_sum: t[n + 1];
+ let sum = 0;
+
+ for i in 0..n {
+ rolling_sum[i] = sum;
+ sum += x[i];
+ }
+ rolling_sum[n] = sum;
+
+ return (sum, rolling_sum);
+}
+
+fn rolling_prod<t: number, n: usize>(x: t[n]) -> t, t[n + 1] {
+ let rolling_prod: t[n + 1];
+ let prod = 1;
+
+ for i in 0..n {
+ rolling_prod[i] = prod;
+ prod *= x[i];
+ }
+ rolling_prod[n] = prod;
+
+ return prod, rolling_prod;
+}
+
+#[entry]
+fn rolling_sum_prod<n: usize>(x: f32[n]) -> f32[n + 1], f32[n + 1] {
+ let rsum = rolling_sum::<_, n>(x).1;
+ let _, rprod = rolling_prod::<_, n>(x);
+ return rsum, rprod;
+}
diff --git a/juno_scheduler/src/pm.rs b/juno_scheduler/src/pm.rs
index d5c0af27e280d523683cec2b39da387a3f2c3f45..94f900486f64a6021cc9e4ee098fd2dd2e8ecd27 100644
--- a/juno_scheduler/src/pm.rs
+++ b/juno_scheduler/src/pm.rs
@@ -2077,16 +2077,32 @@ fn run_pass(
pm.clear_analyses();
}
Pass::InterproceduralSROA => {
- assert!(args.is_empty());
- if let Some(_) = selection {
- return Err(SchedulerError::PassError {
- pass: "interproceduralSROA".to_string(),
- error: "must be applied to the entire module".to_string(),
- });
- }
+ let sroa_with_arrays = match args.get(0) {
+ Some(Value::Boolean { val }) => *val,
+ Some(_) => {
+ return Err(SchedulerError::PassError {
+ pass: "sroa".to_string(),
+ error: "expected boolean argument".to_string(),
+ });
+ }
+ None => false,
+ };
+
+ let selection = selection_of_functions(pm, selection)
+ .ok_or_else(|| {
+ SchedulerError::PassError {
+ pass: "xdot".to_string(),
+ error: "expected coarse-grained selection (can't partially xdot a function)".to_string(),
+ }
+ })?;
+ let mut bool_selection = vec![false; pm.functions.len()];
+ selection.into_iter().for_each(|func| bool_selection[func.idx()] = true);
+
+ pm.make_typing();
+ let typing = pm.typing.take().unwrap();
let mut editors = build_editors(pm);
- interprocedural_sroa(&mut editors);
+ interprocedural_sroa(&mut editors, &typing, &bool_selection, sroa_with_arrays);
for func in editors {
changed |= func.modified();
@@ -2720,21 +2736,15 @@ fn run_pass(
None => true,
};
- let mut bool_selection = vec![];
- if let Some(selection) = selection {
- bool_selection = vec![false; pm.functions.len()];
- for loc in selection {
- let CodeLocation::Function(id) = loc else {
- return Err(SchedulerError::PassError {
+ let selection = selection_of_functions(pm, selection)
+ .ok_or_else(|| {
+ SchedulerError::PassError {
pass: "xdot".to_string(),
error: "expected coarse-grained selection (can't partially xdot a function)".to_string(),
- });
- };
- bool_selection[id.idx()] = true;
- }
- } else {
- bool_selection = vec![true; pm.functions.len()];
- }
+ }
+ })?;
+ let mut bool_selection = vec![false; pm.functions.len()];
+ selection.into_iter().for_each(|func| bool_selection[func.idx()] = true);
pm.make_reverse_postorders();
if force_analyses {