Merge pull request #1687 from fitzgen/sign-extend-immediates

cranelift: Sign extend `Imm64` immediates

cranelift/codegen/meta/src/gen_inst.rs

@@ -874,17 +874,32 @@ fn gen_format_constructor(format: &InstructionFormat, fmt: &mut Formatter) {
         args.join(", ")
     );
 
+    let imms_need_sign_extension = format
+        .imm_fields
+        .iter()
+        .any(|f| f.kind.rust_type == "ir::immediates::Imm64");
+
     fmt.doc_comment(format.to_string());
     fmt.line("#[allow(non_snake_case)]");
     fmtln!(fmt, "fn {} {{", proto);
     fmt.indent(|fmt| {
         // Generate the instruction data.
-        fmtln!(fmt, "let data = ir::InstructionData::{} {{", format.name);
+        fmtln!(
+            fmt,
+            "let{} data = ir::InstructionData::{} {{",
+            if imms_need_sign_extension { " mut" } else { "" },
+            format.name
+        );
         fmt.indent(|fmt| {
             fmt.line("opcode,");
             gen_member_inits(format, fmt);
         });
         fmtln!(fmt, "};");
+
+        if imms_need_sign_extension {
+            fmtln!(fmt, "data.sign_extend_immediates(ctrl_typevar);");
+        }
+
         fmt.line("self.build(data, ctrl_typevar)");
     });
     fmtln!(fmt, "}");

cranelift/codegen/src/ir/immediates.rs

@@ -62,6 +62,21 @@ impl Imm64 {
     pub fn bits(&self) -> i64 {
         self.0
     }
+
+    /// Sign extend this immediate as if it were a signed integer of the given
+    /// power-of-two width.
+    pub fn sign_extend_from_width(&mut self, bit_width: u16) {
+        debug_assert!(bit_width.is_power_of_two());
+
+        if bit_width >= 64 {
+            return;
+        }
+
+        let bit_width = bit_width as i64;
+        let delta = 64 - bit_width;
+        let sign_extended = (self.0 << delta) >> delta;
+        *self = Imm64(sign_extended);
+    }
 }
 
 impl Into<i64> for Imm64 {

cranelift/codegen/src/ir/instructions.rs

@@ -296,6 +296,39 @@ impl InstructionData {
             }
         }
     }
+
+    #[inline]
+    pub(crate) fn sign_extend_immediates(&mut self, ctrl_typevar: Type) {
+        if ctrl_typevar.is_invalid() {
+            return;
+        }
+
+        let bit_width = ctrl_typevar.bits();
+
+        match self {
+            Self::BinaryImm {
+                opcode,
+                arg: _,
+                imm,
+            } => {
+                if matches!(opcode, Opcode::SdivImm | Opcode::SremImm) {
+                    imm.sign_extend_from_width(bit_width);
+                }
+            }
+            Self::IntCompareImm {
+                opcode,
+                arg: _,
+                cond,
+                imm,
+            } => {
+                debug_assert_eq!(*opcode, Opcode::IcmpImm);
+                if cond.unsigned() != *cond {
+                    imm.sign_extend_from_width(bit_width);
+                }
+            }
+            _ => {}
+        }
+    }
 }
 
 /// Information about branch and jump instructions.

cranelift/filetests/filetests/simple_preopt/sign_extend.clif

@@ -0,0 +1,39 @@
+test simple_preopt
+target x86_64
+
+;; Tests for sign-extending immediates.
+
+function %sign_extend_signed_icmp(i8) -> b1 {
+block0(v0: i8):
+    ; 255 = -1 as u8
+    v1 = iconst.i8 255
+    v2 = icmp sge v0, v1
+    ; check: v2 = icmp_imm sge v0, -1
+    return v2
+}
+
+function %do_not_sign_extend_unsigned_icmp(i8) -> b1 {
+block0(v0: i8):
+    v1 = iconst.i8 255
+    v2 = icmp uge v0, v1
+    ; check: v2 = icmp_imm uge v0, 255
+    return v2
+}
+
+function %sign_extend_sdiv(i8) -> i8 {
+block0(v0: i8):
+    ; 255 = -1 as u8
+    v1 = iconst.i8 255
+    v2 = sdiv v0, v1
+    ; check: v2 = sdiv_imm v0, -1
+    return v2
+}
+
+function %sign_extend_srem(i8) -> i8 {
+block0(v0: i8):
+    ; 255 = -1 as u8
+    v1 = iconst.i8 255
+    v2 = srem v0, v1
+    ; check: v2 = srem_imm v0, -1
+    return v2
+}