Add input register shuffling to test calling convention

mkeeter · May 19, 2024 · f96e9c2 · f96e9c2
1 parent 3deb8d1
commit f96e9c2
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 0 deletions.
diff --git a/fidget/src/jit/aarch64/float_slice.rs b/fidget/src/jit/aarch64/float_slice.rs
@@ -88,6 +88,10 @@ impl Assembler for FloatSliceAssembler {
     fn init(mmap: Mmap, slot_count: usize) -> Self {
         let mut out = AssemblerData::new(mmap);
         out.prepare_stack(slot_count, STACK_SIZE as usize);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             // Preserve frame and link register, and set up the frame pointer
             ; stp   x29, x30, [sp, 0x0]

diff --git a/fidget/src/jit/aarch64/grad_slice.rs b/fidget/src/jit/aarch64/grad_slice.rs
@@ -92,6 +92,10 @@ impl Assembler for GradSliceAssembler {
     fn init(mmap: Mmap, slot_count: usize) -> Self {
         let mut out = AssemblerData::new(mmap);
         out.prepare_stack(slot_count, STACK_SIZE as usize);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             // Preserve frame and link register, and set up the frame pointer
             ; stp   x29, x30, [sp, 0x0]

diff --git a/fidget/src/jit/aarch64/interval.rs b/fidget/src/jit/aarch64/interval.rs
@@ -87,6 +87,10 @@ impl Assembler for IntervalAssembler {
     fn init(mmap: Mmap, slot_count: usize) -> Self {
         let mut out = AssemblerData::new(mmap);
         out.prepare_stack(slot_count, STACK_SIZE as usize);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             // Preserve frame and link register, and set up the frame pointer
             ; stp   x29, x30, [sp, 0x0]

diff --git a/fidget/src/jit/aarch64/point.rs b/fidget/src/jit/aarch64/point.rs
@@ -81,6 +81,10 @@ impl Assembler for PointAssembler {
     fn init(mmap: Mmap, slot_count: usize) -> Self {
         let mut out = AssemblerData::new(mmap);
         out.prepare_stack(slot_count, STACK_SIZE as usize);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             // Preserve frame and link register, and set up the frame pointer
             ; stp   x29, x30, [sp, 0x0]

diff --git a/fidget/src/jit/mod.rs b/fidget/src/jit/mod.rs
@@ -322,6 +322,27 @@ impl<T> AssemblerData<T> {
         self.push_stack();
     }
 
+    #[cfg(all(target_arch = "aarch64", test))]
+    fn input_register_shenanigans(&mut self) {
+        dynasm!(self.ops
+            ; mov x3, 0xABCD
+            ; lsl x4, x3, 16
+            ; orr x3, x3, x4
+        );
+    }
+
+    #[cfg(all(target_arch = "x86_64", test))]
+    fn input_register_shenanigans(&mut self) {
+        dynasm!(self.ops
+            ; mov rcx, 0xABCDABCDu32 as i32
+            ; mov r8, 0xABCDABCDu32 as i32
+            ; mov r9, 0xABCDABCDu32 as i32
+        );
+        for i in 4..=7 {
+            dynasm!(self.ops ; mov X(i), x3);
+        }
+    }
+
     #[cfg(target_arch = "aarch64")]
     fn push_stack(&mut self) {
         assert!(self.mem_offset < 4096);

diff --git a/fidget/src/jit/x86_64/float_slice.rs b/fidget/src/jit/x86_64/float_slice.rs
@@ -66,6 +66,10 @@ impl Assembler for FloatSliceAssembler {
             ; mov rbp, rsp
         );
         out.prepare_stack(slot_count, STACK_SIZE_UPPER + STACK_SIZE_LOWER);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             // TODO should there be a `vzeroupper` in here?
 

diff --git a/fidget/src/jit/x86_64/grad_slice.rs b/fidget/src/jit/x86_64/grad_slice.rs
@@ -60,6 +60,10 @@ impl Assembler for GradSliceAssembler {
             ; mov rbp, rsp
         );
         out.prepare_stack(slot_count, STACK_SIZE_UPPER + STACK_SIZE_LOWER);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             ; xor rcx, rcx // set the array offset (rcx) to 0
 

diff --git a/fidget/src/jit/x86_64/interval.rs b/fidget/src/jit/x86_64/interval.rs
@@ -58,6 +58,10 @@ impl Assembler for IntervalAssembler {
             ; mov rbp, rsp
         );
         out.prepare_stack(slot_count, STACK_SIZE_UPPER + STACK_SIZE_LOWER);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             ; vzeroupper
         );

diff --git a/fidget/src/jit/x86_64/point.rs b/fidget/src/jit/x86_64/point.rs
@@ -57,6 +57,10 @@ impl Assembler for PointAssembler {
             ; mov rbp, rsp
         );
         out.prepare_stack(slot_count, STACK_SIZE_UPPER + STACK_SIZE_LOWER);
+
+        #[cfg(test)]
+        out.input_register_shenanigans();
+
         dynasm!(out.ops
             ; vzeroupper
         );