egraph-based midend: draw the rest of the owl (productionized). (#4953)
* egraph-based midend: draw the rest of the owl. * Rename `egg` submodule of cranelift-codegen to `egraph`. * Apply some feedback from @jsharp during code walkthrough. * Remove recursion from find_best_node by doing a single pass. Rather than recursively computing the lowest-cost node for a given eclass and memoizing the answer at each eclass node, we can do a single forward pass; because every eclass node refers only to earlier nodes, this is sufficient. The behavior may slightly differ from the earlier behavior because we cannot short-circuit costs to zero once a node is elaborated; but in practice this should not matter. * Make elaboration non-recursive. Use an explicit stack instead (with `ElabStackEntry` entries, alongside a result stack). * Make elaboration traversal of the domtree non-recursive/stack-safe. * Work analysis logic in Cranelift-side egraph glue into a general analysis framework in cranelift-egraph. * Apply static recursion limit to rule application. * Fix aarch64 wrt dynamic-vector support -- broken rebase. * Topo-sort cranelift-egraph before cranelift-codegen in publish script, like the comment instructs me to! * Fix multi-result call testcase. * Include `cranelift-egraph` in `PUBLISHED_CRATES`. * Fix atomic_rmw: not really a load. * Remove now-unnecessary PartialOrd/Ord derivations. * Address some code-review comments. * Review feedback. * Review feedback. * No overlap in mid-end rules, because we are defining a multi-constructor. * rustfmt * Review feedback. * Review feedback. * Review feedback. * Review feedback. * Remove redundant `mut`. * Add comment noting what rules can do. * Review feedback. * Clarify comment wording. * Update `has_memory_fence_semantics`. * Apply @jameysharp's improved loop-level computation. Co-authored-by: Jamey Sharp <jamey@minilop.net> * Fix suggestion commit. * Fix off-by-one in new loop-nest analysis. * Review feedback. * Review feedback. * Review feedback. * Use `Default`, not `std::default::Default`, as per @fitzgen Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com> * Apply @fitzgen's comment elaboration to a doc-comment. Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com> * Add stat for hitting the rewrite-depth limit. * Some code motion in split prelude to make the diff a little clearer wrt `main`. * Take @jameysharp's suggested `try_into()` usage for blockparam indices. Co-authored-by: Jamey Sharp <jamey@minilop.net> * Take @jameysharp's suggestion to avoid double-match on load op. Co-authored-by: Jamey Sharp <jamey@minilop.net> * Fix suggestion (add import). * Review feedback. * Fix stack_load handling. * Remove redundant can_store case. * Take @jameysharp's suggested improvement to FuncEGraph::build() logic Co-authored-by: Jamey Sharp <jamey@minilop.net> * Tweaks to FuncEGraph::build() on top of suggestion. * Take @jameysharp's suggested clarified condition Co-authored-by: Jamey Sharp <jamey@minilop.net> * Clean up after suggestion (unused variable). * Fix loop analysis. * loop level asserts * Revert constant-space loop analysis -- edge cases were incorrect, so let's go with the simple thing for now. * Take @jameysharp's suggestion re: result_tys Co-authored-by: Jamey Sharp <jamey@minilop.net> * Fix up after suggestion * Take @jameysharp's suggestion to use fold rather than reduce Co-authored-by: Jamey Sharp <jamey@minilop.net> * Fixup after suggestion * Take @jameysharp's suggestion to remove elaborate_eclass_use's return value. * Clarifying comment in terminator insts. Co-authored-by: Jamey Sharp <jamey@minilop.net> Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>
This commit is contained in:
Chris Fallin
297
cranelift/codegen/src/opts.rs
Normal file
297
cranelift/codegen/src/opts.rs
Normal file
@@ -0,0 +1,297 @@
|
||||
//! Optimization driver using ISLE rewrite rules on an egraph.
|
||||
|
||||
use crate::egraph::Analysis;
|
||||
use crate::egraph::FuncEGraph;
|
||||
use crate::egraph::MemoryState;
|
||||
pub use crate::egraph::{Node, NodeCtx};
|
||||
use crate::ir::condcodes;
|
||||
pub use crate::ir::condcodes::{FloatCC, IntCC};
|
||||
pub use crate::ir::immediates::{Ieee32, Ieee64, Imm64, Offset32, Uimm32, Uimm64, Uimm8};
|
||||
pub use crate::ir::types::*;
|
||||
pub use crate::ir::{
|
||||
dynamic_to_fixed, AtomicRmwOp, Block, Constant, DynamicStackSlot, FuncRef, GlobalValue, Heap,
|
||||
Immediate, InstructionImms, JumpTable, MemFlags, Opcode, StackSlot, Table, TrapCode, Type,
|
||||
Value,
|
||||
};
|
||||
use crate::isle_common_prelude_methods;
|
||||
use crate::machinst::isle::*;
|
||||
use crate::trace;
|
||||
pub use cranelift_egraph::{Id, NewOrExisting, NodeIter};
|
||||
use cranelift_entity::{EntityList, EntityRef};
|
||||
use smallvec::SmallVec;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
pub type IdArray = EntityList<Id>;
|
||||
#[allow(dead_code)]
|
||||
pub type Unit = ();
|
||||
pub type Range = (usize, usize);
|
||||
|
||||
pub type ConstructorVec<T> = SmallVec<[T; 8]>;
|
||||
|
||||
mod generated_code;
|
||||
use generated_code::ContextIter;
|
||||
|
||||
struct IsleContext<'a, 'b> {
|
||||
egraph: &'a mut FuncEGraph<'b>,
|
||||
}
|
||||
|
||||
const REWRITE_LIMIT: usize = 5;
|
||||
|
||||
pub fn optimize_eclass<'a>(id: Id, egraph: &mut FuncEGraph<'a>) -> Id {
|
||||
trace!("running rules on eclass {}", id.index());
|
||||
egraph.stats.rewrite_rule_invoked += 1;
|
||||
|
||||
if egraph.rewrite_depth > REWRITE_LIMIT {
|
||||
egraph.stats.rewrite_depth_limit += 1;
|
||||
return id;
|
||||
}
|
||||
egraph.rewrite_depth += 1;
|
||||
|
||||
// Find all possible rewrites and union them in, returning the
|
||||
// union.
|
||||
let mut ctx = IsleContext { egraph };
|
||||
let optimized_ids = generated_code::constructor_simplify(&mut ctx, id);
|
||||
let mut union_id = id;
|
||||
if let Some(mut ids) = optimized_ids {
|
||||
while let Some(new_id) = ids.next(&mut ctx) {
|
||||
if ctx.egraph.subsume_ids.contains(&new_id) {
|
||||
trace!(" -> eclass {} subsumes {}", new_id, id);
|
||||
ctx.egraph.stats.node_subsume += 1;
|
||||
// Merge in the unionfind so canonicalization still
|
||||
// works, but take *only* the subsuming ID, and break
|
||||
// now.
|
||||
ctx.egraph.egraph.unionfind.union(union_id, new_id);
|
||||
union_id = new_id;
|
||||
break;
|
||||
}
|
||||
ctx.egraph.stats.node_union += 1;
|
||||
let old_union_id = union_id;
|
||||
union_id = ctx
|
||||
.egraph
|
||||
.egraph
|
||||
.union(&ctx.egraph.node_ctx, union_id, new_id);
|
||||
trace!(
|
||||
" -> union eclass {} with {} to get {}",
|
||||
new_id,
|
||||
old_union_id,
|
||||
union_id
|
||||
);
|
||||
}
|
||||
}
|
||||
trace!(" -> optimize {} got {}", id, union_id);
|
||||
ctx.egraph.rewrite_depth -= 1;
|
||||
union_id
|
||||
}
|
||||
|
||||
pub(crate) fn store_to_load<'a>(id: Id, egraph: &mut FuncEGraph<'a>) -> Id {
|
||||
// Note that we only examine the latest enode in the eclass: opts
|
||||
// are invoked for every new enode added to an eclass, so
|
||||
// traversing the whole eclass would be redundant.
|
||||
let load_key = egraph.egraph.classes[id].get_node().unwrap();
|
||||
if let Node::Load {
|
||||
op:
|
||||
InstructionImms::Load {
|
||||
opcode: Opcode::Load,
|
||||
offset: load_offset,
|
||||
..
|
||||
},
|
||||
ty: load_ty,
|
||||
addr: load_addr,
|
||||
mem_state: MemoryState::Store(store_inst),
|
||||
..
|
||||
} = load_key.node(&egraph.egraph.nodes)
|
||||
{
|
||||
trace!(" -> got load op for id {}", id);
|
||||
if let Some((store_ty, store_id)) = egraph.store_nodes.get(&store_inst) {
|
||||
trace!(" -> got store id: {} ty: {}", store_id, store_ty);
|
||||
let store_key = egraph.egraph.classes[*store_id].get_node().unwrap();
|
||||
if let Node::Inst {
|
||||
op:
|
||||
InstructionImms::Store {
|
||||
opcode: Opcode::Store,
|
||||
offset: store_offset,
|
||||
..
|
||||
},
|
||||
args: store_args,
|
||||
..
|
||||
} = store_key.node(&egraph.egraph.nodes)
|
||||
{
|
||||
let store_args = store_args.as_slice(&egraph.node_ctx.args);
|
||||
let store_data = store_args[0];
|
||||
let store_addr = store_args[1];
|
||||
if *load_offset == *store_offset
|
||||
&& *load_ty == *store_ty
|
||||
&& egraph.egraph.unionfind.equiv_id_mut(*load_addr, store_addr)
|
||||
{
|
||||
trace!(" -> same offset, type, address; forwarding");
|
||||
egraph.stats.store_to_load_forward += 1;
|
||||
return store_data;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
id
|
||||
}
|
||||
|
||||
struct NodesEtorIter<'a, 'b>
|
||||
where
|
||||
'b: 'a,
|
||||
{
|
||||
root: Id,
|
||||
iter: NodeIter<NodeCtx, Analysis>,
|
||||
_phantom1: PhantomData<&'a ()>,
|
||||
_phantom2: PhantomData<&'b ()>,
|
||||
}
|
||||
|
||||
impl<'a, 'b> generated_code::ContextIter for NodesEtorIter<'a, 'b>
|
||||
where
|
||||
'b: 'a,
|
||||
{
|
||||
type Context = IsleContext<'a, 'b>;
|
||||
type Output = (Type, InstructionImms, IdArray);
|
||||
|
||||
fn next(&mut self, ctx: &mut IsleContext<'a, 'b>) -> Option<Self::Output> {
|
||||
while let Some(node) = self.iter.next(&ctx.egraph.egraph) {
|
||||
trace!("iter from root {}: node {:?}", self.root, node);
|
||||
match node {
|
||||
Node::Pure { op, args, types }
|
||||
| Node::Inst {
|
||||
op, args, types, ..
|
||||
} if types.len() == 1 => {
|
||||
let ty = types.as_slice(&ctx.egraph.node_ctx.types)[0];
|
||||
return Some((ty, op.clone(), args.clone()));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, 'b> generated_code::Context for IsleContext<'a, 'b> {
|
||||
isle_common_prelude_methods!();
|
||||
|
||||
fn eclass_type(&mut self, eclass: Id) -> Option<Type> {
|
||||
let mut iter = self.egraph.egraph.enodes(eclass);
|
||||
while let Some(node) = iter.next(&self.egraph.egraph) {
|
||||
match node {
|
||||
&Node::Pure { types, .. } | &Node::Inst { types, .. } if types.len() == 1 => {
|
||||
return Some(types.as_slice(&self.egraph.node_ctx.types)[0]);
|
||||
}
|
||||
&Node::Load { ty, .. } => return Some(ty),
|
||||
&Node::Result { ty, .. } => return Some(ty),
|
||||
&Node::Param { ty, .. } => return Some(ty),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn at_loop_level(&mut self, eclass: Id) -> (u8, Id) {
|
||||
(
|
||||
self.egraph.egraph.analysis_value(eclass).loop_level.level() as u8,
|
||||
eclass,
|
||||
)
|
||||
}
|
||||
|
||||
type enodes_etor_iter = NodesEtorIter<'a, 'b>;
|
||||
|
||||
fn enodes_etor(&mut self, eclass: Id) -> Option<NodesEtorIter<'a, 'b>> {
|
||||
Some(NodesEtorIter {
|
||||
root: eclass,
|
||||
iter: self.egraph.egraph.enodes(eclass),
|
||||
_phantom1: PhantomData,
|
||||
_phantom2: PhantomData,
|
||||
})
|
||||
}
|
||||
|
||||
fn pure_enode_ctor(&mut self, ty: Type, op: &InstructionImms, args: IdArray) -> Id {
|
||||
let types = self.egraph.node_ctx.types.single(ty);
|
||||
let types = types.freeze(&mut self.egraph.node_ctx.types);
|
||||
let op = op.clone();
|
||||
match self
|
||||
.egraph
|
||||
.egraph
|
||||
.add(Node::Pure { op, args, types }, &mut self.egraph.node_ctx)
|
||||
{
|
||||
NewOrExisting::New(id) => {
|
||||
self.egraph.stats.node_created += 1;
|
||||
self.egraph.stats.node_pure += 1;
|
||||
self.egraph.stats.node_ctor_created += 1;
|
||||
optimize_eclass(id, self.egraph)
|
||||
}
|
||||
NewOrExisting::Existing(id) => {
|
||||
self.egraph.stats.node_ctor_deduped += 1;
|
||||
id
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn id_array_0_etor(&mut self, arg0: IdArray) -> Option<()> {
|
||||
let values = arg0.as_slice(&self.egraph.node_ctx.args);
|
||||
if values.len() == 0 {
|
||||
Some(())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn id_array_0_ctor(&mut self) -> IdArray {
|
||||
EntityList::default()
|
||||
}
|
||||
|
||||
fn id_array_1_etor(&mut self, arg0: IdArray) -> Option<Id> {
|
||||
let values = arg0.as_slice(&self.egraph.node_ctx.args);
|
||||
if values.len() == 1 {
|
||||
Some(values[0])
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn id_array_1_ctor(&mut self, arg0: Id) -> IdArray {
|
||||
EntityList::from_iter([arg0].into_iter(), &mut self.egraph.node_ctx.args)
|
||||
}
|
||||
|
||||
fn id_array_2_etor(&mut self, arg0: IdArray) -> Option<(Id, Id)> {
|
||||
let values = arg0.as_slice(&self.egraph.node_ctx.args);
|
||||
if values.len() == 2 {
|
||||
Some((values[0], values[1]))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn id_array_2_ctor(&mut self, arg0: Id, arg1: Id) -> IdArray {
|
||||
EntityList::from_iter([arg0, arg1].into_iter(), &mut self.egraph.node_ctx.args)
|
||||
}
|
||||
|
||||
fn id_array_3_etor(&mut self, arg0: IdArray) -> Option<(Id, Id, Id)> {
|
||||
let values = arg0.as_slice(&self.egraph.node_ctx.args);
|
||||
if values.len() == 3 {
|
||||
Some((values[0], values[1], values[2]))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn id_array_3_ctor(&mut self, arg0: Id, arg1: Id, arg2: Id) -> IdArray {
|
||||
EntityList::from_iter(
|
||||
[arg0, arg1, arg2].into_iter(),
|
||||
&mut self.egraph.node_ctx.args,
|
||||
)
|
||||
}
|
||||
|
||||
fn remat(&mut self, id: Id) -> Id {
|
||||
trace!("remat: {}", id);
|
||||
self.egraph.remat_ids.insert(id);
|
||||
id
|
||||
}
|
||||
|
||||
fn subsume(&mut self, id: Id) -> Id {
|
||||
trace!("subsume: {}", id);
|
||||
self.egraph.subsume_ids.insert(id);
|
||||
id
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user