From 93db01d38aa2029cafd5e5ccda109cc05eca4282 Mon Sep 17 00:00:00 2001
From: Jakob Stoklund Olesen <jolesen@mozilla.com>
Date: Wed, 26 Apr 2017 10:47:43 -0700
Subject: [PATCH] Assign an affinity to function argument values.

Use the argument types from the function signature to initialize the
affinity of register and stack arguments.
---
 lib/cretonne/src/regalloc/affinity.rs | 12 +++++++++++-
 lib/cretonne/src/regalloc/liveness.rs | 19 +++++++++++++------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/lib/cretonne/src/regalloc/affinity.rs b/lib/cretonne/src/regalloc/affinity.rs
index 529828759e..f56b96d434 100644
--- a/lib/cretonne/src/regalloc/affinity.rs
+++ b/lib/cretonne/src/regalloc/affinity.rs
@@ -8,7 +8,8 @@
 //! subclass. This is just a hint, and the register allocator is allowed to pick a register from a
 //! larger register class instead.
 
-use isa::{RegInfo, RegClassIndex, OperandConstraint, ConstraintKind};
+use ir::{ArgumentType, ArgumentLoc};
+use isa::{TargetIsa, RegInfo, RegClassIndex, OperandConstraint, ConstraintKind};
 
 /// Preferred register allocation for an SSA value.
 #[derive(Clone, Copy)]
@@ -42,6 +43,15 @@ impl Affinity {
         }
     }
 
+    /// Create an affinity that matches an ABI argument for `isa`.
+    pub fn abi(arg: &ArgumentType, isa: &TargetIsa) -> Affinity {
+        match arg.location {
+            ArgumentLoc::Unassigned => Affinity::Any,
+            ArgumentLoc::Reg(_) => Affinity::Reg(isa.regclass_for_abi_type(arg.value_type).into()),
+            ArgumentLoc::Stack(_) => Affinity::Stack,
+        }
+    }
+
     /// Merge an operand constraint into this affinity.
     ///
     /// Note that this does not guarantee that the register allocator will pick a register that
diff --git a/lib/cretonne/src/regalloc/liveness.rs b/lib/cretonne/src/regalloc/liveness.rs
index cb937d67a2..737470b2d1 100644
--- a/lib/cretonne/src/regalloc/liveness.rs
+++ b/lib/cretonne/src/regalloc/liveness.rs
@@ -190,6 +190,7 @@ type LiveRangeSet = SparseMap<Value, LiveRange>;
 /// Create it if necessary.
 fn get_or_create<'a>(lrset: &'a mut LiveRangeSet,
                      value: Value,
+                     isa: &TargetIsa,
                      func: &Function,
                      enc_info: &EncInfo)
                      -> &'a mut LiveRange {
@@ -211,11 +212,17 @@ fn get_or_create<'a>(lrset: &'a mut LiveRangeSet,
                     .map(Affinity::new)
                     .unwrap_or_default();
             }
-            ValueDef::Arg(ebb, _) => {
+            ValueDef::Arg(ebb, num) => {
                 def = ebb.into();
-                // Don't apply any affinity to EBB arguments.
-                // They could be in a register or on the stack.
-                affinity = Default::default();
+                if func.layout.entry_block() == Some(ebb) {
+                    // The affinity for entry block arguments can be inferred from the function
+                    // signature.
+                    affinity = Affinity::abi(&func.signature.argument_types[num], isa);
+                } else {
+                    // Don't apply any affinity to normal EBB arguments.
+                    // They could be in a register or on the stack.
+                    affinity = Default::default();
+                }
             }
         };
         lrset.insert(LiveRange::new(value, def, affinity));
@@ -309,7 +316,7 @@ impl Liveness {
                 // TODO: When we implement DCE, we can use the absence of a live range to indicate
                 // an unused value.
                 for &def in func.dfg.inst_results(inst) {
-                    get_or_create(&mut self.ranges, def, func, &enc_info);
+                    get_or_create(&mut self.ranges, def, isa, func, &enc_info);
                 }
 
                 // Iterator of constraints, one per value operand.
@@ -322,7 +329,7 @@ impl Liveness {
 
                 for &arg in func.dfg.inst_args(inst) {
                     // Get the live range, create it as a dead range if necessary.
-                    let lr = get_or_create(&mut self.ranges, arg, func, &enc_info);
+                    let lr = get_or_create(&mut self.ranges, arg, isa, func, &enc_info);
 
                     // Extend the live range to reach this use.
                     extend_to_use(lr, ebb, inst, &mut self.worklist, func, cfg);