Merge pull request #1537 from ltratt/rev_an_call

Teach the reverse analyser about calls
ykjit · Jan 8, 2025 · c06d8b9 · c06d8b9
2 parents 07c3cae + 354b464
commit c06d8b9
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 4 deletions.
diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs b/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
@@ -963,7 +963,7 @@ impl LSRegAlloc<'_> {
         &mut self,
         asm: &mut Assembler,
         iidx: InstIdx,
-        constraints: [RegConstraint<Rx>; N],
+        mut constraints: [RegConstraint<Rx>; N],
     ) -> [Rx; N] {
         // All constraint operands should be float-typed.
         #[cfg(debug_assertions)]
@@ -1031,6 +1031,35 @@ impl LSRegAlloc<'_> {
             }
         }
 
+        // If we have a hint for a constraint, use it.
+        for (i, cnstr) in constraints.iter_mut().enumerate() {
+            match cnstr {
+                RegConstraint::Output
+                | RegConstraint::OutputCanBeSameAsInput(_)
+                | RegConstraint::InputOutput(_) => {
+                    if let Some(reg_alloc::Register::FP(reg)) =
+                        self.rev_an.reg_hints[usize::from(iidx)]
+                    {
+                        if !avoid.is_set(reg) {
+                            *cnstr = match cnstr {
+                                RegConstraint::Output => RegConstraint::OutputFromReg(reg),
+                                RegConstraint::OutputCanBeSameAsInput(_) => {
+                                    RegConstraint::OutputFromReg(reg)
+                                }
+                                RegConstraint::InputOutput(op) => {
+                                    RegConstraint::InputOutputIntoReg(op.clone(), reg)
+                                }
+                                _ => unreachable!(),
+                            };
+                            asgn[i] = Some(reg);
+                            avoid.set(reg);
+                        }
+                    }
+                }
+                _ => (),
+            }
+        }
+
         // If we already have the value in a register, don't assign a new register.
         for (i, cnstr) in constraints.iter().enumerate() {
             match cnstr {

diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/rev_analyse.rs b/ykrt/src/compile/jitc_yk/codegen/x64/rev_analyse.rs
@@ -13,9 +13,12 @@
 //!      to use `dead_code.rs` as well.
 
 use super::reg_alloc::{Register, VarLocation};
-use crate::compile::jitc_yk::jit_ir::{
-    BinOp, BinOpInst, DynPtrAddInst, ICmpInst, Inst, InstIdx, LoadInst, Module, Operand,
-    PtrAddInst, SExtInst, SelectInst, StoreInst, TraceKind, TruncInst, ZExtInst,
+use crate::compile::jitc_yk::{
+    codegen::x64::{ARG_FP_REGS, ARG_GP_REGS},
+    jit_ir::{
+        BinOp, BinOpInst, DirectCallInst, DynPtrAddInst, ICmpInst, Inst, InstIdx, LoadInst, Module,
+        Operand, PtrAddInst, SExtInst, SelectInst, StoreInst, TraceKind, TruncInst, Ty, ZExtInst,
+    },
 };
 use dynasmrt::x64::Rq;
 use vob::Vob;
@@ -149,6 +152,7 @@ impl<'a> RevAnalyse<'a> {
                     // These are handled in [Self::analyse_header] or [Self::analyse_body].
                 }
                 Inst::BinOp(x) => self.an_binop(iidx, x),
+                Inst::Call(x) => self.an_call(iidx, x),
                 Inst::ICmp(x) => self.an_icmp(iidx, x),
                 Inst::PtrAdd(x) => self.an_ptradd(iidx, x),
                 Inst::DynPtrAdd(x) => self.an_dynptradd(iidx, x),
@@ -252,6 +256,28 @@ impl<'a> RevAnalyse<'a> {
         }
     }
 
+    fn an_call(&mut self, _: InstIdx, cinst: DirectCallInst) {
+        let mut gp_regs = ARG_GP_REGS.iter();
+        let mut fp_regs = ARG_FP_REGS.iter();
+        for aidx in cinst.iter_args_idx() {
+            match self.m.type_(self.m.arg(aidx).tyidx(&self.m)) {
+                Ty::Void => unreachable!(),
+                Ty::Integer(_) | Ty::Ptr => {
+                    if let Some(reg) = gp_regs.next() {
+                        self.push_reg_hint_fixed(self.m.arg(aidx), Register::GP(*reg));
+                    }
+                }
+                Ty::Func(_) => todo!(),
+                Ty::Float(_) => {
+                    if let Some(reg) = fp_regs.next() {
+                        self.push_reg_hint_fixed(self.m.arg(aidx), Register::FP(*reg));
+                    }
+                }
+                Ty::Unimplemented(_) => panic!(),
+            }
+        }
+    }
+
     fn an_icmp(&mut self, iidx: InstIdx, icinst: ICmpInst) {
         self.push_reg_hint(iidx, icinst.lhs(self.m));
     }