If-conversion

src/lib/egraph/extractor/extractor_core.ml

@@ -91,6 +91,13 @@ let rec pp_ext ppf = function
     Format.fprintf ppf "(%a %a %a)" pp_prov p Enode.pp_op op
       (Format.pp_print_list ~pp_sep pp_ext) args
 
+(* Pure cost of an operation.
+
+   It's important to note that some of these costs are relative
+   to what we can rewrite the operation into. For example,
+   integer division/remainder has an enormous cost because we
+   want to incentivize rewrites a la Hacker's Delight.
+*)
 let op_cost : Enode.op -> cost = function
   | Oint (i, t) ->
     (* In practice, a negative constant might need some work to
@@ -116,7 +123,7 @@ let op_cost : Enode.op -> cost = function
   | Ovastart _ -> Cost.pure 0
   | Obr | Ovar _ -> Cost.pure 2
   | Osw _ | Obinop #Insn.bitwise_binop | Ounop _ -> Cost.pure 3
-  | Obinop (`div _ | `udiv _ | `rem _ | `urem _) -> Cost.pure 90
+  | Obinop (`div _ | `udiv _ | `rem _ | `urem _) -> Cost.pure 97
   | Obinop (`mul _) -> Cost.pure 42
   | Obinop (`mulh _ | `umulh _) -> Cost.pure 11
   | Obinop _ -> Cost.pure 4

src/lib/machine/x86/x86_amd64_common.ml

@@ -631,7 +631,7 @@ module Insn = struct
 
          This is morally equivalent to:
 
-         `x := cc ? b : a`
+         `a := cc ? b : a`
       *)
       ->
       Set.union (rset' [`rflags])

src/lib/magic_div.ml

@@ -67,10 +67,11 @@ let signed d t =
         then B.(q2 + one, r2 - ad)
         else q2, r2 in
       let delta = B.(ad - r2) in
-      if Bv.(q1 < delta || (q1 = delta && r1 = zero)) then
+      if Bv.(q1 < delta || (q1 = delta && r1 = zero))
+      then loop p q1 r1 q2 r2
+      else
         let m = B.(q2 + one) in
-        (if B.msb d then B.neg m else m), p - sz
-      else loop p q1 r1 q2 r2 in
+        (if B.msb d then B.neg m else m), p - sz in
     let q1 = B.(mins / anc) in
     let r1 = B.(mins - q1 * anc) in
     let q2 = B.(mins / ad) in

src/lib/passes/simplify_cfg/simplify_cfg.ml

@@ -2,6 +2,7 @@ open Core
 open Virtual
 open Simplify_cfg_common
 
+module Ifc = Simplify_cfg_ifc
 module Brsel = Simplify_cfg_brsel
 module Merge_blks = Simplify_cfg_merge_blks
 module Contract = Simplify_cfg_contract
@@ -26,6 +27,7 @@ let run tenv fn =
       || Contract.run env
       || Short_circ.run env
       || Duplicate_br.run env
+      || Ifc.run tenv env fn
       || brsel
       then loop @@ recompute_cfg env @@ update_fn env fn
       else !!fn in

src/lib/passes/simplify_cfg/simplify_cfg_common.ml

@@ -42,14 +42,9 @@ let is_ret env l = match env.ret with
   | None -> false
 
 let update_fn env fn =
-  Func.blks fn |>
-  Seq.fold ~init:[] ~f:(fun acc b ->
-      let l = Blk.label b in
-      if Hashtbl.mem env.blks l then acc
-      else l :: acc) |>
-  Func.remove_blks_exn fn |>
-  Func.map_blks ~f:(fun b ->
-      Hashtbl.find_exn env.blks @@ Blk.label b)
+  Func.blks fn |> Seq.filter_map ~f:(fun b ->
+      Hashtbl.find env.blks @@ Blk.label b) |>
+  Seq.to_list |> Func.with_blks_exn fn
 
 let not_pseudo = Fn.non Label.is_pseudo
 

src/lib/passes/simplify_cfg/simplify_cfg_ifc.ml

@@ -0,0 +1,182 @@
+(* If-conversion canonicalization.
+
+   For the `Brsel` pass, we handle the "degenerate diamond"
+   case for if-conversion, where the join block is targeted
+   by both arms of the branch for our header block.
+
+   For other cases (e.g. a proper diamond, or a triangle),
+   we want to canonicalize the shape enough for other passes
+   to fire, so that `Brsel` can do the final part of if-conversion.
+*)
+
+open Core
+open Regular.Std
+open Virtual
+open Simplify_cfg_common
+
+module BR = Simplify_cfg_brsel
+
+type join =
+  | Diamond of {
+      pred1 : Label.t;
+      pred2 : Label.t;
+      hdr   : Label.t;
+    }
+  | Triangle of {
+      pred : Label.t;
+      hdr  : Label.t;
+    }
+
+let diamond p1 p2 h = Diamond {pred1 = p1; pred2 = p2; hdr = h}
+let triangle p h = Triangle {pred = p; hdr = h}
+
+let pp_join ppf = function
+  | Diamond {pred1; pred2; hdr} ->
+    Format.fprintf ppf "(diamond (hdr %a) (pred1 %a) (pred2 %a))"
+      Label.pp hdr Label.pp pred1 Label.pp pred2
+  | Triangle {pred; hdr} ->
+    Format.fprintf ppf "(triangle (hdr %a) (pred %a))"
+      Label.pp hdr Label.pp pred
+
+exception Not_hoistable
+
+let insn_cost : Insn.op -> int = function
+  (* Provably non-trapping div/rem gets a high cost. *)
+  | `bop (_, (`div #Type.imm | `rem #Type.imm), _, `int (i, _))
+    when Bv.(i <> zero) -> 4
+  (* Effectful or potentially trapping instructions cannot
+     be hoisted. *)
+  | `bop (_, (`div #Type.imm | `rem #Type.imm), _, _)
+  | `load _
+  | `store _
+  | #Insn.variadic
+  | `call _
+    -> raise_notrace Not_hoistable
+  (* Cheap operations. *)
+  | `uop (_, #Insn.copy, _)
+  | `bop (_, #Insn.arith_binop, _, _)
+  | `bop (_, #Insn.bitwise_binop, _, _)
+  | `bop (_, #Insn.cmp, _, _)
+  | `uop (_, #Insn.arith_unop, _)
+  | `uop (_, #Insn.bitwise_unop, _)
+    -> 1
+  (* Moderate operations. *)
+  | `uop (_, #Insn.cast, _)
+  | `sel _
+    -> 2
+
+let blk_cost b =
+  Blk.insns b |> Seq.map ~f:Insn.op |>
+  Seq.sum (module Int) ~f:insn_cost
+
+let max_join_args = 2
+let max_pred_cost = 4
+
+let check_args tenv env fn b =
+  let args = Seq.to_list @@ Blk.args b in
+  List.length args <= max_join_args && try
+    List.iter args ~f:(fun x ->
+        ignore @@ BR.basicty tenv env fn x);
+    true
+  with BR.Non_basic -> false
+
+let (.?[]) env l = Hashtbl.find env.blks l
+let (.![]) env l = Hashtbl.find_exn env.blks l
+
+let check_blk env l j = match env.?[l] with
+  | None -> false
+  | Some b ->
+    try match Blk.ctrl b with
+      | `jmp `label (l', args) when Label.(l' = j) ->
+        (* Ensure that this block has no parameters.
+           We have an invariant that, in its canonical
+           form, the only time a block can have parameters
+           is if it exists as a join point. *)
+        Seq.is_empty (Blk.args b) &&
+        (* Ensure correct arity *)
+        Seq.length (Blk.args env.![j]) = List.length args &&
+        (* Ensure that the block isn't too expensive to copy
+           to the header. *)
+        blk_cost b <= max_pred_cost
+      | _ -> false
+    with Not_hoistable -> false
+
+let check_hdr env l s1 s2 = match env.?[l] with
+  | None -> false
+  | Some b -> match Blk.ctrl b with
+    | `br (_, `label (l1, _), `label (l2, _)) ->
+      Label.((l1 = s1 && l2 = s2) || (l1 = s2 && l2 = s1))
+    | _ -> false
+
+let plist env l = Seq.to_list @@ Cfg.Node.preds l env.cfg
+
+let find_join tenv env fn =
+  with_return @@ fun {return} ->
+  Hashtbl.iteri env.blks ~f:(fun ~key ~data:b ->
+      let found j = return @@ Some (key, j) in
+      if check_args tenv env fn b then match plist env key with
+        | [p1; p2] when Label.(p1 <> p2) ->
+          begin match plist env p1, plist env p2 with
+            | [p1'], _ when Label.(p1' = p2) ->
+              if check_blk env p1 key
+              && check_hdr env p2 p1 key
+              then found @@ triangle p1 p2
+            | _, [p2'] when Label.(p2' = p1) ->
+              if check_blk env p2 key
+              && check_hdr env p1 p2 key
+              then found @@ triangle p2 p1
+            | [p1'], [p2'] when Label.(p1' = p2') ->
+              if check_blk env p1 key
+              && check_blk env p2 key
+              && check_hdr env p1' p1 p2
+              then found @@ diamond p1 p2 p1'
+            | _ -> ()
+          end
+        | _ -> ());
+  None
+
+let canonicalize_diamond env p1 p2 hdr =
+  let b1 = env.![p1] and b2 = env.![p2] and h = env.![hdr] in
+  let ctrl = match Blk.(ctrl h, ctrl b1, ctrl b2) with
+    | `br (c, `label (l1, _), `label (l2, _)),
+      `jmp (`label _ as d1),
+      `jmp (`label _ as d2) ->
+      if Label.(l1 = p1 && l2 = p2) then
+        `br (c, d1, d2)
+      else if Label.(l1 = p2 && l2 = p1) then
+        `br (c, d2, d1)
+      else assert false
+    | _ -> assert false in
+  let i1 = Seq.append (Blk.insns b1) (Blk.insns b2) in
+  let h' = Blk.append_insns h @@ Seq.to_list i1 in
+  let h' = Blk.with_ctrl h' ctrl in
+  Hashtbl.set env.blks ~key:hdr ~data:h'
+
+let canonicalize_triangle env key pred hdr =
+  let b = env.![pred] and h = env.![hdr] in
+  let ctrl = match Blk.(ctrl h, ctrl b) with
+    | `br (c, `label (l1, a1), `label (l2, a2)),
+      `jmp (`label _ as d) ->
+      if Label.(key = l1 && pred = l2) then
+        `br (c, `label (l1, a1), d)
+      else if Label.(key = l2 && pred = l1) then
+        `br (c, d, `label (l2, a2))
+      else assert false
+    | _ -> assert false in
+  let h' = Blk.append_insns h @@ Seq.to_list @@ Blk.insns b in
+  let h' = Blk.with_ctrl h' ctrl in
+  Hashtbl.set env.blks ~key:hdr ~data:h'
+
+let canonicalize env key = function
+  | Diamond {pred1; pred2; hdr} -> canonicalize_diamond env pred1 pred2 hdr
+  | Triangle {pred; hdr} -> canonicalize_triangle env key pred hdr
+
+let run tenv env fn = match find_join tenv env fn with
+  | None -> false
+  | Some (key, data) ->
+    canonicalize env key data;
+    Logs.debug (fun m ->
+        m "%s: join in $%s at %a: %a%!"
+          __FUNCTION__ (Func.name fn)
+          Label.pp key pp_join data);
+    true

src/test/data/opt/and_test.vir.opt

@@ -2,45 +2,40 @@ module and_test
 
 export function w $foo(w %x, w %y) {
 @2:
-  %0 = and.w %x, 0x3_w ; @29
-  %1 = eq.w %0, 0x0_w ; @30
-  %11 = sel.w %1, %y, %x ; @40
-  ret %11
+  %2 = and.w %x, 0x3_w ; @31
+  %3 = eq.w %2, 0x0_w ; @32
+  %4 = sel.w %3, %y, %x ; @33
+  ret %4
 }
 
 export function w $bar(w %x, w %y) {
 @10:
-  %2 = and.w %x, 0x3_w ; @31
-  %3 = eq.w %2, 0x0_w ; @32
-  br %3, @11, @12
-@11:
-  %4 = add.w %y, 0x1_w ; @33
-  jmp @15(%4)
-@12:
-  %5 = add.w %x, 0x1_w ; @34
-  jmp @15(%5)
-@15(%r.1):
-  ret %r.1
+  %5 = and.w %x, 0x3_w ; @34
+  %6 = eq.w %5, 0x0_w ; @35
+  %7 = add.w %y, 0x1_w ; @36
+  %8 = add.w %x, 0x1_w ; @37
+  %9 = sel.w %6, %7, %8 ; @38
+  ret %9
 }
 
 export function l $baz(l %x) {
 @18:
-  %6 = and.l %x, 0x3_l ; @35
-  switch.l %6, @19 [0x1_l -> @20,
-                    0x2_l -> @21,
-                    0x3_l -> @22]
+  %10 = and.l %x, 0x3_l ; @39
+  switch.l %10, @19 [0x1_l -> @20,
+                     0x2_l -> @21,
+                     0x3_l -> @22]
 @19:
-  %7 = add.l %x, 0x1_l ; @36
-  jmp @24(%7)
+  %11 = add.l %x, 0x1_l ; @40
+  jmp @24(%11)
 @20:
-  %8 = add.l %x, 0x2_l ; @37
-  jmp @24(%8)
+  %12 = add.l %x, 0x2_l ; @41
+  jmp @24(%12)
 @21:
-  %9 = add.l %x, 0x3_l ; @38
-  jmp @24(%9)
+  %13 = add.l %x, 0x3_l ; @42
+  jmp @24(%13)
 @22:
-  %10 = add.l %x, 0x4_l ; @39
-  jmp @24(%10)
+  %14 = add.l %x, 0x4_l ; @43
+  jmp @24(%14)
 @24(%r.1):
   ret %r.1
 }