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 memarg_imm_from_u16 (SImm20) u16)
+(extern extractor infallible memarg_imm_from_u16 memarg_imm_from_u16)
+
 ;; 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))

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,16 @@ 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_u16(&mut self, imm: u16) -> SImm20 {
+        SImm20::maybe_from_i64(imm as i64).unwrap()
+    }
+
     #[inline]
     fn memarg_reg_plus_reg(&mut self, x: Reg, y: Reg, bias: u8, flags: MemFlags) -> MemArg {
         MemArg::BXD12 {
@@ -689,6 +700,25 @@ 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