Optimize load/store with an iadd_imm operand.

Fold the immediate into the load/store offset when possible.

cranelift/filetests/postopt/fold_offset_into_address.clif

@@ -0,0 +1,18 @@
+test postopt
+target x86_64
+
+; Fold the immediate of an iadd_imm into an address offset.
+
+function u0:0(i64 vmctx) -> i64 {
+ebb0(v0: i64):
+                     v1 = iadd_imm.i64 v0, 16
+[RexOp1ldDisp8#808b] v2 = load.i64 notrap aligned v1
+[Op1ret#c3]          return v2
+}
+
+; sameln: function u0:0(i64 vmctx) -> i64 fast {
+; nextln: ebb0(v0: i64):
+; nextln:                                     v1 = iadd_imm v0, 16
+; nextln: [RexOp1ldDisp8#808b]                v2 = load.i64 notrap aligned v0+16
+; nextln: [Op1ret#c3]                         return v2
+; nextln: }

lib/codegen/src/ir/immediates.rs

@@ -224,6 +224,16 @@ impl Offset32 {
             None
         }
     }
+
+    /// Add in the signed number `x` if possible.
+    pub fn try_add_i64(self, x: i64) -> Option<Self> {
+        let casted = x as i32;
+        if casted as i64 == x {
+            self.0.checked_add(casted).map(Self::new)
+        } else {
+            None
+        }
+    }
 }
 
 impl Into<i32> for Offset32 {

lib/codegen/src/postopt.rs

@@ -179,6 +179,7 @@ struct MemOpInfo {
 }
 
 fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa) {
+    // Look for simple loads and stores we can optimize.
     let info = match pos.func.dfg[inst] {
         InstructionData::Load {
             opcode,
@@ -209,30 +210,27 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
         _ => return,
     };
 
-    let add_args = if let ValueDef::Result(result_inst, _) = pos.func.dfg.value_def(info.arg) {
+    // Examine the instruction that defines the address operand.
+    if let ValueDef::Result(result_inst, _) = pos.func.dfg.value_def(info.arg) {
         match pos.func.dfg[result_inst] {
             InstructionData::Binary {
                 opcode: Opcode::Iadd,
                 args,
-            } => args,
+            } => match info.opcode {
-            _ => return,
+                // Operand is an iadd. Fold it into a memory address with a complex address mode.
-        }
-    } else {
-        return;
-    };
-
-    match info.opcode {
                 Opcode::Load => {
-            pos.func
+                    pos.func.dfg.replace(inst).load_complex(
-                .dfg
+                        info.itype,
-                .replace(inst)
+                        info.flags,
-                .load_complex(info.itype, info.flags, &add_args, info.offset);
+                        &args,
+                        info.offset,
+                    );
                 }
                 Opcode::Uload8 => {
                     pos.func.dfg.replace(inst).uload8_complex(
                         info.itype,
                         info.flags,
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -240,7 +238,7 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).sload8_complex(
                         info.itype,
                         info.flags,
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -248,7 +246,7 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).uload16_complex(
                         info.itype,
                         info.flags,
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -256,7 +254,7 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).sload16_complex(
                         info.itype,
                         info.flags,
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -264,19 +262,19 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func
                         .dfg
                         .replace(inst)
-                .uload32_complex(info.flags, &add_args, info.offset);
+                        .uload32_complex(info.flags, &args, info.offset);
                 }
                 Opcode::Sload32 => {
                     pos.func
                         .dfg
                         .replace(inst)
-                .sload32_complex(info.flags, &add_args, info.offset);
+                        .sload32_complex(info.flags, &args, info.offset);
                 }
                 Opcode::Store => {
                     pos.func.dfg.replace(inst).store_complex(
                         info.flags,
                         info.st_arg.unwrap(),
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -284,7 +282,7 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).istore8_complex(
                         info.flags,
                         info.st_arg.unwrap(),
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -292,7 +290,7 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).istore16_complex(
                         info.flags,
                         info.st_arg.unwrap(),
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
@@ -300,12 +298,56 @@ fn optimize_complex_addresses(pos: &mut EncCursor, inst: Inst, isa: &TargetIsa)
                     pos.func.dfg.replace(inst).istore32_complex(
                         info.flags,
                         info.st_arg.unwrap(),
-                &add_args,
+                        &args,
                         info.offset,
                     );
                 }
-        _ => return,
+                _ => panic!("Unsupported load or store opcode"),
+            },
+            InstructionData::BinaryImm {
+                opcode: Opcode::IaddImm,
+                arg,
+                imm,
+            } => match pos.func.dfg[inst] {
+                // Operand is an iadd_imm. Fold the immediate into the offset if possible.
+                InstructionData::Load {
+                    arg: ref mut load_arg,
+                    ref mut offset,
+                    ..
+                } => {
+                    if let Some(imm) = offset.try_add_i64(imm.into()) {
+                        *load_arg = arg;
+                        *offset = imm;
+                    } else {
+                        // Overflow.
+                        return;
                     }
+                }
+                InstructionData::Store {
+                    args: ref mut store_args,
+                    ref mut offset,
+                    ..
+                } => {
+                    if let Some(imm) = offset.try_add_i64(imm.into()) {
+                        store_args[0] = arg;
+                        *offset = imm;
+                    } else {
+                        // Overflow.
+                        return;
+                    }
+                }
+                _ => panic!(),
+            },
+            _ => {
+                // Address value is defined by some other kind of instruction.
+                return;
+            }
+        }
+    } else {
+        // Address value is not the result of an instruction.
+        return;
+    }
+
     let ok = pos.func.update_encoding(inst, isa).is_ok();
     debug_assert!(ok);
 }