diff --git a/cranelift/codegen/src/egraph/cost.rs b/cranelift/codegen/src/egraph/cost.rs index ee4a40ade376..8deef914222e 100644 --- a/cranelift/codegen/src/egraph/cost.rs +++ b/cranelift/codegen/src/egraph/cost.rs @@ -1,6 +1,87 @@ //! Cost functions for egraph representation. -use crate::ir::Opcode; +use crate::ir::{Inst, Opcode}; +use cranelift_entity::EntityRef; + +/// Approximate cost of an expression as a DAG of instructions. +/// +/// In addition to the saturating total cost, this tracks an approximate +/// footprint of the instructions that contribute to the expression. When an +/// operand's whole footprint is already covered, we don't charge its total +/// again. This catches common shared-DAG shapes like `iadd x, x` without +/// allocating precise instruction sets in the egraph extraction hot path. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct ExprCost { + total: Cost, + inst_buckets: u32, +} + +impl ExprCost { + pub(crate) fn zero() -> Self { + Self { + total: Cost::zero(), + inst_buckets: 0, + } + } + + pub(crate) fn infinity() -> Self { + Self { + total: Cost::infinity(), + inst_buckets: 0, + } + } + + pub(crate) fn for_inst(inst: Inst, op: Opcode) -> Self { + Self { + total: Cost::of_opcode(op), + inst_buckets: Self::inst_bucket(inst), + } + } + + /// Compute the cost of the operation and its given operands. + /// + /// Caller is responsible for checking that the opcode came from an instruction + /// that satisfies `inst_predicates::is_pure_for_egraph()`. + pub(crate) fn of_pure_op( + inst: Inst, + op: Opcode, + operand_costs: impl IntoIterator, + ) -> Self { + let mut cost = Self::for_inst(inst, op); + for operand_cost in operand_costs { + cost.add_operand(operand_cost); + } + cost + } +} + +impl ExprCost { + fn inst_bucket(inst: Inst) -> u32 { + let index = u32::try_from(inst.index()).unwrap(); + let hash = index.wrapping_mul(0x9e37_79b9); + 1u32 << (hash >> 27) + } + + fn add_operand(&mut self, other: Self) { + let new_buckets = other.inst_buckets & !self.inst_buckets; + if new_buckets != 0 { + self.total = self.total + other.total; + } + self.inst_buckets |= other.inst_buckets; + } +} + +impl PartialOrd for ExprCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ExprCost { + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + self.total.cmp(&other.total) + } +} /// A cost of computing some value in the program. /// @@ -175,4 +256,22 @@ mod tests { assert_eq!(a + b, Cost::infinity()); assert_eq!(b + a, Cost::infinity()); } + + #[test] + fn expr_cost_skips_fully_covered_operand() { + let x = ExprCost::for_inst(Inst::new(0), Opcode::Iconst); + let add = ExprCost::of_pure_op(Inst::new(1), Opcode::Iadd, [x, x]); + + assert_eq!(add.total, Cost::new(4)); + } + + #[test] + fn expr_cost_grows_linearly_for_repeated_self_adds() { + let mut cost = ExprCost::for_inst(Inst::new(0), Opcode::Iconst); + for index in 1..4 { + cost = ExprCost::of_pure_op(Inst::new(index), Opcode::Iadd, [cost, cost]); + } + + assert_eq!(cost.total, Cost::new(10)); + } } diff --git a/cranelift/codegen/src/egraph/elaborate.rs b/cranelift/codegen/src/egraph/elaborate.rs index 813320663440..f73ccecac38e 100644 --- a/cranelift/codegen/src/egraph/elaborate.rs +++ b/cranelift/codegen/src/egraph/elaborate.rs @@ -2,7 +2,7 @@ //! in CFG nodes. use super::Stats; -use super::cost::Cost; +use super::cost::{Cost, ExprCost}; use crate::ctxhash::NullCtx; use crate::dominator_tree::DominatorTree; use crate::hash_map::Entry as HashEntry; @@ -121,6 +121,28 @@ impl Ord for BestEntry { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ExprBestEntry(ExprCost, Value); + +impl PartialOrd for ExprBestEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ExprBestEntry { + #[inline] + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + self.0.cmp(&other.0).then_with(|| { + // Note that this comparison is reversed. When costs are equal, + // prefer the value with the bigger index. This is a heuristic that + // prefers results of rewrites to the original value, since we + // expect that our rewrites are generally improvements. + self.1.cmp(&other.1).reverse() + }) + } +} + #[derive(Clone, Copy, Debug)] struct ElaboratedValue { in_block: Block, @@ -345,8 +367,6 @@ impl<'a> Elaborator<'a> { fn compute_best_values(&mut self) { let sorted_values = self.topo_sorted_values(); - let best = &mut self.value_to_best_value; - // We can't make random decisions inside the fixpoint loop below because // that could cause values to change on every iteration of the loop, // which would make the loop never terminate. So in chaos testing @@ -364,6 +384,21 @@ impl<'a> Elaborator<'a> { } ); + let saw_infinity = self.compute_best_values_with_scalar_costs(&sorted_values, use_worst); + + if saw_infinity { + self.compute_best_values_with_expr_costs(&sorted_values, use_worst); + } + } + + fn compute_best_values_with_scalar_costs( + &mut self, + sorted_values: &[Value], + use_worst: bool, + ) -> bool { + let best = &mut self.value_to_best_value; + let mut saw_infinity = false; + // Because the values are topologically sorted, we know that we will see // defs before uses, so an instruction's operands' costs will already be // computed by the time we are computing the cost for the current value @@ -414,6 +449,9 @@ impl<'a> Elaborator<'a> { best[value].0 }), ); + if cost == Cost::infinity() { + saw_infinity = true; + } best[value] = BestEntry(cost, value); trace!(" -> cost of value {} = {:?}", value, cost); } @@ -440,13 +478,67 @@ impl<'a> Elaborator<'a> { // v6 = iadd v5, v5 ;; cost = 3 + 125 + 125 // ;; etc... // - // Such a chain can cause cost to saturate to infinity. How do we - // choose which e-node is best when there are multiple that have - // saturated to infinity? It doesn't matter. As long as invariant - // (2) for optimization rules is upheld by our rule set (see - // `cranelift/codegen/src/opts/README.md`) it is safe to choose - // *any* e-node in the e-class. At worst we will produce suboptimal - // code, but never an incorrectness. + // If this happens, the caller will recompute the same values with + // a heavier sharing-aware fallback cost. Otherwise, the common path + // remains the simple scalar cost. + } + + saw_infinity + } + + fn compute_best_values_with_expr_costs(&mut self, sorted_values: &[Value], use_worst: bool) { + let mut expr_best = SecondaryMap::with_default(ExprBestEntry( + ExprCost::infinity(), + Value::reserved_value(), + )); + expr_best.resize(self.func.dfg.num_values()); + + for value in sorted_values.iter().copied() { + let def = self.func.dfg.value_def(value); + trace!( + "recomputing sharing-aware best for value {:?} def {:?}", + value, def + ); + + match def { + ValueDef::Union(x, y) => { + debug_assert!(!expr_best[x].1.is_reserved_value()); + debug_assert!(!expr_best[y].1.is_reserved_value()); + expr_best[value] = if use_worst { + core::cmp::max(expr_best[x], expr_best[y]) + } else { + core::cmp::min(expr_best[x], expr_best[y]) + }; + trace!( + " -> sharing-aware best of union({:?}, {:?}) = {:?}", + expr_best[x], expr_best[y], expr_best[value] + ); + } + + ValueDef::Param(_, _) => { + expr_best[value] = ExprBestEntry(ExprCost::zero(), value); + } + + ValueDef::Result(inst, _) => { + if let Some(_) = self.func.layout.inst_block(inst) { + expr_best[value] = ExprBestEntry(ExprCost::zero(), value); + } else { + let inst_data = &self.func.dfg.insts[inst]; + let cost = ExprCost::of_pure_op( + inst, + inst_data.opcode(), + self.func.dfg.inst_values(inst).map(|value| { + debug_assert!(!expr_best[value].1.is_reserved_value()); + expr_best[value].0 + }), + ); + expr_best[value] = ExprBestEntry(cost, value); + trace!(" -> sharing-aware cost of value {} = {:?}", value, cost); + } + } + } + + self.value_to_best_value[value].1 = expr_best[value].1; } } diff --git a/cranelift/filetests/filetests/egraph/cost-function.clif b/cranelift/filetests/filetests/egraph/cost-function.clif new file mode 100644 index 000000000000..10d992b576fa --- /dev/null +++ b/cranelift/filetests/filetests/egraph/cost-function.clif @@ -0,0 +1,87 @@ +;; A chain of `iadd x, x` saturates a scalar cost that recounts shared +;; operands. Extraction must still see through the identity +;; `(x * 2) - x` once that chain is costed as a DAG. + +test optimize precise-output +set opt_level=speed_and_size +target x86_64 + +function %f(i64) -> i64 { + block0(v0: i64): + v1 = iadd v0, v0 + v2 = iadd v1, v1 + v3 = iadd v2, v2 + v4 = iadd v3, v3 + v5 = iadd v4, v4 + v6 = iadd v5, v5 + v7 = iadd v6, v6 + v8 = iadd v7, v7 + v9 = iadd v8, v8 + v10 = iadd v9, v9 + v11 = iadd v10, v10 + v12 = iadd v11, v11 + v13 = iadd v12, v12 + v14 = iadd v13, v13 + v15 = iadd v14, v14 + v16 = iadd v15, v15 + v17 = iadd v16, v16 + v18 = iadd v17, v17 + v19 = iadd v18, v18 + v20 = iadd v19, v19 + v21 = iadd v20, v20 + v22 = iadd v21, v21 + v23 = iadd v22, v22 + v24 = iadd v23, v23 + v25 = iadd v24, v24 + v26 = iadd v25, v25 + v27 = iadd v26, v26 + v28 = iadd v27, v27 + v29 = iadd v28, v28 + v30 = iadd v29, v29 + v31 = iadd v30, v30 + v32 = iadd v31, v31 + v33 = iadd v32, v32 + + v34 = iconst.i64 2 + v35 = imul v33, v34 + v36 = isub v35, v33 + return v36 +} + +; function %f(i64) -> i64 fast { +; block0(v0: i64): +; v1 = iadd v0, v0 +; v2 = iadd v1, v1 +; v3 = iadd v2, v2 +; v4 = iadd v3, v3 +; v5 = iadd v4, v4 +; v6 = iadd v5, v5 +; v7 = iadd v6, v6 +; v8 = iadd v7, v7 +; v9 = iadd v8, v8 +; v10 = iadd v9, v9 +; v11 = iadd v10, v10 +; v12 = iadd v11, v11 +; v13 = iadd v12, v12 +; v14 = iadd v13, v13 +; v15 = iadd v14, v14 +; v16 = iadd v15, v15 +; v17 = iadd v16, v16 +; v18 = iadd v17, v17 +; v19 = iadd v18, v18 +; v20 = iadd v19, v19 +; v21 = iadd v20, v20 +; v22 = iadd v21, v21 +; v23 = iadd v22, v22 +; v24 = iadd v23, v23 +; v25 = iadd v24, v24 +; v26 = iadd v25, v25 +; v27 = iadd v26, v26 +; v28 = iadd v27, v27 +; v29 = iadd v28, v28 +; v30 = iadd v29, v29 +; v31 = iadd v30, v30 +; v32 = iadd v31, v31 +; v33 = iadd v32, v32 +; return v33 +; } diff --git a/cranelift/filetests/filetests/isa/x64/iminmax-i128.clif b/cranelift/filetests/filetests/isa/x64/iminmax-i128.clif index 851b38bc02a7..46df59e7e140 100644 --- a/cranelift/filetests/filetests/isa/x64/iminmax-i128.clif +++ b/cranelift/filetests/filetests/isa/x64/iminmax-i128.clif @@ -19,19 +19,18 @@ block0: ; pushq %rbp ; movq %rsp, %rbp ; block0: -; uninit %rax -; xorq %rax, %rax ; uninit %rdx ; xorq %rdx, %rdx -; movq %rax, %rdi -; subq $0x0, %rdi -; movq %rdx, %r9 -; sbbq $0x0, %r9 -; cmpq %rax, %rdi -; movq %r9, %r8 +; movq %rdx, %rsi +; subq $0x0, %rsi +; movq %rdx, %rdi +; sbbq $0x0, %rdi +; cmpq %rdx, %rsi +; movq %rdi, %r8 ; sbbq %rdx, %r8 -; cmovbq %rdi, %rax -; cmovbq %r9, %rdx +; movq %rdx, %rax +; cmovbq %rsi, %rax +; cmovbq %rdi, %rdx ; movq %rbp, %rsp ; popq %rbp ; retq @@ -41,17 +40,17 @@ block0: ; pushq %rbp ; movq %rsp, %rbp ; block1: ; offset 0x4 -; xorq %rax, %rax ; xorq %rdx, %rdx -; movq %rax, %rdi -; subq $0, %rdi -; movq %rdx, %r9 -; sbbq $0, %r9 -; cmpq %rax, %rdi -; movq %r9, %r8 +; movq %rdx, %rsi +; subq $0, %rsi +; movq %rdx, %rdi +; sbbq $0, %rdi +; cmpq %rdx, %rsi +; movq %rdi, %r8 ; sbbq %rdx, %r8 -; cmovbq %rdi, %rax -; cmovbq %r9, %rdx +; movq %rdx, %rax +; cmovbq %rsi, %rax +; cmovbq %rdi, %rdx ; movq %rbp, %rsp ; popq %rbp ; retq diff --git a/tests/disas/gc/array-new-default-i16.wat b/tests/disas/gc/array-new-default-i16.wat index f2caa8e7f063..9cdfe652ae6d 100644 --- a/tests/disas/gc/array-new-default-i16.wat +++ b/tests/disas/gc/array-new-default-i16.wat @@ -36,10 +36,9 @@ ;; ;; block0(v0: i64, v1: i64, v2: i32): ;; @001f v4 = uextend.i64 v2 -;; v91 = iconst.i64 1 -;; v92 = ishl v4, v91 ; v91 = 1 +;; v90 = iadd v4, v4 ;; @001f v7 = iconst.i64 32 -;; @001f v8 = ushr v92, v7 ; v7 = 32 +;; @001f v8 = ushr v90, v7 ; v7 = 32 ;; @001f trapnz v8, user18 ;; @001f v3 = iconst.i32 20 ;; v96 = iadd v2, v2 @@ -108,12 +107,12 @@ ;; @001f v76 = load.i64 notrap aligned region12 v145+40 ;; @001f v64 = iconst.i64 20 ;; @001f v65 = iadd v49, v64 ; v64 = 20 -;; @001f v78 = uadd_overflow_trap v65, v92, user2 +;; @001f v78 = uadd_overflow_trap v65, v90, user2 ;; @001f v77 = iadd v146, v76 ;; @001f v79 = icmp ugt v78, v77 ;; @001f trapnz v79, user2 ;; @001f v44 = iconst.i32 0 -;; @001f call fn1(v0, v65, v44, v92), stack_map=[i32 @ ss0+0] ; v44 = 0 +;; @001f call fn1(v0, v65, v44, v90), stack_map=[i32 @ ss0+0] ; v44 = 0 ;; @0022 jump block1 ;; ;; block1: diff --git a/tests/disas/x64-optimize-vector-types.wat b/tests/disas/x64-optimize-vector-types.wat index 28316f6823fc..d99db4ff2bda 100644 --- a/tests/disas/x64-optimize-vector-types.wat +++ b/tests/disas/x64-optimize-vector-types.wat @@ -201,10 +201,12 @@ ;; wasm[0]::function[8]: ;; pushq %rbp ;; movq %rsp, %rbp -;; por %xmm0, %xmm1 ;; pcmpeqd %xmm7, %xmm7 +;; movdqa %xmm1, %xmm2 +;; pxor %xmm7, %xmm2 +;; pandn %xmm0, %xmm1 ;; movdqa %xmm1, %xmm0 -;; pcmpeqb %xmm7, %xmm0 +;; pcmpeqb %xmm2, %xmm0 ;; movq %rbp, %rsp ;; popq %rbp ;; retq