pulley: superinstructions for pushing/popping list of registers (#9099)

* pulley: add `push` and `pop` instructions

Add `xpush{32, 64}` and `xpop{32, 64}` for pushing/popping XRegs from the stack,
and `push_frame`/`pop_frame` for saving/restoring LR and FP in function
prologue/epilogue.

Copyright (c) 2024, Arm Limited.

Signed-off-by: Karl Meakin <[email protected]>

* cranelift-bitset: more impls for `ScalarBitset`

Implement `Arbitrary` for `ScalarBitset`.

Also implement `DoubleEndedIterator` and `ExactSizeIterator` for `Iter`.
The `pop_min` method was added to help implement `DoubleEndedIterator`,
and `pop` was renamed to `pop_max` for consistency.

Copyright (c) 2024, Arm Limited.

Signed-off-by: Karl Meakin <[email protected]>

* pulley: add instructions for pushing/popping list of registers

Add `xpush{32,64}_many` and `xpop{32,64}_many` for pushing/popping any
combination of all 32 XRegs.

Copyright (c) 2024, Arm Limited.

Signed-off-by: Karl Meakin <[email protected]>

---------

Signed-off-by: Karl Meakin <[email protected]>

cranelift/bitset/Cargo.toml

@@ -13,8 +13,10 @@ rust-version.workspace = true
 workspace = true
 
 [dependencies]
+arbitrary = { workspace = true, optional = true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
 
 [features]
 enable-serde = ["dep:serde", "dep:serde_derive"]
+arbitrary = ["dep:arbitrary"]

cranelift/bitset/src/scalar.rs

@@ -350,6 +350,33 @@ where
         self.0 = T::from(0);
     }
 
+    /// Remove and return the smallest value in the bitset.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use cranelift_bitset::ScalarBitSet;
+    ///
+    /// let mut bitset = ScalarBitSet::<u64>::new();
+    ///
+    /// bitset.insert(0);
+    /// bitset.insert(24);
+    /// bitset.insert(13);
+    /// bitset.insert(36);
+    ///
+    /// assert_eq!(bitset.pop_min(), Some(0));
+    /// assert_eq!(bitset.pop_min(), Some(13));
+    /// assert_eq!(bitset.pop_min(), Some(24));
+    /// assert_eq!(bitset.pop_min(), Some(36));
+    /// assert_eq!(bitset.pop_min(), None);
+    /// ```
+    #[inline]
+    pub fn pop_min(&mut self) -> Option<u8> {
+        let min = self.min()?;
+        self.remove(min);
+        Some(min)
+    }
+
     /// Remove and return the largest value in the bitset.
     ///
     /// # Example
@@ -364,14 +391,14 @@ where
     /// bitset.insert(13);
     /// bitset.insert(36);
     ///
-    /// assert_eq!(bitset.pop(), Some(36));
-    /// assert_eq!(bitset.pop(), Some(24));
-    /// assert_eq!(bitset.pop(), Some(13));
-    /// assert_eq!(bitset.pop(), Some(0));
-    /// assert_eq!(bitset.pop(), None);
+    /// assert_eq!(bitset.pop_max(), Some(36));
+    /// assert_eq!(bitset.pop_max(), Some(24));
+    /// assert_eq!(bitset.pop_max(), Some(13));
+    /// assert_eq!(bitset.pop_max(), Some(0));
+    /// assert_eq!(bitset.pop_max(), None);
     /// ```
     #[inline]
-    pub fn pop(&mut self) -> Option<u8> {
+    pub fn pop_max(&mut self) -> Option<u8> {
         let max = self.max()?;
         self.remove(max);
         Some(max)
@@ -458,11 +485,8 @@ where
     /// );
     /// ```
     #[inline]
-    pub fn iter(&self) -> Iter<T> {
-        Iter {
-            value: self.0,
-            index: 0,
-        }
+    pub fn iter(self) -> Iter<T> {
+        Iter { bitset: self }
     }
 }
 
@@ -494,6 +518,12 @@ where
     }
 }
 
+impl<T: ScalarBitSetStorage> From<T> for ScalarBitSet<T> {
+    fn from(bits: T) -> Self {
+        Self(bits)
+    }
+}
+
 /// A trait implemented by all integers that can be used as the backing storage
 /// for a [`ScalarBitSet`].
 ///
@@ -550,8 +580,7 @@ impl_storage!(usize);
 
 /// An iterator over the elements in a [`ScalarBitSet`].
 pub struct Iter<T> {
-    value: T,
-    index: u8,
+    bitset: ScalarBitSet<T>,
 }
 
 impl<T> Iterator for Iter<T>
@@ -562,14 +591,36 @@ where
 
     #[inline]
     fn next(&mut self) -> Option<u8> {
-        if self.value == T::from(0) {
-            None
-        } else {
-            let trailing_zeros = self.value.trailing_zeros();
-            let elem = self.index + trailing_zeros;
-            self.index += trailing_zeros + 1;
-            self.value = self.value >> (trailing_zeros + 1);
-            Some(elem)
-        }
+        self.bitset.pop_min()
+    }
+}
+
+impl<T> DoubleEndedIterator for Iter<T>
+where
+    T: ScalarBitSetStorage,
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.bitset.pop_max()
+    }
+}
+
+impl<T> ExactSizeIterator for Iter<T>
+where
+    T: ScalarBitSetStorage,
+{
+    #[inline]
+    fn len(&self) -> usize {
+        usize::from(self.bitset.len())
+    }
+}
+
+#[cfg(feature = "arbitrary")]
+impl<'a, T> arbitrary::Arbitrary<'a> for ScalarBitSet<T>
+where
+    T: ScalarBitSetStorage + arbitrary::Arbitrary<'a>,
+{
+    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
+        T::arbitrary(u).map(Self)
     }
 }

cranelift/filetests/filetests/isa/pulley32/get_stack_pointer.clif

@@ -13,6 +13,5 @@ block0:
 ;   ret
 ;
 ; Disassembled:
-;        0: 33 02 00 00                     get_sp x0
+;        0: 3d 02 00 00                     get_sp x0
 ;        4: 00                              ret
-

cranelift/filetests/filetests/isa/pulley32/trap.clif

@@ -11,7 +11,7 @@ block0:
 ;   trap // code = User(0)
 ;
 ; Disassembled:
-;        0: 33 00 00                        trap
+;        0: 3d 00 00                        trap
 
 function %trapnz(i64) {
 block0(v0: i64):
@@ -36,7 +36,7 @@ block0(v0: i64):
 ;        3: 14 03 0c                        xeq64 x3, x0, x3
 ;        6: 03 03 07 00 00 00               br_if x3, 0x7    // target = 0xd
 ;        c: 00                              ret
-;        d: 33 00 00                        trap
+;        d: 3d 00 00                        trap
 
 function %trapz(i64) {
 block0(v0: i64):
@@ -61,5 +61,4 @@ block0(v0: i64):
 ;        3: 14 03 0c                        xeq64 x3, x0, x3
 ;        6: 04 03 07 00 00 00               br_if_not x3, 0x7    // target = 0xd
 ;        c: 00                              ret
-;        d: 33 00 00                        trap
-
+;        d: 3d 00 00                        trap

cranelift/filetests/filetests/isa/pulley64/get_stack_pointer.clif

@@ -13,5 +13,5 @@ block0:
 ;   ret
 ;
 ; Disassembled:
-;        0: 33 02 00 00                     get_sp x0
+;        0: 3d 02 00 00                     get_sp x0
 ;        4: 00                              ret

cranelift/filetests/filetests/isa/pulley64/trap.clif

@@ -11,7 +11,7 @@ block0:
 ;   trap // code = User(0)
 ;
 ; Disassembled:
-;        0: 33 00 00                        trap
+;        0: 3d 00 00                        trap
 
 function %trapnz(i64) {
 block0(v0: i64):
@@ -36,7 +36,7 @@ block0(v0: i64):
 ;        3: 14 03 0c                        xeq64 x3, x0, x3
 ;        6: 03 03 07 00 00 00               br_if x3, 0x7    // target = 0xd
 ;        c: 00                              ret
-;        d: 33 00 00                        trap
+;        d: 3d 00 00                        trap
 
 function %trapz(i64) {
 block0(v0: i64):
@@ -61,5 +61,4 @@ block0(v0: i64):
 ;        3: 14 03 0c                        xeq64 x3, x0, x3
 ;        6: 04 03 07 00 00 00               br_if_not x3, 0x7    // target = 0xd
 ;        c: 00                              ret
-;        d: 33 00 00                        trap
-
+;        d: 3d 00 00                        trap

pulley/Cargo.toml

@@ -14,6 +14,7 @@ workspace = true
 
 [dependencies]
 arbitrary = { workspace = true, optional = true }
+cranelift-bitset = { workspace = true } 
 log = { workspace = true }
 sptr = { workspace = true }
 
@@ -22,7 +23,7 @@ env_logger = { workspace = true }
 
 [features]
 std = []
-arbitrary = ["dep:arbitrary", "arbitrary/derive", "std"]
+arbitrary = ["dep:arbitrary", "arbitrary/derive", "std", "cranelift-bitset/arbitrary"]
 encode = []
 decode = []
 disas = ["decode"]

pulley/README.md

@@ -47,27 +47,16 @@ to change, instructions to appear and disappear, and APIs to be overhauled.
 Here is the disassembly of `f(a, b) = a + b` in Pulley today:
 
 ```
-       0: 0e 1a f0                        xconst8 x26, -16
-       3: 12 7b 6b                        xadd32 sp, sp, x26
-       6: 2c 1b 08 1c                     store64_offset8 sp, 8, lr
-       a: 2a 1b 1d                        store64 sp, fp
-       d: 0b 1d 1b                        xmov fp, sp
-      10: 12 00 04                        xadd32 x0, x0, x1
-      13: 0b 1b 1d                        xmov sp, fp
-      16: 25 1c 1b 08                     load64_offset8 lr, sp, 8
-      1a: 22 1d 1b                        load64 fp, sp
-      1d: 0e 1a 10                        xconst8 x26, 16
-      20: 12 7b 6b                        xadd32 sp, sp, x26
-      23: 00                              ret
+       0: 2f                              push_frame
+       1: 12 00 04                        xadd32 x0, x0, x1
+       4: 30                              pop_frame
+       5: 00                              ret
 ```
 
 Note that there are a number of things that could be improved here:
 
 * We could avoid allocating and deallocating a stack frame because this function's
   body doesn't use any stack slots.
-* We could collapse the whole prologue and epilogue instruction sequences into
-  super-instructions, since they are identical (modulo the frame size immediate)
-  for all functions.
 
 As mentioned above, Pulley is very much a work in progress.
 

pulley/fuzz/src/interp.rs

@@ -105,6 +105,11 @@ fn op_is_safe_for_fuzzing(op: &Op) -> bool {
         | Op::Xslteq32(Xslteq32 { operands, .. })
         | Op::Xult32(Xult32 { operands, .. })
         | Op::Xulteq32(Xulteq32 { operands, .. }) => !operands.dst.is_special(),
+        Op::PushFrame(_) | Op::PopFrame(_) => false,
+        Op::XPush32(_) | Op::XPush64(_) => false,
+        Op::XPop32(_) | Op::XPop64(_) => false,
+        Op::XPush32Many(_) | Op::XPush64Many(_) => false,
+        Op::XPop32Many(_) | Op::XPop64Many(_) => false,
     }
 }
 

pulley/src/decode.rs

@@ -1,6 +1,8 @@
 //! Decoding support for pulley bytecode.
 
 use alloc::vec::Vec;
+use cranelift_bitset::scalar::ScalarBitSetStorage;
+use cranelift_bitset::ScalarBitSet;
 
 use crate::imms::*;
 use crate::opcode::*;
@@ -384,6 +386,24 @@ impl<R: Reg> Decode for BinaryOperands<R> {
     }
 }
 
+impl<S: Decode + ScalarBitSetStorage> Decode for ScalarBitSet<S> {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        S::decode(bytecode).map(ScalarBitSet::from)
+    }
+}
+
+impl<R: Reg + Decode> Decode for RegSet<R> {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        ScalarBitSet::decode(bytecode).map(Self::from)
+    }
+}
+
 /// A Pulley bytecode decoder.
 ///
 /// Does not materialize bytecode instructions, instead all decoding methods are

pulley/src/disas.rs

@@ -127,13 +127,26 @@ impl Disas for PcRelOffset {
     }
 }
 
-impl<R: Disas> Disas for BinaryOperands<R> {
-    fn disas(&self, position: usize, disas: &mut String) {
-        self.dst.disas(position, disas);
-        write!(disas, ", ").unwrap();
-        self.src1.disas(position, disas);
+fn disas_list<T: Disas>(position: usize, disas: &mut String, iter: impl IntoIterator<Item = T>) {
+    let mut iter = iter.into_iter();
+    let Some(first) = iter.next() else { return };
+    first.disas(position, disas);
+
+    for item in iter {
         write!(disas, ", ").unwrap();
-        self.src2.disas(position, disas);
+        item.disas(position, disas);
+    }
+}
+
+impl<R: Reg + Disas> Disas for BinaryOperands<R> {
+    fn disas(&self, position: usize, disas: &mut String) {
+        disas_list(position, disas, [self.dst, self.src1, self.src2])
+    }
+}
+
+impl<R: Reg + Disas> Disas for RegSet<R> {
+    fn disas(&self, position: usize, disas: &mut String) {
+        disas_list(position, disas, *self)
     }
 }