aarch64: skip the extend of a scalar comparison's result

cranelift/codegen/src/isa/aarch64/lower.isle

@@ -1275,6 +1275,16 @@
       (if-let mem_op (is_sinkable_inst x))
       (load_acquire in flags (sink_atomic_load mem_op)))
 
+;; A scalar comparison is lowered by `lower_cond_result_bool`, whose every arm
+;; already has all upper bits clear, so extending it is a no-op.
+;;
+;; This does not hold true for vector comparisons, which instead produce a
+;; per-lane mask of either all ones or zeroes. rather than 0 or 1.
+(rule 1 (lower (uextend (fits_in_64 _) x @ (icmp _ _ (value_type (ty_int _)) _)))
+      (value_regs_get x 0))
+(rule 1 (lower (uextend (fits_in_64 _) x @ (fcmp _ _ (value_type (ty_scalar_float _)) _)))
+      (value_regs_get x 0))
+
 ;; Conversion to 128-bit needs a zero-extension of the lower bits and the upper
 ;; bits are all zero.
 (attr rule uextend_i128 (tag i128))
@@ -2172,6 +2182,8 @@
   (lower_cond_result_bool (emit_icmp cond x y)))
 
 ;; Helper to put a `CondResult` into a general register.
+;;
+;; Upper bits of the result register are zeroed.
 (decl lower_cond_result_bool (CondResult) Reg)
 (rule (lower_cond_result_bool (CondResult.Zero reg size))
   (value_regs_get (with_flags (cmp_imm size reg (u8_into_imm12 0)) (cset (Cond.Eq))) 0))

cranelift/filetests/filetests/isa/aarch64/uextend-sextend.clif

@@ -194,3 +194,181 @@ block0(v0: i32):
 ;   sxtw x0, w0
 ;   ret
 
+
+;; A comparison's result is produced by `cset`, which already zeroes the whole
+;; register, so extending it needs no instruction of its own.
+function %f_u_icmp_32(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+  v2 = icmp eq v0, v1
+  v3 = uextend.i32 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   subs wzr, w0, w1
+;   cset x0, eq
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp w0, w1
+;   cset x0, eq
+;   ret
+
+;; Any output width, since `cset` defines all 64 bits.
+function %f_u_icmp_64(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+  v2 = icmp slt v0, v1
+  v3 = uextend.i64 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   subs wzr, w0, w1
+;   cset x0, lt
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp w0, w1
+;   cset x0, lt
+;   ret
+
+function %f_u_icmp_16(i64, i64) -> i16 {
+block0(v0: i64, v1: i64):
+  v2 = icmp ugt v0, v1
+  v3 = uextend.i16 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   subs xzr, x0, x1
+;   cset x0, hi
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp x0, x1
+;   cset x0, hi
+;   ret
+
+;; An `i128` comparison takes several instructions to compute its flags, but
+;; still ends in one `cset`.
+function %f_u_icmp_i128(i128, i128) -> i32 {
+block0(v0: i128, v1: i128):
+  v2 = icmp ult v0, v1
+  v3 = uextend.i32 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   subs xzr, x0, x2
+;   sbcs xzr, x1, x3
+;   cset x0, lo
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp x0, x2
+;   sbcs xzr, x1, x3
+;   cset x0, lo
+;   ret
+
+;; `fcmp` lowers through the same helper.
+function %f_u_fcmp(f64, f64) -> i32 {
+block0(v0: f64, v1: f64):
+  v2 = fcmp lt v0, v1
+  v3 = uextend.i32 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   fcmp d0, d1
+;   cset x0, mi
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   fcmp d0, d1
+;   cset x0, mi
+;   ret
+
+;; `fcmp one` needs two `cset`s combined with an `orr`, which is also a 32-bit
+;; write of 0 or 1.
+function %f_u_fcmp_one(f32, f32) -> i64 {
+block0(v0: f32, v1: f32):
+  v2 = fcmp one v0, v1
+  v3 = uextend.i64 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   fcmp s0, s1
+;   cset x4, mi
+;   cset x6, gt
+;   orr w0, w4, w6
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   fcmp s0, s1
+;   cset x4, mi
+;   cset x6, gt
+;   orr w0, w4, w6
+;   ret
+
+;; Several extends of one comparison share its register rather than each
+;; re-emitting the comparison.
+function %f_u_icmp_multi_use(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+  v2 = icmp eq v0, v1
+  v3 = uextend.i64 v2
+  v4 = uextend.i64 v2
+  v5 = ishl v4, v3
+  return v5
+}
+
+; VCode:
+; block0:
+;   subs wzr, w0, w1
+;   cset x5, eq
+;   lsl x0, x5, x5
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp w0, w1
+;   cset x5, eq
+;   lsl x0, x5, x5
+;   ret
+
+;; Extending to `i128` needs a zero upper half, so it keeps its own rule.
+function %f_u_icmp_i128_out(i32, i32) -> i128 {
+block0(v0: i32, v1: i32):
+  v2 = icmp eq v0, v1
+  v3 = uextend.i128 v2
+  return v3
+}
+
+; VCode:
+; block0:
+;   subs wzr, w0, w1
+;   cset x6, eq
+;   uxtb w0, w6
+;   movz x1, #0
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   cmp w0, w1
+;   cset x6, eq
+;   uxtb w0, w6
+;   mov x1, #0
+;   ret
+

cranelift/filetests/filetests/runtests/extend-of-compare.clif

@@ -0,0 +1,152 @@
+test interpret
+test run
+target x86_64
+target s390x
+target aarch64
+target riscv64
+target pulley64
+target pulley64be
+
+;; A comparison's result is 0 or 1, so extending it must produce exactly 0 or 1
+;; in the wider type -- a backend that skips the extend because its comparison
+;; already wrote the whole register must leave no other bits set.
+
+function %uext_icmp_eq_32(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+    v2 = icmp eq v0, v1
+    v3 = uextend.i32 v2
+    return v3
+}
+; run: %uext_icmp_eq_32(0, 0) == 1
+; run: %uext_icmp_eq_32(1, 0) == 0
+; run: %uext_icmp_eq_32(-1, -1) == 1
+; run: %uext_icmp_eq_32(-1, 0x7fffffff) == 0
+
+function %uext_icmp_slt_64(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+    v2 = icmp slt v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_icmp_slt_64(-1, 0) == 1
+; run: %uext_icmp_slt_64(0, -1) == 0
+; run: %uext_icmp_slt_64(0x80000000, 0x7fffffff) == 1
+; run: %uext_icmp_slt_64(0, 0) == 0
+
+function %uext_icmp_ult_64(i64, i64) -> i64 {
+block0(v0: i64, v1: i64):
+    v2 = icmp ult v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_icmp_ult_64(0, -1) == 1
+; run: %uext_icmp_ult_64(-1, 0) == 0
+; run: %uext_icmp_ult_64(-1, -1) == 0
+
+;; A narrow output width, and one narrower than the comparison's operands.
+function %uext_icmp_16(i64, i64) -> i16 {
+block0(v0: i64, v1: i64):
+    v2 = icmp ugt v0, v1
+    v3 = uextend.i16 v2
+    return v3
+}
+; run: %uext_icmp_16(1, 0) == 1
+; run: %uext_icmp_16(0, 1) == 0
+
+;; Narrow comparison operands: the `i8` inputs must not leak their upper bits
+;; into the result.
+function %uext_icmp_i8_operands(i8, i8) -> i64 {
+block0(v0: i8, v1: i8):
+    v2 = icmp eq v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_icmp_i8_operands(-1, -1) == 1
+; run: %uext_icmp_i8_operands(-1, 0xff) == 1
+; run: %uext_icmp_i8_operands(0x7f, 0x80) == 0
+
+;; The extended value has to be usable as an integer, not just compared against
+;; 0 or 1: an `imul` by it exposes any stray upper bit.
+function %uext_icmp_arith(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+    v2 = icmp eq v0, v1
+    v3 = uextend.i64 v2
+    v4 = iconst.i64 0x1234_5678_9abc_def0
+    v5 = imul v3, v4
+    return v5
+}
+; run: %uext_icmp_arith(0, 0) == 0x123456789abcdef0
+; run: %uext_icmp_arith(0, 1) == 0
+
+;; Several extends of one comparison.
+function %uext_icmp_multi_use(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+    v2 = icmp eq v0, v1
+    v3 = uextend.i64 v2
+    v4 = uextend.i64 v2
+    v5 = iadd v3, v4
+    return v5
+}
+; run: %uext_icmp_multi_use(7, 7) == 2
+; run: %uext_icmp_multi_use(7, 8) == 0
+
+;; Sign-extending a comparison is the same as zero-extending it, which is a
+;; mid-end rewrite; check it still holds end to end.
+function %sext_icmp(i32, i32) -> i64 {
+block0(v0: i32, v1: i32):
+    v2 = icmp eq v0, v1
+    v3 = sextend.i64 v2
+    return v3
+}
+; run: %sext_icmp(5, 5) == 1
+; run: %sext_icmp(5, 6) == 0
+
+;; Floating point, where the unordered cases are what a 0-or-1 assumption is
+;; most likely to get wrong.
+function %uext_fcmp_lt(f64, f64) -> i64 {
+block0(v0: f64, v1: f64):
+    v2 = fcmp lt v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_fcmp_lt(0x0.0, 0x1.0) == 1
+; run: %uext_fcmp_lt(0x1.0, 0x0.0) == 0
+; run: %uext_fcmp_lt(-0x0.0, 0x0.0) == 0
+; run: %uext_fcmp_lt(+NaN, 0x0.0) == 0
+; run: %uext_fcmp_lt(0x0.0, +NaN) == 0
+
+;; `ne` is true for NaN.
+function %uext_fcmp_ne(f32, f32) -> i64 {
+block0(v0: f32, v1: f32):
+    v2 = fcmp ne v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_fcmp_ne(0x1.0, 0x1.0) == 0
+; run: %uext_fcmp_ne(0x1.0, 0x2.0) == 1
+; run: %uext_fcmp_ne(+NaN, +NaN) == 1
+; run: %uext_fcmp_ne(-0x0.0, 0x0.0) == 0
+
+;; `one` and `ueq` are the two conditions aarch64 builds from a pair of `cset`s
+;; combined with `orr`.
+function %uext_fcmp_one(f32, f32) -> i64 {
+block0(v0: f32, v1: f32):
+    v2 = fcmp one v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_fcmp_one(0x1.0, 0x2.0) == 1
+; run: %uext_fcmp_one(0x1.0, 0x1.0) == 0
+; run: %uext_fcmp_one(+NaN, 0x1.0) == 0
+; run: %uext_fcmp_one(+NaN, +NaN) == 0
+
+function %uext_fcmp_ueq(f64, f64) -> i64 {
+block0(v0: f64, v1: f64):
+    v2 = fcmp ueq v0, v1
+    v3 = uextend.i64 v2
+    return v3
+}
+; run: %uext_fcmp_ueq(0x1.0, 0x1.0) == 1
+; run: %uext_fcmp_ueq(0x1.0, 0x2.0) == 0
+; run: %uext_fcmp_ueq(+NaN, 0x1.0) == 1
+; run: %uext_fcmp_ueq(-0x0.0, 0x0.0) == 1