Reorganize to put stack finalization in one place

mkeeter · Feb 2, 2025 · 998f99b · 998f99b
1 parent 8297eb1
commit 998f99b
Show file tree

Hide file tree

Showing 10 changed files with 45 additions and 68 deletions.
diff --git a/fidget/src/core/eval/test/float_slice.rs b/fidget/src/core/eval/test/float_slice.rs
@@ -209,7 +209,7 @@ impl<F: Function + MathFunction> TestFloatSlice<F> {
     }
 
     pub fn test_f_stress() {
-        for n in [4, 8, 12, 16, 32] {
+        for n in [4, 8, 12, 16, 32, 256, 512] {
             Self::test_f_stress_n(n);
         }
     }

diff --git a/fidget/src/jit/aarch64/float_slice.rs b/fidget/src/jit/aarch64/float_slice.rs
@@ -454,13 +454,8 @@ impl Assembler for FloatSliceAssembler {
             ; ldr x22, [sp, 0x210]
             ; ldr x23, [sp, 0x218]
             ; ldr x24, [sp, 0x220]
-
-            // Fix up the stack
-            ; add sp, sp, self.0.mem_offset as u32
-            ; ret
         );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/aarch64/grad_slice.rs b/fidget/src/jit/aarch64/grad_slice.rs
@@ -539,14 +539,8 @@ impl Assembler for GradSliceAssembler {
             ; ldr x21, [sp, 0x208]
             ; ldr x22, [sp, 0x210]
             ; ldr x23, [sp, 0x218]
-
-            // Fix up the stack
-            ; add sp, sp, self.0.mem_offset as u32
-            ; ret
-
         );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/aarch64/interval.rs b/fidget/src/jit/aarch64/interval.rs
@@ -772,24 +772,7 @@ impl Assembler for IntervalAssembler {
             ; ldp   d12, d13, [sp, 0x30]
             ; ldp   d14, d15, [sp, 0x40]
         );
-
-        // Fix up the stack
-        if self.0.mem_offset < 4096 {
-            dynasm!(self.0.ops
-                ; add sp, sp, self.0.mem_offset as u32
-            );
-        } else {
-            dynasm!(self.0.ops
-                ; mov w28, self.0.mem_offset as u64
-                ; add sp, sp, w28
-            );
-        }
-
-        dynasm!(self.0.ops
-            ; ret
-        );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/aarch64/point.rs b/fidget/src/jit/aarch64/point.rs
@@ -472,14 +472,8 @@ impl Assembler for PointAssembler {
             ; ldp   d10, d11, [sp, 0x20]
             ; ldp   d12, d13, [sp, 0x30]
             ; ldp   d14, d15, [sp, 0x40]
-
-            // Fix up the stack
-            ; add sp, sp, (self.0.mem_offset as u32)
-
-            ; ret
         );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/mod.rs b/fidget/src/jit/mod.rs
@@ -331,7 +331,7 @@ impl<T> AssemblerData<T> {
 
     #[cfg(target_arch = "aarch64")]
     fn push_stack(&mut self) {
-        if self.mem_offset < 4096 {    
+        if self.mem_offset < 4096 {
             dynasm!(self.ops
                 ; sub sp, sp, self.mem_offset as u32
             );
@@ -340,7 +340,7 @@ impl<T> AssemblerData<T> {
                 ; mov w28, self.mem_offset as u64
                 ; sub sp, sp, w28
             );
-      }
+        }
     }
 
     #[cfg(target_arch = "x86_64")]
@@ -354,6 +354,40 @@ impl<T> AssemblerData<T> {
         assert!(slot >= REGISTER_LIMIT as u32);
         (slot - REGISTER_LIMIT as u32) * std::mem::size_of::<T>() as u32
     }
+
+    #[cfg(target_arch = "aarch64")]
+    fn finalize(mut self) -> Result<Mmap, Error> {
+        // Fix up the stack
+        if self.mem_offset < 4096 {
+            dynasm!(self.ops
+                ; add sp, sp, self.mem_offset as u32
+            );
+        } else if self.mem_offset < 65536 {
+            dynasm!(self.ops
+                ; mov w9, self.mem_offset as u64
+                ; add sp, sp, w9
+            );
+        } else {
+            panic!("invalid mem offset: {}", self.mem_offset);
+        }
+
+        dynasm!(self.ops
+            ; ret
+        );
+        self.ops.finalize()
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    fn finalize(mut self) -> Result<Mmap, Error> {
+        dynasm!(self.ops
+            ; add rsp, self.mem_offset as i32
+            ; pop rbp
+            ; emms
+            ; vzeroall
+            ; ret
+        );
+        self.ops.finalize()
+    }
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/fidget/src/jit/x86_64/float_slice.rs b/fidget/src/jit/x86_64/float_slice.rs
@@ -383,14 +383,8 @@ impl Assembler for FloatSliceAssembler {
 
             // Finalization code, which happens after all evaluation is complete
             ; ->X:
-            ; add rsp, self.0.mem_offset as i32
-            ; pop rbp
-            ; emms
-            ; vzeroall
-            ; ret
         );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/x86_64/grad_slice.rs b/fidget/src/jit/x86_64/grad_slice.rs
@@ -482,13 +482,8 @@ impl Assembler for GradSliceAssembler {
 
             // Finalization code, which happens after all evaluation is complete
             ; -> X:
-            ; add rsp, self.0.mem_offset as i32
-            ; pop rbp
-            ; emms
-            ; ret
         );
-
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/x86_64/interval.rs b/fidget/src/jit/x86_64/interval.rs
@@ -771,13 +771,7 @@ impl Assembler for IntervalAssembler {
                 ; mov r15, [rbp - 0x20]
             );
         }
-        dynasm!(self.0.ops
-            ; add rsp, self.0.mem_offset as i32
-            ; pop rbp
-            ; emms
-            ; ret
-        );
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }
 

diff --git a/fidget/src/jit/x86_64/point.rs b/fidget/src/jit/x86_64/point.rs
@@ -419,13 +419,7 @@ impl Assembler for PointAssembler {
                 ; mov r15, [rbp - 0x20]
             );
         }
-        dynasm!(self.0.ops
-            ; add rsp, self.0.mem_offset as i32
-            ; pop rbp
-            ; emms
-            ; ret
-        );
-        self.0.ops.finalize()
+        self.0.finalize()
     }
 }