Skip to content
Snippets Groups Projects
Commit e95929fd authored by Aaron Councilman's avatar Aaron Councilman
Browse files

Merge branch 'matmul_paper_schedules' into fork-reshape

parents 39e2ea93 10f93736
No related branches found
No related tags found
1 merge request!205Fork reshape
Pipeline #201930 passed
......@@ -676,50 +676,47 @@ pub fn fork_coalesce_helper(
// CHECKME / FIXME: Might need to be added the other way.
new_factors.append(&mut inner_dims.to_vec());
for tid in inner_tids {
let (fork, dim) = editor.func().nodes[tid.idx()].try_thread_id().unwrap();
let new_tid = Node::ThreadID {
control: fork,
dimension: dim + num_outer_dims,
};
let mut new_fork = NodeID::new(0);
let new_join = inner_join; // We'll reuse the inner join as the join of the new fork
let success = editor.edit(|mut edit| {
for tid in inner_tids {
let (fork, dim) = edit.get_node(tid).try_thread_id().unwrap();
let new_tid = Node::ThreadID {
control: fork,
dimension: dim + num_outer_dims,
};
editor.edit(|mut edit| {
let new_tid = edit.add_node(new_tid);
let mut edit = edit.replace_all_uses(tid, new_tid)?;
edit = edit.replace_all_uses(tid, new_tid)?;
edit.sub_edit(tid, new_tid);
Ok(edit)
});
}
// Fuse Reductions
for (outer_reduce, inner_reduce) in pairs {
let (_, outer_init, _) = editor.func().nodes[outer_reduce.idx()]
.try_reduce()
.unwrap();
let (_, inner_init, _) = editor.func().nodes[inner_reduce.idx()]
.try_reduce()
.unwrap();
let success = editor.edit(|mut edit| {
}
// Fuse Reductions
for (outer_reduce, inner_reduce) in pairs {
let (_, outer_init, _) = edit.get_node(outer_reduce)
.try_reduce()
.unwrap();
let (_, inner_init, _) = edit.get_node(inner_reduce)
.try_reduce()
.unwrap();
// Set inner init to outer init.
edit =
edit.replace_all_uses_where(inner_init, outer_init, |usee| *usee == inner_reduce )?;
edit = edit.replace_all_uses(outer_reduce, inner_reduce)?;
edit = edit.delete_node(outer_reduce)?;
}
Ok(edit)
});
}
let mut new_fork = NodeID::new(0);
let new_join = inner_join; // We reuse the inner join as the join of the new fork
editor.edit(|mut edit| {
let new_fork_node = Node::Fork {
control: outer_pred,
factors: new_factors.into(),
};
new_fork = edit.add_node(new_fork_node);
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)?;
......@@ -730,7 +727,11 @@ pub fn fork_coalesce_helper(
Ok(edit)
});
Some((new_fork, new_join))
if success {
Some((new_fork, new_join))
} else {
None
}
}
pub fn split_any_fork(
......@@ -1277,9 +1278,14 @@ pub fn fork_interchange(
edit = edit.delete_node(old_id)?;
}
let new_fork = edit.add_node(new_fork);
new_fork_id = Some(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)
edit = edit.delete_node(fork)?;
new_fork_id = Some(new_fork);
Ok(edit)
});
new_fork_id
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment