Remove sig data arg and ret fields to reduce size (#5319)

* Remove sig data arg and ret fields to reduce size

* Update cranelift/codegen/src/machinst/abi.rs

Co-authored-by: Jamey Sharp <jamey@minilop.net>

* Update cranelift/codegen/src/machinst/abi.rs

Co-authored-by: Jamey Sharp <jamey@minilop.net>

* Fix offsets

* Add comment

Co-authored-by: Jamey Sharp <jamey@minilop.net>

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

@@ -126,7 +126,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
     fn defs_init(&mut self, abi: &Sig) -> CallRetList {
         // Allocate writable registers for all retval regs, except for StructRet args.
         let mut defs = smallvec![];
-        for i in 0..self.lower_ctx.sigs()[*abi].num_rets() {
+        for i in 0..self.lower_ctx.sigs().num_rets(*abi) {
             if let &ABIArg::Slots {
                 ref slots, purpose, ..
             } = &self.lower_ctx.sigs().get_ret(*abi, i)
@@ -169,7 +169,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         // Return the index of the first actual return value, excluding
         // any StructReturn that might have been added to Sig.
         let sig = &self.lower_ctx.dfg().signatures[sig_ref];
-        self.lower_ctx.sigs()[*abi].num_rets() - sig.returns.len()
+        self.lower_ctx.sigs().num_rets(*abi) - sig.returns.len()
     }
 
     fn abi_lane_order(&mut self, abi: &Sig) -> LaneOrder {

cranelift/codegen/src/machinst/abi.rs

@@ -608,26 +608,42 @@ pub trait ABIMachineSpec {
 pub struct Sig(u32);
 cranelift_entity::entity_impl!(Sig);
 
+impl Sig {
+    fn prev(self) -> Option<Sig> {
+        self.0.checked_sub(1).map(Sig)
+    }
+}
+
 /// ABI information shared between body (callee) and caller.
 #[derive(Clone, Debug)]
 pub struct SigData {
-    /// Argument location starting offset (regs or stack slots). Stack offsets are relative to
+    /// Currently both return values and arguments are stored in a continuous space vector
+    /// in `SigSet::abi_args`.
+    ///
+    /// ```plain
+    ///                  +----------------------------------------------+
+    ///                  | return values                                |
+    ///                  | ...                                          |
+    ///   rets_end   --> +----------------------------------------------+
+    ///                  | arguments                                    |
+    ///                  | ...                                          |
+    ///   args_end   --> +----------------------------------------------+
+    ///
+    /// ```
+    ///
+    /// Note we only store two offsets as rets_end == args_start, and rets_start == prev.args_end.
+    ///
+    /// Argument location ending offset (regs or stack slots). Stack offsets are relative to
     /// SP on entry to function.
     ///
     /// This is a index into the `SigSet::abi_args`.
-    args_start: u32,
+    args_end: u32,
 
-    /// Amount of arguments in the stack.
+    /// Return-value location ending offset. Stack offsets are relative to the return-area
-    args_length: u16,
-
-    /// Return-value location starting offset. Stack offsets are relative to the return-area
     /// pointer.
     ///
     /// This is a index into the `SigSet::abi_args`.
-    rets_start: u32,
+    rets_end: u32,
-
-    /// Amount of Return value locations in the stack.
-    rets_length: u16,
 
     /// Space on stack used to store arguments.
     sized_stack_arg_space: i64,
@@ -643,27 +659,11 @@ pub struct SigData {
 }
 
 impl SigData {
-    /// Get the number of arguments expected.
-    pub fn num_args(&self) -> usize {
-        let num = if self.stack_ret_arg.is_some() {
-            self.args_length - 1
-        } else {
-            self.args_length
-        };
-
-        usize::from(num)
-    }
-
     /// Get total stack space required for arguments.
     pub fn sized_stack_arg_space(&self) -> i64 {
         self.sized_stack_arg_space
     }
 
-    /// Get the number of return values expected.
-    pub fn num_rets(&self) -> usize {
-        usize::from(self.rets_length)
-    }
-
     /// Get total stack space required for return values.
     pub fn sized_stack_ret_space(&self) -> i64 {
         self.sized_stack_ret_space
@@ -806,7 +806,10 @@ impl SigSet {
 
         // Compute args and retvals from signature. Handle retvals first,
         // because we may need to add a return-area arg to the args.
-        let rets_start = u32::try_from(self.abi_args.len()).unwrap();
+
+        // NOTE: We rely on the order of the args (rets -> args) inserted to compute the offsets in
+        // `SigSet::args()` and `SigSet::rets()`. Therefore, we cannot change the two
+        // compute_arg_locs order.
         let (sized_stack_ret_space, _) = M::compute_arg_locs(
             sig.call_conv,
             flags,
@@ -818,7 +821,6 @@ impl SigSet {
         let rets_end = u32::try_from(self.abi_args.len()).unwrap();
 
         let need_stack_return_area = sized_stack_ret_space > 0;
-        let args_start = u32::try_from(self.abi_args.len()).unwrap();
         let (sized_stack_arg_space, stack_ret_arg) = M::compute_arg_locs(
             sig.call_conv,
             flags,
@@ -829,17 +831,12 @@ impl SigSet {
         )?;
         let args_end = u32::try_from(self.abi_args.len()).unwrap();
 
-        let args_length = u16::try_from(args_end - args_start).unwrap();
-        let rets_length = u16::try_from(rets_end - rets_start).unwrap();
-
         trace!(
-            "ABISig: sig {:?} => args start = {} args len = {} rets start = {} rets len = {}
+            "ABISig: sig {:?} => args end = {} rets end = {}
              arg stack = {} ret stack = {} stack_ret_arg = {:?}",
             sig,
-            args_start,
+            args_end,
-            args_length,
+            rets_end,
-            rets_start,
-            rets_length,
             sized_stack_arg_space,
             sized_stack_ret_space,
             need_stack_return_area,
@@ -847,10 +844,8 @@ impl SigSet {
 
         let stack_ret_arg = stack_ret_arg.map(|s| u16::try_from(s).unwrap());
         Ok(SigData {
-            args_start,
+            args_end,
-            args_length,
+            rets_end,
-            rets_start,
-            rets_length,
             sized_stack_arg_space,
             sized_stack_ret_space,
             stack_ret_arg,
@@ -861,8 +856,9 @@ impl SigSet {
     /// Get this signature's ABI arguments.
     pub fn args(&self, sig: Sig) -> &[ABIArg] {
         let sig_data = &self.sigs[sig];
-        let start = usize::try_from(sig_data.args_start).unwrap();
+        // Please see comments in `SigSet::from_func_sig` of how we store the offsets.
-        let end = usize::try_from(sig_data.args_start + u32::from(sig_data.args_length)).unwrap();
+        let start = usize::try_from(sig_data.rets_end).unwrap();
+        let end = usize::try_from(sig_data.args_end).unwrap();
         &self.abi_args[start..end]
     }
 
@@ -885,8 +881,9 @@ impl SigSet {
     /// Get this signature's ABI returns.
     pub fn rets(&self, sig: Sig) -> &[ABIArg] {
         let sig_data = &self.sigs[sig];
-        let start = usize::try_from(sig_data.rets_start).unwrap();
+        // Please see comments in `SigSet::from_func_sig` of how we store the offsets.
-        let end = usize::try_from(sig_data.rets_start + u32::from(sig_data.rets_length)).unwrap();
+        let start = usize::try_from(sig.prev().map_or(0, |prev| self.sigs[prev].args_end)).unwrap();
+        let end = usize::try_from(sig_data.rets_end).unwrap();
         &self.abi_args[start..end]
     }
 
@@ -927,6 +924,21 @@ impl SigSet {
 
         clobbers
     }
+
+    /// Get the number of arguments expected.
+    pub fn num_args(&self, sig: Sig) -> usize {
+        let len = self.args(sig).len();
+        if self.sigs[sig].stack_ret_arg.is_some() {
+            len - 1
+        } else {
+            len
+        }
+    }
+
+    /// Get the number of return values expected.
+    pub fn num_rets(&self, sig: Sig) -> usize {
+        self.rets(sig).len()
+    }
 }
 
 // NB: we do _not_ implement `IndexMut` because these signatures are
@@ -2097,12 +2109,7 @@ fn adjust_stack_and_nominal_sp<M: ABIMachineSpec>(ctx: &mut Lower<M::I>, off: i3
 impl<M: ABIMachineSpec> Caller<M> {
     /// Get the number of arguments expected.
     pub fn num_args(&self, sigs: &SigSet) -> usize {
-        let len = sigs.args(self.sig).len();
+        sigs.num_args(self.sig)
-        if sigs[self.sig].stack_ret_arg.is_some() {
-            len - 1
-        } else {
-            len
-        }
     }
 
     /// Emit code to pre-adjust the stack, prior to argument copies and call.
@@ -2402,6 +2409,6 @@ mod tests {
     fn sig_data_size() {
         // The size of `SigData` is performance sensitive, so make sure
         // we don't regress it unintentionally.
-        assert_eq!(std::mem::size_of::<SigData>(), 40);
+        assert_eq!(std::mem::size_of::<SigData>(), 32);
     }
 }

cranelift/codegen/src/machinst/isle.rs

@@ -432,7 +432,7 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_num_args(&mut self, abi: &Sig) -> usize {
-            self.lower_ctx.sigs()[*abi].num_args()
+            self.lower_ctx.sigs().num_args(*abi)
         }
 
         fn abi_get_arg(&mut self, abi: &Sig, idx: usize) -> ABIArg {
@@ -440,7 +440,7 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_num_rets(&mut self, abi: &Sig) -> usize {
-            self.lower_ctx.sigs()[*abi].num_rets()
+            self.lower_ctx.sigs().num_rets(*abi)
         }
 
         fn abi_get_ret(&mut self, abi: &Sig, idx: usize) -> ABIArg {
@@ -681,7 +681,7 @@ macro_rules! isle_prelude_method_helpers {
         ) -> InstOutput {
             caller.emit_stack_pre_adjust(self.lower_ctx);
 
-            let num_args = self.lower_ctx.sigs()[abi].num_args();
+            let num_args = self.lower_ctx.sigs().num_args(abi);
 
             assert_eq!(
                 inputs.len(&self.lower_ctx.dfg().value_lists) - off,
@@ -710,9 +710,9 @@ macro_rules! isle_prelude_method_helpers {
             let mut retval_insts: crate::machinst::abi::SmallInstVec<_> = smallvec::smallvec![];
             // We take the *last* `num_rets` returns of the sig:
             // this skips a StructReturn, if any, that is present.
-            let sigdata = &self.lower_ctx.sigs()[abi];
+            let sigdata_num_rets = self.lower_ctx.sigs().num_rets(abi);
-            debug_assert!(num_rets <= sigdata.num_rets());
+            debug_assert!(num_rets <= sigdata_num_rets);
-            for i in (sigdata.num_rets() - num_rets)..sigdata.num_rets() {
+            for i in (sigdata_num_rets - num_rets)..sigdata_num_rets {
                 // Borrow `sigdata` again so we don't hold a `self`
                 // borrow across the `&mut self` arg to
                 // `abi_arg_slot_regs()` below.