Repository: bytecodealliance/wasmtime
Pull request #7280 (2023-10-18 04:23:30)

PCC: merge/propagate facts through egraph opts.
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.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/egraph.rs b/cranelift/codegen/src/egraph.rs
index e12b6fed3..001687e65 100644
--- a/cranelift/codegen/src/egraph.rs
+++ b/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);
         }
 
diff --git a/cranelift/codegen/src/ir/dfg.rs b/cranelift/codegen/src/ir/dfg.rs
index c903b18f4..411aa73bd 100644
--- a/cranelift/codegen/src/ir/dfg.rs
+++ b/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.
diff --git a/cranelift/codegen/src/ir/pcc.rs b/cranelift/codegen/src/ir/pcc.rs
index 406228000..b2efb1dc8 100644
--- a/cranelift/codegen/src/ir/pcc.rs
+++ b/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"),
         }
     }
 }
diff --git a/cranelift/codegen/src/isa/riscv64/inst.isle b/cranelift/codegen/src/isa/riscv64/inst.isle
index 6da824e62..921e144ae 100644
--- a/cranelift/codegen/src/isa/riscv64/inst.isle
+++ b/cranelift/codegen/src/isa/riscv64/inst.isle
@@ -1558,6 +1558,12 @@
 (rule (rv_orn rs1 rs2)
   (alu_rrr (AluOPRRR.Orn) rs1 rs2))
 
+;; Helper for emitting the `xnor` ("Exclusive NOR") instruction.
+;; rd ← ~(rs1 ^ rs2)
+(decl rv_xnor (XReg XReg) XReg)
+(rule (rv_xnor rs1 rs2)
+  (alu_rrr (AluOPRRR.Xnor) rs1 rs2))
+
 ;; Helper for emitting the `clz` ("Count Leading Zero Bits") instruction.
 (decl rv_clz (XReg) XReg)
 (rule (rv_clz rs1)
@@ -2008,23 +2014,6 @@
 (rule (select_addi (fits_in_64 ty)) (AluOPRRI.Addi))
 
 
-(decl gen_bnot (Type ValueRegs) ValueRegs)
-(rule 2 (gen_bnot (ty_scalar_float ty) x)
-  (let ((val FReg (value_regs_get x 0))
-        (x_val XReg (move_f_to_x val ty))
-        (inverted XReg (rv_not x_val))
-        (res FReg (move_x_to_f inverted (float_int_of_same_size ty))))
-    (value_reg res)))
-
-(rule 1 (gen_bnot $I128 x)
-  (let ((lo XReg (rv_not (value_regs_get x 0)))
-        (hi XReg (rv_not (value_regs_get x 1))))
-    (value_regs lo hi)))
-
-(rule 0 (gen_bnot (ty_int_ref_scalar_64 _) x)
-  (rv_not (value_regs_get x 0)))
-
-
 (decl gen_andi (XReg u64) XReg)
 (rule 1 (gen_andi x (imm12_from_u64 y))
   (rv_andi x y))
diff --git a/cranelift/codegen/src/isa/riscv64/lower.isle b/cranelift/codegen/src/isa/riscv64/lower.isle
index 021a3503d..c5bebc4c5 100644
--- a/cranelift/codegen/src/isa/riscv64/lower.isle
+++ b/cranelift/codegen/src/isa/riscv64/lower.isle
@@ -934,12 +934,25 @@
 (extern constructor binvi_imm binvi_imm)
 
 ;;;; Rules for `bnot` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
-(rule 0 (lower (has_type (ty_scalar ty) (bnot x)))
-  (gen_bnot ty x))
 
-(rule 1 (lower (has_type (ty_vec_fits_in_register ty) (bnot x)))
+(rule 0 (lower (has_type (ty_int_ref_scalar_64 _) (bnot x)))
+  (rv_not x))
+
+(rule 1 (lower (has_type (ty_scalar_float ty) (bnot x)))
+  (move_x_to_f (rv_not (move_f_to_x x ty)) (float_int_of_same_size ty)))
+
+(rule 2 (lower (has_type $I128 (bnot x)))
+  (value_regs
+    (rv_not (value_regs_get x 0))
+    (rv_not (value_regs_get x 1))))
+
+(rule 3 (lower (has_type (ty_vec_fits_in_register ty) (bnot x)))
   (rv_vnot_v x (unmasked) ty))
 
+(rule 4 (lower (has_type (ty_int_ref_scalar_64 _) (bnot (bxor x y))))
+  (if-let $true (has_zbb))
+  (rv_xnor x y))
+
 ;;;; Rules for `bit_reverse` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 (rule (lower (has_type (fits_in_64 (ty_int ty)) (bitrev x)))
   (lower_bit_reverse x ty))
diff --git a/cranelift/reader/src/parser.rs b/cranelift/reader/src/parser.rs
index b088af004..dbd84d857 100644
--- a/cranelift/reader/src/parser.rs
+++ b/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")),
         }
     }
 


Test patch:
diff --git a/cranelift/filetests/filetests/isa/riscv64/wasm/zbb.wat b/cranelift/filetests/filetests/isa/riscv64/wasm/zbb.wat
index ed51298d7..e307e6a6a 100644
--- a/cranelift/filetests/filetests/isa/riscv64/wasm/zbb.wat
+++ b/cranelift/filetests/filetests/isa/riscv64/wasm/zbb.wat
@@ -20,8 +20,16 @@
     (i32.rotr (local.get 0) (i32.const 100)))
   (func (export "rori") (param i64) (result i64)
     (i64.rotr (local.get 0) (i64.const 40)))
-)
 
+  (func (export "xnor32_1") (param i32 i32) (result i32)
+    (i32.xor (i32.xor (local.get 0) (local.get 1)) (i32.const -1)))
+  (func (export "xnor32_2") (param i32 i32) (result i32)
+    (i32.xor (i32.const -1) (i32.xor (local.get 0) (local.get 1))))
+  (func (export "xnor64_1") (param i64 i64) (result i64)
+    (i64.xor (i64.xor (local.get 0) (local.get 1)) (i64.const -1)))
+  (func (export "xnor64_2") (param i64 i64) (result i64)
+    (i64.xor (i64.const -1) (i64.xor (local.get 0) (local.get 1))))
+)
 ;; function u0:0:
 ;; block0:
 ;;   j label1
@@ -77,3 +85,31 @@
 ;; block1:
 ;;   rori a0,a0,40
 ;;   ret
+;;
+;; function u0:8:
+;; block0:
+;;   j label1
+;; block1:
+;;   xnor a0,a0,a1
+;;   ret
+;;
+;; function u0:9:
+;; block0:
+;;   j label1
+;; block1:
+;;   xnor a0,a0,a1
+;;   ret
+;;
+;; function u0:10:
+;; block0:
+;;   j label1
+;; block1:
+;;   xnor a0,a0,a1
+;;   ret
+;;
+;; function u0:11:
+;; block0:
+;;   j label1
+;; block1:
+;;   xnor a0,a0,a1
+;;   ret
diff --git a/cranelift/filetests/filetests/pcc/fail/opt.clif b/cranelift/filetests/filetests/pcc/fail/opt.clif
new file mode 100644
index 000000000..fa628583a
--- /dev/null
+++ b/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
+}
diff --git a/cranelift/filetests/filetests/pcc/succeed/opt.clif b/cranelift/filetests/filetests/pcc/succeed/opt.clif
new file mode 100644
index 000000000..cd1abebd9
--- /dev/null
+++ b/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
+}