Less unnecessary zero and sign extensions

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

@@ -348,6 +348,45 @@ fn put_input_in_rse<C: LowerCtx<I = Inst>>(
         let out_ty = ctx.output_ty(insn, 0);
         let out_bits = ty_bits(out_ty);
 
+        // Is this a zero-extend or sign-extend and can we handle that with a register-mode operator?
+        if op == Opcode::Uextend || op == Opcode::Sextend {
+            let sign_extend = op == Opcode::Sextend;
+            let inner_ty = ctx.input_ty(insn, 0);
+            let inner_bits = ty_bits(inner_ty);
+            assert!(inner_bits < out_bits);
+            if match (sign_extend, narrow_mode) {
+                // A single zero-extend or sign-extend is equal to itself.
+                (_, NarrowValueMode::None) => true,
+                // Two zero-extends or sign-extends in a row is equal to a single zero-extend or sign-extend.
+                (false, NarrowValueMode::ZeroExtend32) | (false, NarrowValueMode::ZeroExtend64) => {
+                    true
+                }
+                (true, NarrowValueMode::SignExtend32) | (true, NarrowValueMode::SignExtend64) => {
+                    true
+                }
+                // A zero-extend and a sign-extend in a row is not equal to a single zero-extend or sign-extend
+                (false, NarrowValueMode::SignExtend32) | (false, NarrowValueMode::SignExtend64) => {
+                    false
+                }
+                (true, NarrowValueMode::ZeroExtend32) | (true, NarrowValueMode::ZeroExtend64) => {
+                    false
+                }
+            } {
+                let extendop = match (sign_extend, inner_bits) {
+                    (true, 8) => ExtendOp::SXTB,
+                    (false, 8) => ExtendOp::UXTB,
+                    (true, 16) => ExtendOp::SXTH,
+                    (false, 16) => ExtendOp::UXTH,
+                    (true, 32) => ExtendOp::SXTW,
+                    (false, 32) => ExtendOp::UXTW,
+                    _ => unreachable!(),
+                };
+                let reg =
+                    put_input_in_reg(ctx, InsnInput { insn, input: 0 }, NarrowValueMode::None);
+                return ResultRSE::RegExtend(reg, extendop);
+            }
+        }
+
         // If `out_ty` is smaller than 32 bits and we need to zero- or sign-extend,
         // then get the result into a register and return an Extend-mode operand on
         // that register.
@@ -381,28 +420,6 @@ fn put_input_in_rse<C: LowerCtx<I = Inst>>(
             };
             return ResultRSE::RegExtend(reg, extendop);
         }
-
-        // Is this a zero-extend or sign-extend and can we handle that with a register-mode operator?
-        if op == Opcode::Uextend || op == Opcode::Sextend {
-            assert!(out_bits == 32 || out_bits == 64);
-            let sign_extend = op == Opcode::Sextend;
-            let inner_ty = ctx.input_ty(insn, 0);
-            let inner_bits = ty_bits(inner_ty);
-            assert!(inner_bits < out_bits);
-            let extendop = match (sign_extend, inner_bits) {
-                (true, 1) => ExtendOp::SXTB,
-                (false, 1) => ExtendOp::UXTB,
-                (true, 8) => ExtendOp::SXTB,
-                (false, 8) => ExtendOp::UXTB,
-                (true, 16) => ExtendOp::SXTH,
-                (false, 16) => ExtendOp::UXTH,
-                (true, 32) => ExtendOp::SXTW,
-                (false, 32) => ExtendOp::UXTW,
-                _ => unreachable!(),
-            };
-            let reg = put_input_in_reg(ctx, input, NarrowValueMode::None);
-            return ResultRSE::RegExtend(reg, extendop);
-        }
     }
 
     ResultRSE::from_rs(put_input_in_rs(ctx, input, narrow_mode))

cranelift/filetests/filetests/vcode/aarch64/extend-op.clif

@@ -11,7 +11,6 @@ block0(v0: i8):
 
 ; check:  stp fp, lr, [sp, #-16]!
 ; nextln:  mov fp, sp
-; nextln:  sxtb x0, w0
 ; nextln:  movz x1, #42
 ; nextln:  add x0, x1, x0, SXTB
 ; nextln:  mov sp, fp

cranelift/filetests/filetests/vcode/aarch64/heap_addr.clif

@@ -20,8 +20,7 @@ block0(v0: i64, v1: i32):
 ; nextln: subs wzr, w1, w2
 ; nextln: b.ls label1 ; b label2
 ; check: Block 1:
-; check: mov w3, w1
+; check: add x0, x0, x1, UXTW
-; check: add x0, x0, x3, UXTW
 ; nextln: subs wzr, w1, w2
 ; nextln: movz x1, #0
 ; nextln: csel x0, x1, x0, hi
@@ -46,8 +45,7 @@ block0(v0: i64, v1: i32):
 ; nextln: subs wzr, w1, #65536
 ; nextln: b.ls label1 ; b label2
 ; check: Block 1:
-; check: mov w2, w1
+; check: add x0, x0, x1, UXTW
-; check: add x0, x0, x2, UXTW
 ; nextln: subs wzr, w1, #65536
 ; nextln: movz x1, #0
 ; nextln: csel x0, x1, x0, hi

cranelift/filetests/filetests/vcode/aarch64/narrow-arithmetic.clif

@@ -49,7 +49,6 @@ block0(v0: i32, v1: i8):
 
 ; check: stp fp, lr, [sp, #-16]!
 ; nextln: mov fp, sp
-; nextln: sxtb w1, w1
 ; nextln: add w0, w0, w1, SXTB
 ; nextln: mov sp, fp
 ; nextln: ldp fp, lr, [sp], #16
@@ -64,7 +63,6 @@ block0(v0: i64, v1: i32):
 
 ; check: stp fp, lr, [sp, #-16]!
 ; nextln: mov fp, sp
-; nextln: sxtw x1, w1
 ; nextln: add x0, x0, x1, SXTW
 ; nextln: mov sp, fp
 ; nextln: ldp fp, lr, [sp], #16