s390x: Emit 20bit immedate variants of indexed loads

Previously, cranelift did not emit memory operations using
20 bit immedates for reg + reg memory modes. This patch
enables that in a type safe manner by exporting a few
extractors for 20 bit immediates and changing the argument
to the MemArg constructor for reg + reg addressing mode to
accept that 20 bit offset instead of the u8 offset.

Note that the '_bias' suffixed rule is not modified, as it's
primarily used by vector instructions that don't have BXD20
addressing, only BXD12.

Author: theotherjimmy

cranelift/codegen/src/isa/s390x/inst.isle

@@ -1773,6 +1773,14 @@
 (decl i64_from_offset (i64) Offset32)
 (extern extractor infallible i64_from_offset i64_from_offset)
 
+(type SImm20 extern (enum))
+
+(decl pure partial memarg_imm_from_offset (SImm20) Offset32)
+(extern extractor memarg_imm_from_offset memarg_imm_from_offset)
+
+(decl pure partial memarg_imm_from_offset_plus_bias (Offset32 u8) SImm20)
+(extern constructor memarg_imm_from_offset_plus_bias memarg_imm_from_offset_plus_bias)
+
 ;; Accessors for `MemFlags`.
 
 (decl littleendian () MemFlags)
@@ -1788,6 +1796,9 @@
 
 (type MemArg extern (enum))
 
+(decl memarg_reg_plus_reg_plus_off (Reg Reg SImm20 MemFlags) MemArg)
+(extern constructor memarg_reg_plus_reg_plus_off memarg_reg_plus_reg_plus_off)
+
 (decl memarg_reg_plus_reg (Reg Reg u8 MemFlags) MemArg)
 (extern constructor memarg_reg_plus_reg memarg_reg_plus_reg)
 
@@ -1828,8 +1839,8 @@
 (rule (lower_address flags addr @ (value_type (ty_addr64 _)) (i64_from_offset offset))
       (memarg_reg_plus_off addr offset 0 flags))
 
-(rule 1 (lower_address flags (has_type (ty_addr64 _) (iadd x y)) (i64_from_offset 0))
-      (memarg_reg_plus_reg x y 0 flags))
+(rule 1 (lower_address flags (has_type (ty_addr64 _) (iadd x y)) (memarg_imm_from_offset z))
+      (memarg_reg_plus_reg_plus_off x y z flags))
 
 (rule 1 (lower_address flags
                      (symbol_value (symbol_value_data name (RelocDistance.Near) sym_offset))
@@ -1845,8 +1856,9 @@
 (rule (lower_address_bias flags addr @ (value_type $I64) (i64_from_offset offset) bias)
       (memarg_reg_plus_off addr offset bias flags))
 
-(rule 1 (lower_address_bias flags (has_type $I64 (iadd x y)) (i64_from_offset 0) bias)
-      (memarg_reg_plus_reg x y bias flags))
+(rule 1 (lower_address_bias flags (has_type $I64 (iadd x y)) z bias)
+      (if-let final_offset (memarg_imm_from_offset_plus_bias z bias))
+      (memarg_reg_plus_reg_plus_off x y final_offset flags))
 
 
 ;; Test whether a `load` address will be lowered to a `MemArg::Symbol`.

cranelift/codegen/src/isa/s390x/inst/imms.rs

@@ -21,6 +21,15 @@ impl UImm12 {
         }
     }
 
+    pub fn maybe_from_simm20(value: SImm20) -> Option<UImm12> {
+        let SImm20 { value } = value;
+        if value >= 0 {
+            Self::maybe_from_u64(value as u64)
+        } else {
+            None
+        }
+    }
+
     /// Create a zero immediate of this format.
     pub fn zero() -> UImm12 {
         UImm12 { value: 0 }

cranelift/codegen/src/isa/s390x/lower/isle.rs

@@ -8,8 +8,9 @@ use crate::ir::ExternalName;
 use crate::isa::s390x::S390xBackend;
 use crate::isa::s390x::abi::REG_SAVE_AREA_SIZE;
 use crate::isa::s390x::inst::{
-    CallInstDest, Cond, Inst as MInst, LaneOrder, MemArg, RegPair, ReturnCallInfo, SymbolReloc,
-    UImm12, UImm16Shifted, UImm32Shifted, WritableRegPair, gpr, stack_reg, writable_gpr, zero_reg,
+    CallInstDest, Cond, Inst as MInst, LaneOrder, MemArg, RegPair, ReturnCallInfo, SImm20,
+    SymbolReloc, UImm12, UImm16Shifted, UImm32Shifted, WritableRegPair, gpr, stack_reg,
+    writable_gpr, zero_reg,
 };
 use crate::machinst::isle::*;
 use crate::machinst::{CallInfo, MachLabel, Reg, TryCallInfo, non_writable_value_regs};
@@ -679,6 +680,17 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, S390xBackend> {
         MemFlags::trusted()
     }
 
+    #[inline]
+    fn memarg_imm_from_offset(&mut self, imm: Offset32) -> Option<SImm20> {
+        SImm20::maybe_from_i64(i64::from(imm))
+    }
+
+    #[inline]
+    fn memarg_imm_from_offset_plus_bias(&mut self, imm: Offset32, bias: u8) -> Option<SImm20> {
+        let final_offset = i64::from(imm) + bias as i64;
+        SImm20::maybe_from_i64(final_offset)
+    }
+
     #[inline]
     fn memarg_reg_plus_reg(&mut self, x: Reg, y: Reg, bias: u8, flags: MemFlags) -> MemArg {
         MemArg::BXD12 {
@@ -689,6 +701,31 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, S390xBackend> {
         }
     }
 
+    #[inline]
+    fn memarg_reg_plus_reg_plus_off(
+        &mut self,
+        x: Reg,
+        y: Reg,
+        offset: &SImm20,
+        flags: MemFlags,
+    ) -> MemArg {
+        if let Some(imm) = UImm12::maybe_from_simm20(*offset) {
+            MemArg::BXD12 {
+                base: x,
+                index: y,
+                disp: imm,
+                flags,
+            }
+        } else {
+            MemArg::BXD20 {
+                base: x,
+                index: y,
+                disp: *offset,
+                flags,
+            }
+        }
+    }
+
     #[inline]
     fn memarg_reg_plus_off(&mut self, reg: Reg, off: i64, bias: u8, flags: MemFlags) -> MemArg {
         MemArg::reg_plus_off(reg, off + (bias as i64), flags)

cranelift/filetests/filetests/isa/s390x/load.clif

@@ -52,6 +52,43 @@ block0(v0: i64):
 ;   llgc %r2, 0(%r2) ; trap: heap_oob
 ;   br %r14
 
+function %uload8_i64_i64_offset(i64, i64) -> i64 {
+block0(v0: i64, v1: i64):
+  v2 = iadd v0, v1
+  v3 = uload8.i64 v2+100000
+  return v3
+}
+
+; VCode:
+; block0:
+;   llgc %r2, 100000(%r3,%r2)
+;   br %r14
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   llgc %r2, 0x186a0(%r3, %r2) ; trap: heap_oob
+;   br %r14
+
+function %uload8_i64_i64_offset_too_big(i64, i64) -> i64 {
+block0(v0: i64, v1: i64):
+  v2 = iadd v0, v1
+  v3 = uload8.i64 v2+10000000
+  return v3
+}
+
+; VCode:
+; block0:
+;   agr %r2, %r3
+;   lgfi %r1, 10000000 ; llgc %r2, 0(%r1,%r2)
+;   br %r14
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   agr %r2, %r3
+;   lgfi %r1, 0x989680
+;   llgc %r2, 0(%r1, %r2) ; trap: heap_oob
+;   br %r14
+
 function %sload8_i64(i64) -> i64 {
 block0(v0: i64):
   v1 = sload8.i64 v0