PCC: merge/propagate facts through egraph opts. (#7280)

This turns out to be quite simple: the fundamental operation during
egraph-based optimization is to *merge* eclasses, which is an assertion
that their value is equal. Since the values of either side of the merge
are equal, a fact about one side is a fact about the other, and
vice-versa.

Since we only support *one* fact at most per value, we can't take the
union of all facts; instead, if one side is missing a fact, or if both
sides have exactly the same fact, we keep that one; otherwise we go to a
special "conflict" fact that can't support any check. This is edging
closer to a fact-lattice, but I opted not to go there with a full
meet-function merge yet, for simplicity. (It's a little complex because
of the "minimum fact" we can know about a value based on its type -- if
we're going to do something better, I think we should account for that
too.)

In any case, that complexity mostly isn't needed, because the two cases
that happen in reality are (i) opt rules rewrite to a new node, which
will have no facts at all, so facts just propagate; or (ii) GVN merges
two values in the input program into one, but if both are the same
value, in practice the Wasm PCC frontend (for example) should be
producing the same facts on both values, so the merge is trivial.

cranelift/codegen/src/egraph.rs

@@ -162,6 +162,7 @@ where
                 let result = self.func.dfg.first_result(inst);
                 self.value_to_opt_value[result] = orig_result;
                 self.eclasses.union(result, orig_result);
+                self.func.dfg.merge_facts(result, orig_result);
                 self.stats.union += 1;
                 result
             } else {
@@ -256,6 +257,11 @@ where
                 // still works, but take *only* the subsuming
                 // value, and break now.
                 isle_ctx.ctx.eclasses.union(optimized_value, union_value);
+                isle_ctx
+                    .ctx
+                    .func
+                    .dfg
+                    .merge_facts(optimized_value, union_value);
                 union_value = optimized_value;
                 break;
             }
@@ -273,6 +279,11 @@ where
                 .ctx
                 .eclasses
                 .union(old_union_value, optimized_value);
+            isle_ctx
+                .ctx
+                .func
+                .dfg
+                .merge_facts(old_union_value, optimized_value);
             isle_ctx.ctx.eclasses.union(old_union_value, union_value);
         }
 

cranelift/codegen/src/ir/dfg.rs

@@ -1285,6 +1285,28 @@ impl DataFlowGraph {
     pub fn detach_block_params(&mut self, block: Block) -> ValueList {
         self.blocks[block].params.take()
     }
+
+    /// Merge the facts for two values. If both values have facts and
+    /// they differ, both values get a special "conflict" fact that is
+    /// never satisfied.
+    pub fn merge_facts(&mut self, a: Value, b: Value) {
+        let a = self.resolve_aliases(a);
+        let b = self.resolve_aliases(b);
+        match (&self.facts[a], &self.facts[b]) {
+            (Some(a), Some(b)) if a == b => { /* nothing */ }
+            (None, None) => { /* nothing */ }
+            (Some(a), None) => {
+                self.facts[b] = Some(a.clone());
+            }
+            (None, Some(b)) => {
+                self.facts[a] = Some(b.clone());
+            }
+            _ => {
+                self.facts[a] = Some(Fact::Conflict);
+                self.facts[b] = Some(Fact::Conflict);
+            }
+        }
+    }
 }
 
 /// Contents of a basic block.

cranelift/codegen/src/ir/pcc.rs

@@ -164,6 +164,11 @@ pub enum Fact {
         /// The maximum offset into the memory type, inclusive.
         max_offset: u64,
     },
+
+    /// A "conflict fact": this fact results from merging two other
+    /// facts, and it can never be satisfied -- checking any value
+    /// against this fact will fail.
+    Conflict,
 }
 
 impl fmt::Display for Fact {
@@ -179,6 +184,7 @@ impl fmt::Display for Fact {
                 min_offset,
                 max_offset,
             } => write!(f, "mem({}, {:#x}, {:#x})", ty, min_offset, max_offset),
+            Fact::Conflict => write!(f, "conflict"),
         }
     }
 }

cranelift/filetests/filetests/pcc/fail/opt.clif

@@ -0,0 +1,30 @@
+test compile expect-fail
+set enable_pcc=true
+set opt_level=speed
+target aarch64
+
+;; Failed merge of two facts.
+function %f0(i64, i32) {
+    mt0 = struct 4 { 0: i32 ! range(32, 1, 3) }
+
+block0(v0 ! mem(mt0, 0, 0): i64, v1 ! range(32, 0, 1): i32):
+    v2 ! range(32, 1, 1) = iconst.i32 1
+    v3 ! range(32, 1, 2) = iadd.i32 v1, v2
+
+    v4 ! range(32, 1, 1) = iconst.i32 1
+    v5 ! range(32, 1, 3) = iadd.i32 v1, v4  ;; should GVN onto v3.
+
+    ;; v3/v5's facts should merge to `conflict`. Even though we *could*
+    ;; take the union of possible ranges, we don't; so even though the
+    ;; stored range *would* subsume the field's fact if we could
+    ;; compute it more precisely, here we should expect a failure. We
+    ;; can revisit this if we decide to admit lattice-meet merging
+    ;; later.
+    ;;
+    ;; (Note that we have to use both v3 and v5 here so one doesn't
+    ;; get DCE'd away before opt merges them.)
+    store.i32 checked v3, v0
+    store.i32 checked v5, v0
+
+    return
+}

cranelift/filetests/filetests/pcc/succeed/opt.clif

@@ -0,0 +1,47 @@
+test compile
+set enable_pcc=true
+set opt_level=speed
+target aarch64
+
+;; Equivalent to a Wasm `i64.load` from a static memory, but with some
+;; redundant stuff that should be optimized away (x+0 -> x).
+function %f0(i64, i32) -> i64 {
+    ;; mock vmctx struct:
+    mt0 = struct 8 { 0: i64 readonly ! mem(mt1, 0, 0) }
+    ;; mock static memory: 4GiB range, plus 2GiB guard
+    mt1 = memory 0x1_8000_0000
+
+block0(v0 ! mem(mt0, 0, 0): i64, v1: i32):
+    v2 ! mem(mt1, 0, 0) = load.i64 checked v0+0
+    v3 ! range(64, 0, 0xffff_ffff) = uextend.i64 v1
+    v4 = iconst.i64 0
+    v5 = iadd.i64 v3, v4
+    v6 ! mem(mt1, 0, 0xffff_ffff) = iadd.i64 v2, v5
+    v7 = load.i64 checked v6
+    return v7
+}
+
+;; GVN opportunity.
+function %f0(i64, i32) -> i64 {
+    ;; mock vmctx struct:
+    mt0 = struct 8 { 0: i64 readonly ! mem(mt1, 0, 0) }
+    ;; mock static memory: 4GiB range, plus 2GiB guard
+    mt1 = memory 0x1_8000_0000
+
+block0(v0 ! mem(mt0, 0, 0): i64, v1: i32):
+    v2 ! mem(mt1, 0, 0) = load.i64 checked notrap readonly v0+0
+    v3 ! range(64, 0, 0xffff_ffff) = uextend.i64 v1
+    v4 = iconst.i64 0
+    v5 = iadd.i64 v3, v4
+    v6 ! mem(mt1, 0, 0xffff_ffff) = iadd.i64 v2, v5
+    v7 = load.i64 checked v6
+
+    v8 = load.i64 checked notrap readonly v0+0
+    v9 = uextend.i64 v1
+    v10 ! mem(mt1, 0, 0xffff_ffff) = iadd.i64 v8, v9
+    v11 = load.i64 checked v10
+
+    v12 = iadd.i64 v7, v11
+
+    return v12
+}

cranelift/reader/src/parser.rs

@@ -2167,6 +2167,7 @@ impl<'a> Parser<'a> {
     //
     // fact ::= "range" "(" bit-width "," min-value "," max-value ")"
     //        | "mem" "(" memory-type "," mt-offset "," mt-offset ")"
+    //        | "conflict"
     // bit-width ::= uimm64
     // min-value ::= uimm64
     // max-value ::= uimm64
@@ -2237,7 +2238,11 @@ impl<'a> Parser<'a> {
                     max_offset,
                 })
             }
-            _ => Err(self.error("expected a `range` or `mem` fact")),
+            Some(Token::Identifier("conflict")) => {
+                self.consume();
+                Ok(Fact::Conflict)
+            }
+            _ => Err(self.error("expected a `range`, `mem` or `conflict` fact")),
         }
     }