cranelift: fold ctz/clz directly into brif cond via simplify_skeleton

cranelift/codegen/src/egraph/mod.rs

@@ -811,6 +811,14 @@ where
                     );
                     (inst, Some(val))
                 }
+                // `ReplaceBranchCond` is unconditionally accepted — the
+                // opcode and successors don't change, so we can't use the
+                // cost-based ranking the other variants do. The first such
+                // candidate wins; ISLE rule ordering picks the form.
+                SkeletonInstSimplification::ReplaceBranchCond { cond } => {
+                    log::trace!(" -> simplify_skeleton: replace `brif` cond with {cond}");
+                    return Some(SkeletonInstSimplification::ReplaceBranchCond { cond });
+                }
             };
 
             if cfg!(debug_assertions) {
@@ -1119,6 +1127,16 @@ impl<'a> EgraphPass<'a> {
             }
             SkeletonInstSimplification::Replace { inst } => (inst, None),
             SkeletonInstSimplification::ReplaceWithVal { inst, val } => (inst, Some(val)),
+            SkeletonInstSimplification::ReplaceBranchCond { cond } => {
+                // Swap the condition operand of the existing `brif` in
+                // place. Successors stay; CFG is preserved.
+                debug_assert_eq!(
+                    cursor.func.dfg.insts[old_inst].opcode(),
+                    crate::ir::Opcode::Brif,
+                );
+                cursor.func.dfg.inst_args_mut(old_inst)[0] = cond;
+                return;
+            }
         };
 
         // Replace the old instruction with the new one.

cranelift/codegen/src/opts/icmp.isle

@@ -450,3 +450,18 @@
 ;; clz(X) != 0  iff  the MSB of X is 0, i.e.  X is signed-non-negative.
 (rule (simplify (ne result_ty (clz x_ty X) (iconst_u _ 0)))
       (sge result_ty X (iconst_u x_ty 0)))
+
+;;;;; Same simplifications applied directly when the count-leading /
+;;;;; trailing-zeros value is consumed as a `brif` condition (the 2-op
+;;;;; wasm pattern `i32.ctz; br_if`, with no `i32.eqz` interposed for the
+;;;;; rules above to hinge on). We rewrite only the condition operand in
+;;;;; place; the brif's successors are unchanged so the CFG is preserved.
+
+;; `brif (ctz X) bt be` branches when `ctz(X) != 0`, i.e. LSB(X) == 0.
+(rule (simplify_skeleton (brif (ctz x_ty X) _ _))
+      (replace_branch_cond
+        (eq $I8 (band x_ty X (iconst_u x_ty 1)) (iconst_u x_ty 0))))
+
+;; `brif (clz X) bt be` branches when `clz(X) != 0`, i.e. MSB(X) == 0.
+(rule (simplify_skeleton (brif (clz x_ty X) _ _))
+      (replace_branch_cond (sge $I8 X (iconst_u x_ty 0))))

cranelift/codegen/src/prelude_opt.isle

@@ -92,7 +92,15 @@
         ;; The old instruction must define a single result value and `val` must
         ;; match its type. The new instruction need not define the same number
         ;; or types of results as the old instruction.
-        (ReplaceWithVal (inst Inst) (val Value))))
+        (ReplaceWithVal (inst Inst) (val Value))
+
+        ;; Narrow rewrite: replace the condition operand of a `brif`
+        ;; (argument 0) with the given value. The opcode and successor
+        ;; blocks stay the same, so the CFG is preserved. Use this to
+        ;; rewrite branch conditions whose producer simplifies in
+        ;; boolean-test context (e.g. `(brif (ctz x))` → branch on the
+        ;; LSB of x).
+        (ReplaceBranchCond (cond Value))))
 
 (decl pure inst_to_skeleton_inst_simplification (Inst) SkeletonInstSimplification)
 (rule (inst_to_skeleton_inst_simplification inst)
@@ -108,6 +116,10 @@
 (decl pure replace_with_val (Inst Value) SkeletonInstSimplification)
 (rule (replace_with_val inst val) (SkeletonInstSimplification.ReplaceWithVal inst val))
 
+(decl pure replace_branch_cond (Value) SkeletonInstSimplification)
+(rule (replace_branch_cond new_cond)
+      (SkeletonInstSimplification.ReplaceBranchCond new_cond))
+
 (convert Inst SkeletonInstSimplification inst_to_skeleton_inst_simplification)
 (convert Value SkeletonInstSimplification value_to_skeleton_inst_simplification)
 

cranelift/filetests/filetests/egraph/brif-cnt-cond.clif

@@ -0,0 +1,132 @@
+test optimize precise-output
+set opt_level=speed
+target x86_64
+
+;; `brif (ctz X)` / `brif (clz X)` — 2-op forms with no `icmp ne 0`
+;; interposed for the value-level rules in `opts/icmp.isle` to hinge on.
+;; The new `simplify_skeleton` rule rewrites the brif's cond operand
+;; in place via `ReplaceBranchCond`; successors are preserved so the
+;; CFG is unchanged.
+
+;; `brif (ctz X)` branches when `ctz(X) != 0`, i.e. LSB(X) == 0.
+function %brif_ctz_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = ctz v0
+    brif v1, block1, block2
+
+block1:
+    v2 = iconst.i32 100
+    return v2
+
+block2:
+    v3 = iconst.i32 200
+    return v3
+}
+
+; function %brif_ctz_i32(i32) -> i32 fast {
+; block0(v0: i32):
+;     v4 = iconst.i32 1
+;     v5 = band v0, v4  ; v4 = 1
+;     v6 = iconst.i32 0
+;     v7 = icmp eq v5, v6  ; v6 = 0
+;     brif v7, block1, block2
+;
+; block1:
+;     v2 = iconst.i32 100
+;     return v2  ; v2 = 100
+;
+; block2:
+;     v3 = iconst.i32 200
+;     return v3  ; v3 = 200
+; }
+
+;; `brif (clz X)` branches when `clz(X) != 0`, i.e. MSB(X) == 0.
+function %brif_clz_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = clz v0
+    brif v1, block1, block2
+
+block1:
+    v2 = iconst.i32 100
+    return v2
+
+block2:
+    v3 = iconst.i32 200
+    return v3
+}
+
+; function %brif_clz_i32(i32) -> i32 fast {
+; block0(v0: i32):
+;     v4 = iconst.i32 0
+;     v5 = icmp sge v0, v4  ; v4 = 0
+;     brif v5, block1, block2
+;
+; block1:
+;     v2 = iconst.i32 100
+;     return v2  ; v2 = 100
+;
+; block2:
+;     v3 = iconst.i32 200
+;     return v3  ; v3 = 200
+; }
+
+;; Same for i64.
+function %brif_ctz_i64(i64) -> i32 {
+block0(v0: i64):
+    v1 = ctz v0
+    brif v1, block1, block2
+
+block1:
+    v2 = iconst.i32 100
+    return v2
+
+block2:
+    v3 = iconst.i32 200
+    return v3
+}
+
+; function %brif_ctz_i64(i64) -> i32 fast {
+; block0(v0: i64):
+;     v4 = iconst.i64 1
+;     v5 = band v0, v4  ; v4 = 1
+;     v6 = iconst.i64 0
+;     v7 = icmp eq v5, v6  ; v6 = 0
+;     brif v7, block1, block2
+;
+; block1:
+;     v2 = iconst.i32 100
+;     return v2  ; v2 = 100
+;
+; block2:
+;     v3 = iconst.i32 200
+;     return v3  ; v3 = 200
+; }
+
+function %brif_clz_i64(i64) -> i32 {
+block0(v0: i64):
+    v1 = clz v0
+    brif v1, block1, block2
+
+block1:
+    v2 = iconst.i32 100
+    return v2
+
+block2:
+    v3 = iconst.i32 200
+    return v3
+}
+
+; function %brif_clz_i64(i64) -> i32 fast {
+; block0(v0: i64):
+;     v4 = iconst.i64 0
+;     v5 = icmp sge v0, v4  ; v4 = 0
+;     brif v5, block1, block2
+;
+; block1:
+;     v2 = iconst.i32 100
+;     return v2  ; v2 = 100
+;
+; block2:
+;     v3 = iconst.i32 200
+;     return v3  ; v3 = 200
+; }

tests/disas/ctz-clz-bool-condition.wat

@@ -108,14 +108,11 @@
 ;; wasm[0]::function[2]::if_ctz_bare_i32:
 ;;       pushq   %rbp
 ;;       movq    %rsp, %rbp
-;;       movl    $0x20, %esi
-;;       bsfl    %edx, %r9d
-;;       cmovel  %esi, %r9d
-;;       testl   %r9d, %r9d
-;;       jne     0xa4
-;;   9a: movl    $0xc8, %eax
-;;       jmp     0xa9
-;;   a4: movl    $0x64, %eax
+;;       testl   $1, %edx
+;;       je      0x9a
+;;   90: movl    $0xc8, %eax
+;;       jmp     0x9f
+;;   9a: movl    $0x64, %eax
 ;;       movq    %rbp, %rsp
 ;;       popq    %rbp
 ;;       retq
@@ -216,16 +213,11 @@
 ;; wasm[0]::function[11]::if_clz_bare_i32:
 ;;       pushq   %rbp
 ;;       movq    %rsp, %rbp
-;;       movq    $18446744073709551615, %rsi
-;;       bsrl    %edx, %r9d
-;;       cmovel  %esi, %r9d
-;;       movl    $0x1f, %eax
-;;       subl    %r9d, %eax
-;;       testl   %eax, %eax
-;;       jne     0x24d
-;;  243: movl    $0xc8, %eax
-;;       jmp     0x252
-;;  24d: movl    $0x64, %eax
+;;       testl   %edx, %edx
+;;       jge     0x236
+;;  22c: movl    $0xc8, %eax
+;;       jmp     0x23b
+;;  236: movl    $0x64, %eax
 ;;       movq    %rbp, %rsp
 ;;       popq    %rbp
 ;;       retq
@@ -244,10 +236,10 @@
 ;;       pushq   %rbp
 ;;       movq    %rsp, %rbp
 ;;       testq   %rdx, %rdx
-;;       jl      0x297
-;;  28d: movl    $0xc8, %eax
-;;       jmp     0x29c
-;;  297: movl    $0x64, %eax
+;;       jl      0x277
+;;  26d: movl    $0xc8, %eax
+;;       jmp     0x27c
+;;  277: movl    $0x64, %eax
 ;;       movq    %rbp, %rsp
 ;;       popq    %rbp
 ;;       retq
@@ -256,10 +248,10 @@
 ;;       pushq   %rbp
 ;;       movq    %rsp, %rbp
 ;;       testq   %rdx, %rdx
-;;       jge     0x2d7
-;;  2cd: movl    $0xc8, %eax
-;;       jmp     0x2dc
-;;  2d7: movl    $0x64, %eax
+;;       jge     0x2b7
+;;  2ad: movl    $0xc8, %eax
+;;       jmp     0x2bc
+;;  2b7: movl    $0x64, %eax
 ;;       movq    %rbp, %rsp
 ;;       popq    %rbp
 ;;       retq
@@ -271,10 +263,10 @@
 ;;       bsfl    %edx, %r9d
 ;;       cmovel  %esi, %r9d
 ;;       cmpl    $4, %r9d
-;;       je      0x325
-;;  31b: movl    $0xc8, %eax
-;;       jmp     0x32a
-;;  325: movl    $0x64, %eax
+;;       je      0x305
+;;  2fb: movl    $0xc8, %eax
+;;       jmp     0x30a
+;;  305: movl    $0x64, %eax
 ;;       movq    %rbp, %rsp
 ;;       popq    %rbp
 ;;       retq