Skip to content

Commit f6a7288

Browse files
ggreifclaude
andauthored
cranelift: fold ctz/clz directly into brif cond via simplify_skeleton (bytecodealliance#13343)
* cranelift: fold `ctz`/`clz` directly into `brif` cond via `simplify_skeleton` The mid-end rules added in bytecodealliance#13332 hinge on an `icmp eq/ne (ctz/clz X) 0` shape — i.e. the wasm 3-op pattern `i32.ctz; i32.eqz; br_if`. Frontends that emit the 2-op form `i32.ctz; br_if` (e.g. Motoko's `moc` after its `and 1; eqz; br_if` → `ctz; br_if` byte-size peephole) feed `(brif (ctz X))` into cranelift with no `icmp` for the existing rules to match. This commit extends `simplify_skeleton` to rewrite the *condition operand* of an existing `brif` in place, without touching its opcode or successor blocks (CFG-preserving by construction). A new `SkeletonInstSimplification` variant `ReplaceBranchCond(Value)` carries the new condition; the egraph driver applies it by writing through `inst_args_mut`. Two ISLE rules in `opts/icmp.isle` rewrite `(brif (ctz X) bt be)` and `(brif (clz X) bt be)` to brifs over the equivalent bit-extract form: brif (ctz X) bt be → brif (eq (band X 1) 0) bt be brif (clz X) bt be → brif (sge X 0) bt be End-to-end lowering on the resulting brif then composes with existing backend `icmp+brif` fusion to produce: x86_64 brif (ctz X): `testl $1, %edi; je` x86_64 brif (clz X): `testl %edi, %edi; jge` aarch64 brif (ctz X): `tbz w0, #0` — single-instruction test-and-branch This subsumes the backend-side x64 rules added in bytecodealliance#13334 and the aarch64 rules in bytecodealliance#13336 (and yields tighter aarch64 code than bytecodealliance#13336 did). The driver still rejects non-`brif` branches and rejects non-`ReplaceBranchCond` simplification variants on `brif` (a `Replace inst` of a brif would risk changing successor block IDs and is left to a future, broader extension). Filetest `egraph/brif-cnt-cond.clif` covers ctz/clz over i32/i64 in the 2-op form. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * rustfmt: collapse `is_branch` && opcode-guard onto one line * tests/disas: re-bless ctz/clz-bool-condition for new mid-end fold The new `simplify_skeleton`-on-`brif` rule rewrites the 2-op `if (ctz/clz x)` cases that bytecodealliance#13332's commentary noted were the non-icmp-mediated holdouts. Bare-form lowering shrinks from ~9 instructions (bsf/bsr + cmov + test + jne + …) to `testl $1, %edx; je` (ctz) and `testl %edx, %edx; jge` (clz). Offsets on the subsequent non-bare functions shift down to match. --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent eef7379 commit f6a7288

5 files changed

Lines changed: 200 additions & 31 deletions

File tree

cranelift/codegen/src/egraph/mod.rs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -811,6 +811,14 @@ where
811811
);
812812
(inst, Some(val))
813813
}
814+
// `ReplaceBranchCond` is unconditionally accepted — the
815+
// opcode and successors don't change, so we can't use the
816+
// cost-based ranking the other variants do. The first such
817+
// candidate wins; ISLE rule ordering picks the form.
818+
SkeletonInstSimplification::ReplaceBranchCond { cond } => {
819+
log::trace!(" -> simplify_skeleton: replace `brif` cond with {cond}");
820+
return Some(SkeletonInstSimplification::ReplaceBranchCond { cond });
821+
}
814822
};
815823

816824
if cfg!(debug_assertions) {
@@ -1119,6 +1127,16 @@ impl<'a> EgraphPass<'a> {
11191127
}
11201128
SkeletonInstSimplification::Replace { inst } => (inst, None),
11211129
SkeletonInstSimplification::ReplaceWithVal { inst, val } => (inst, Some(val)),
1130+
SkeletonInstSimplification::ReplaceBranchCond { cond } => {
1131+
// Swap the condition operand of the existing `brif` in
1132+
// place. Successors stay; CFG is preserved.
1133+
debug_assert_eq!(
1134+
cursor.func.dfg.insts[old_inst].opcode(),
1135+
crate::ir::Opcode::Brif,
1136+
);
1137+
cursor.func.dfg.inst_args_mut(old_inst)[0] = cond;
1138+
return;
1139+
}
11221140
};
11231141

11241142
// Replace the old instruction with the new one.

cranelift/codegen/src/opts/icmp.isle

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,3 +450,18 @@
450450
;; clz(X) != 0 iff the MSB of X is 0, i.e. X is signed-non-negative.
451451
(rule (simplify (ne result_ty (clz x_ty X) (iconst_u _ 0)))
452452
(sge result_ty X (iconst_u x_ty 0)))
453+
454+
;;;;; Same simplifications applied directly when the count-leading /
455+
;;;;; trailing-zeros value is consumed as a `brif` condition (the 2-op
456+
;;;;; wasm pattern `i32.ctz; br_if`, with no `i32.eqz` interposed for the
457+
;;;;; rules above to hinge on). We rewrite only the condition operand in
458+
;;;;; place; the brif's successors are unchanged so the CFG is preserved.
459+
460+
;; `brif (ctz X) bt be` branches when `ctz(X) != 0`, i.e. LSB(X) == 0.
461+
(rule (simplify_skeleton (brif (ctz x_ty X) _ _))
462+
(replace_branch_cond
463+
(eq $I8 (band x_ty X (iconst_u x_ty 1)) (iconst_u x_ty 0))))
464+
465+
;; `brif (clz X) bt be` branches when `clz(X) != 0`, i.e. MSB(X) == 0.
466+
(rule (simplify_skeleton (brif (clz x_ty X) _ _))
467+
(replace_branch_cond (sge $I8 X (iconst_u x_ty 0))))

cranelift/codegen/src/prelude_opt.isle

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,15 @@
9292
;; The old instruction must define a single result value and `val` must
9393
;; match its type. The new instruction need not define the same number
9494
;; or types of results as the old instruction.
95-
(ReplaceWithVal (inst Inst) (val Value))))
95+
(ReplaceWithVal (inst Inst) (val Value))
96+
97+
;; Narrow rewrite: replace the condition operand of a `brif`
98+
;; (argument 0) with the given value. The opcode and successor
99+
;; blocks stay the same, so the CFG is preserved. Use this to
100+
;; rewrite branch conditions whose producer simplifies in
101+
;; boolean-test context (e.g. `(brif (ctz x))` → branch on the
102+
;; LSB of x).
103+
(ReplaceBranchCond (cond Value))))
96104

97105
(decl pure inst_to_skeleton_inst_simplification (Inst) SkeletonInstSimplification)
98106
(rule (inst_to_skeleton_inst_simplification inst)
@@ -108,6 +116,10 @@
108116
(decl pure replace_with_val (Inst Value) SkeletonInstSimplification)
109117
(rule (replace_with_val inst val) (SkeletonInstSimplification.ReplaceWithVal inst val))
110118

119+
(decl pure replace_branch_cond (Value) SkeletonInstSimplification)
120+
(rule (replace_branch_cond new_cond)
121+
(SkeletonInstSimplification.ReplaceBranchCond new_cond))
122+
111123
(convert Inst SkeletonInstSimplification inst_to_skeleton_inst_simplification)
112124
(convert Value SkeletonInstSimplification value_to_skeleton_inst_simplification)
113125

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
test optimize precise-output
2+
set opt_level=speed
3+
target x86_64
4+
5+
;; `brif (ctz X)` / `brif (clz X)` — 2-op forms with no `icmp ne 0`
6+
;; interposed for the value-level rules in `opts/icmp.isle` to hinge on.
7+
;; The new `simplify_skeleton` rule rewrites the brif's cond operand
8+
;; in place via `ReplaceBranchCond`; successors are preserved so the
9+
;; CFG is unchanged.
10+
11+
;; `brif (ctz X)` branches when `ctz(X) != 0`, i.e. LSB(X) == 0.
12+
function %brif_ctz_i32(i32) -> i32 {
13+
block0(v0: i32):
14+
v1 = ctz v0
15+
brif v1, block1, block2
16+
17+
block1:
18+
v2 = iconst.i32 100
19+
return v2
20+
21+
block2:
22+
v3 = iconst.i32 200
23+
return v3
24+
}
25+
26+
; function %brif_ctz_i32(i32) -> i32 fast {
27+
; block0(v0: i32):
28+
; v4 = iconst.i32 1
29+
; v5 = band v0, v4 ; v4 = 1
30+
; v6 = iconst.i32 0
31+
; v7 = icmp eq v5, v6 ; v6 = 0
32+
; brif v7, block1, block2
33+
;
34+
; block1:
35+
; v2 = iconst.i32 100
36+
; return v2 ; v2 = 100
37+
;
38+
; block2:
39+
; v3 = iconst.i32 200
40+
; return v3 ; v3 = 200
41+
; }
42+
43+
;; `brif (clz X)` branches when `clz(X) != 0`, i.e. MSB(X) == 0.
44+
function %brif_clz_i32(i32) -> i32 {
45+
block0(v0: i32):
46+
v1 = clz v0
47+
brif v1, block1, block2
48+
49+
block1:
50+
v2 = iconst.i32 100
51+
return v2
52+
53+
block2:
54+
v3 = iconst.i32 200
55+
return v3
56+
}
57+
58+
; function %brif_clz_i32(i32) -> i32 fast {
59+
; block0(v0: i32):
60+
; v4 = iconst.i32 0
61+
; v5 = icmp sge v0, v4 ; v4 = 0
62+
; brif v5, block1, block2
63+
;
64+
; block1:
65+
; v2 = iconst.i32 100
66+
; return v2 ; v2 = 100
67+
;
68+
; block2:
69+
; v3 = iconst.i32 200
70+
; return v3 ; v3 = 200
71+
; }
72+
73+
;; Same for i64.
74+
function %brif_ctz_i64(i64) -> i32 {
75+
block0(v0: i64):
76+
v1 = ctz v0
77+
brif v1, block1, block2
78+
79+
block1:
80+
v2 = iconst.i32 100
81+
return v2
82+
83+
block2:
84+
v3 = iconst.i32 200
85+
return v3
86+
}
87+
88+
; function %brif_ctz_i64(i64) -> i32 fast {
89+
; block0(v0: i64):
90+
; v4 = iconst.i64 1
91+
; v5 = band v0, v4 ; v4 = 1
92+
; v6 = iconst.i64 0
93+
; v7 = icmp eq v5, v6 ; v6 = 0
94+
; brif v7, block1, block2
95+
;
96+
; block1:
97+
; v2 = iconst.i32 100
98+
; return v2 ; v2 = 100
99+
;
100+
; block2:
101+
; v3 = iconst.i32 200
102+
; return v3 ; v3 = 200
103+
; }
104+
105+
function %brif_clz_i64(i64) -> i32 {
106+
block0(v0: i64):
107+
v1 = clz v0
108+
brif v1, block1, block2
109+
110+
block1:
111+
v2 = iconst.i32 100
112+
return v2
113+
114+
block2:
115+
v3 = iconst.i32 200
116+
return v3
117+
}
118+
119+
; function %brif_clz_i64(i64) -> i32 fast {
120+
; block0(v0: i64):
121+
; v4 = iconst.i64 0
122+
; v5 = icmp sge v0, v4 ; v4 = 0
123+
; brif v5, block1, block2
124+
;
125+
; block1:
126+
; v2 = iconst.i32 100
127+
; return v2 ; v2 = 100
128+
;
129+
; block2:
130+
; v3 = iconst.i32 200
131+
; return v3 ; v3 = 200
132+
; }

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

Lines changed: 22 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -108,14 +108,11 @@
108108
;; wasm[0]::function[2]::if_ctz_bare_i32:
109109
;; pushq %rbp
110110
;; movq %rsp, %rbp
111-
;; movl $0x20, %esi
112-
;; bsfl %edx, %r9d
113-
;; cmovel %esi, %r9d
114-
;; testl %r9d, %r9d
115-
;; jne 0xa4
116-
;; 9a: movl $0xc8, %eax
117-
;; jmp 0xa9
118-
;; a4: movl $0x64, %eax
111+
;; testl $1, %edx
112+
;; je 0x9a
113+
;; 90: movl $0xc8, %eax
114+
;; jmp 0x9f
115+
;; 9a: movl $0x64, %eax
119116
;; movq %rbp, %rsp
120117
;; popq %rbp
121118
;; retq
@@ -216,16 +213,11 @@
216213
;; wasm[0]::function[11]::if_clz_bare_i32:
217214
;; pushq %rbp
218215
;; movq %rsp, %rbp
219-
;; movq $18446744073709551615, %rsi
220-
;; bsrl %edx, %r9d
221-
;; cmovel %esi, %r9d
222-
;; movl $0x1f, %eax
223-
;; subl %r9d, %eax
224-
;; testl %eax, %eax
225-
;; jne 0x24d
226-
;; 243: movl $0xc8, %eax
227-
;; jmp 0x252
228-
;; 24d: movl $0x64, %eax
216+
;; testl %edx, %edx
217+
;; jge 0x236
218+
;; 22c: movl $0xc8, %eax
219+
;; jmp 0x23b
220+
;; 236: movl $0x64, %eax
229221
;; movq %rbp, %rsp
230222
;; popq %rbp
231223
;; retq
@@ -244,10 +236,10 @@
244236
;; pushq %rbp
245237
;; movq %rsp, %rbp
246238
;; testq %rdx, %rdx
247-
;; jl 0x297
248-
;; 28d: movl $0xc8, %eax
249-
;; jmp 0x29c
250-
;; 297: movl $0x64, %eax
239+
;; jl 0x277
240+
;; 26d: movl $0xc8, %eax
241+
;; jmp 0x27c
242+
;; 277: movl $0x64, %eax
251243
;; movq %rbp, %rsp
252244
;; popq %rbp
253245
;; retq
@@ -256,10 +248,10 @@
256248
;; pushq %rbp
257249
;; movq %rsp, %rbp
258250
;; testq %rdx, %rdx
259-
;; jge 0x2d7
260-
;; 2cd: movl $0xc8, %eax
261-
;; jmp 0x2dc
262-
;; 2d7: movl $0x64, %eax
251+
;; jge 0x2b7
252+
;; 2ad: movl $0xc8, %eax
253+
;; jmp 0x2bc
254+
;; 2b7: movl $0x64, %eax
263255
;; movq %rbp, %rsp
264256
;; popq %rbp
265257
;; retq
@@ -271,10 +263,10 @@
271263
;; bsfl %edx, %r9d
272264
;; cmovel %esi, %r9d
273265
;; cmpl $4, %r9d
274-
;; je 0x325
275-
;; 31b: movl $0xc8, %eax
276-
;; jmp 0x32a
277-
;; 325: movl $0x64, %eax
266+
;; je 0x305
267+
;; 2fb: movl $0xc8, %eax
268+
;; jmp 0x30a
269+
;; 305: movl $0x64, %eax
278270
;; movq %rbp, %rsp
279271
;; popq %rbp
280272
;; retq

0 commit comments

Comments
 (0)