Merge "Setup policies for register allocation."

6 files changed, 496 insertions, 107 deletions

diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index babb1f5628..ff316e5b04 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -55,9 +55,105 @@ void CodeGenerator::CompileBlock(HBasicBlock* block) {
   }
 }
 
+size_t CodeGenerator::AllocateFreeRegisterInternal(
+    bool* blocked_registers, size_t number_of_registers) const {
+  for (size_t regno = 0; regno < number_of_registers; regno++) {
+    if (!blocked_registers[regno]) {
+      blocked_registers[regno] = true;
+      return regno;
+    }
+  }
+  LOG(FATAL) << "Unreachable";
+  return -1;
+}
+
+
+void CodeGenerator::AllocateRegistersLocally(HInstruction* instruction) const {
+  LocationSummary* locations = instruction->GetLocations();
+  if (locations == nullptr) return;
+
+  for (size_t i = 0, e = GetNumberOfRegisters(); i < e; ++i) {
+    blocked_registers_[i] = false;
+  }
+
+  // Mark all fixed input, temp and output registers as used.
+  for (size_t i = 0, e = locations->GetInputCount(); i < e; ++i) {
+    Location loc = locations->InAt(i);
+    if (loc.IsRegister()) {
+      // Check that a register is not specified twice in the summary.
+      DCHECK(!blocked_registers_[loc.GetEncoding()]);
+      blocked_registers_[loc.GetEncoding()] = true;
+    }
+  }
+
+  for (size_t i = 0, e = locations->GetTempCount(); i < e; ++i) {
+    Location loc = locations->GetTemp(i);
+    if (loc.IsRegister()) {
+      // Check that a register is not specified twice in the summary.
+      DCHECK(!blocked_registers_[loc.GetEncoding()]);
+      blocked_registers_[loc.GetEncoding()] = true;
+    }
+  }
+
+  SetupBlockedRegisters(blocked_registers_);
+
+  // Allocate all unallocated input locations.
+  for (size_t i = 0, e = locations->GetInputCount(); i < e; ++i) {
+    Location loc = locations->InAt(i);
+    HInstruction* input = instruction->InputAt(i);
+    if (loc.IsUnallocated()) {
+      if (loc.GetPolicy() == Location::kRequiresRegister) {
+        loc = Location::RegisterLocation(
+            AllocateFreeRegister(input->GetType(), blocked_registers_));
+      } else {
+        DCHECK_EQ(loc.GetPolicy(), Location::kAny);
+        HLoadLocal* load = input->AsLoadLocal();
+        if (load != nullptr) {
+          loc = GetStackLocation(load);
+        } else {
+          loc = Location::RegisterLocation(
+              AllocateFreeRegister(input->GetType(), blocked_registers_));
+        }
+      }
+      locations->SetInAt(i, loc);
+    }
+  }
+
+  // Allocate all unallocated temp locations.
+  for (size_t i = 0, e = locations->GetTempCount(); i < e; ++i) {
+    Location loc = locations->GetTemp(i);
+    if (loc.IsUnallocated()) {
+      DCHECK_EQ(loc.GetPolicy(), Location::kRequiresRegister);
+      // TODO: Adjust handling of temps. We currently consider temps to use
+      // core registers. They may also use floating point registers at some point.
+      loc = Location::RegisterLocation(static_cast<ManagedRegister>(
+          AllocateFreeRegister(Primitive::kPrimInt, blocked_registers_)));
+      locations->SetTempAt(i, loc);
+    }
+  }
+
+  Location result_location = locations->Out();
+  if (result_location.IsUnallocated()) {
+    switch (result_location.GetPolicy()) {
+      case Location::kAny:
+      case Location::kRequiresRegister:
+        result_location = Location::RegisterLocation(
+            AllocateFreeRegister(instruction->GetType(), blocked_registers_));
+        break;
+      case Location::kSameAsFirstInput:
+        result_location = locations->InAt(0);
+        break;
+    }
+    locations->SetOut(result_location);
+  }
+}
+
 void CodeGenerator::InitLocations(HInstruction* instruction) {
-  if (instruction->GetLocations() == nullptr) return;
-  for (size_t i = 0; i < instruction->InputCount(); i++) {
+  if (instruction->GetLocations() == nullptr) {
+    return;
+  }
+  AllocateRegistersLocally(instruction);
+  for (size_t i = 0, e = instruction->InputCount(); i < e; ++i) {
     Location location = instruction->GetLocations()->InAt(i);
     if (location.IsValid()) {
       // Move the input to the desired location.
diff --git a/compiler/optimizing/code_generator.h b/compiler/optimizing/code_generator.h
index 54f9e70695..74cbccc4b8 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -62,6 +62,12 @@ class Location : public ValueObject {
     // bits are in a stack slot. The kQuickParameter kind is for
     // handling this special case.
     kQuickParameter = 4,
+
+    // Unallocated location represents a location that is not fixed and can be
+    // allocated by a register allocator.  Each unallocated location has
+    // a policy that specifies what kind of location is suitable. Payload
+    // contains register allocation policy.
+    kUnallocated = 5,
   };
 
   Location() : value_(kInvalid) {
@@ -166,10 +172,50 @@ class Location : public ValueObject {
       case kStackSlot: return "S";
       case kDoubleStackSlot: return "DS";
       case kQuickParameter: return "Q";
+      case kUnallocated: return "U";
     }
     return "?";
   }
 
+  // Unallocated locations.
+  enum Policy {
+    kAny,
+    kRequiresRegister,
+    kSameAsFirstInput,
+  };
+
+  bool IsUnallocated() const {
+    return GetKind() == kUnallocated;
+  }
+
+  static Location UnallocatedLocation(Policy policy) {
+    return Location(kUnallocated, PolicyField::Encode(policy));
+  }
+
+  // Any free register is suitable to replace this unallocated location.
+  static Location Any() {
+    return UnallocatedLocation(kAny);
+  }
+
+  static Location RequiresRegister() {
+    return UnallocatedLocation(kRequiresRegister);
+  }
+
+  // The location of the first input to the instruction will be
+  // used to replace this unallocated location.
+  static Location SameAsFirstInput() {
+    return UnallocatedLocation(kSameAsFirstInput);
+  }
+
+  Policy GetPolicy() const {
+    DCHECK(IsUnallocated());
+    return PolicyField::Decode(GetPayload());
+  }
+
+  uword GetEncoding() const {
+    return GetPayload();
+  }
+
  private:
   // Number of bits required to encode Kind value.
   static constexpr uint32_t kBitsForKind = 4;
@@ -187,6 +233,9 @@ class Location : public ValueObject {
   typedef BitField<Kind, 0, kBitsForKind> KindField;
   typedef BitField<uword, kBitsForKind, kBitsForPayload> PayloadField;
 
+  // Layout for kUnallocated locations payload.
+  typedef BitField<Policy, 0, 3> PolicyField;
+
   // Layout for stack slots.
   static const intptr_t kStackIndexBias =
       static_cast<intptr_t>(1) << (kBitsForPayload - 1);
@@ -208,40 +257,52 @@ class Location : public ValueObject {
 class LocationSummary : public ArenaObject {
  public:
   explicit LocationSummary(HInstruction* instruction)
-      : inputs(instruction->GetBlock()->GetGraph()->GetArena(), instruction->InputCount()),
-        temps(instruction->GetBlock()->GetGraph()->GetArena(), 0) {
-    inputs.SetSize(instruction->InputCount());
+      : inputs_(instruction->GetBlock()->GetGraph()->GetArena(), instruction->InputCount()),
+        temps_(instruction->GetBlock()->GetGraph()->GetArena(), 0) {
+    inputs_.SetSize(instruction->InputCount());
     for (size_t i = 0; i < instruction->InputCount(); i++) {
-      inputs.Put(i, Location());
+      inputs_.Put(i, Location());
     }
   }
 
   void SetInAt(uint32_t at, Location location) {
-    inputs.Put(at, location);
+    inputs_.Put(at, location);
   }
 
   Location InAt(uint32_t at) const {
-    return inputs.Get(at);
+    return inputs_.Get(at);
+  }
+
+  size_t GetInputCount() const {
+    return inputs_.Size();
   }
 
   void SetOut(Location location) {
-    output = Location(location);
+    output_ = Location(location);
   }
 
   void AddTemp(Location location) {
-    temps.Add(location);
+    temps_.Add(location);
   }
 
   Location GetTemp(uint32_t at) const {
-    return temps.Get(at);
+    return temps_.Get(at);
+  }
+
+  void SetTempAt(uint32_t at, Location location) {
+    temps_.Put(at, location);
   }
 
-  Location Out() const { return output; }
+  size_t GetTempCount() const {
+    return temps_.Size();
+  }
+
+  Location Out() const { return output_; }
 
  private:
-  GrowableArray<Location> inputs;
-  GrowableArray<Location> temps;
-  Location output;
+  GrowableArray<Location> inputs_;
+  GrowableArray<Location> temps_;
+  Location output_;
 
   DISALLOW_COPY_AND_ASSIGN(LocationSummary);
 };
@@ -286,15 +347,33 @@ class CodeGenerator : public ArenaObject {
       std::vector<uint8_t>* vector, const DexCompilationUnit& dex_compilation_unit) const;
 
  protected:
-  explicit CodeGenerator(HGraph* graph)
+  CodeGenerator(HGraph* graph, size_t number_of_registers)
       : frame_size_(0),
         graph_(graph),
         block_labels_(graph->GetArena(), 0),
-        pc_infos_(graph->GetArena(), 32) {
+        pc_infos_(graph->GetArena(), 32),
+        blocked_registers_(static_cast<bool*>(
+            graph->GetArena()->Alloc(number_of_registers * sizeof(bool), kArenaAllocData))) {
     block_labels_.SetSize(graph->GetBlocks()->Size());
   }
   ~CodeGenerator() { }
 
+  // Register allocation logic.
+  void AllocateRegistersLocally(HInstruction* instruction) const;
+
+  // Backend specific implementation for allocating a register.
+  virtual ManagedRegister AllocateFreeRegister(Primitive::Type type,
+                                               bool* blocked_registers) const = 0;
+
+  // Raw implementation of allocating a register: loops over blocked_registers to find
+  // the first available register.
+  size_t AllocateFreeRegisterInternal(bool* blocked_registers, size_t number_of_registers) const;
+
+  virtual void SetupBlockedRegisters(bool* blocked_registers) const = 0;
+  virtual size_t GetNumberOfRegisters() const = 0;
+
+  virtual Location GetStackLocation(HLoadLocal* load) const = 0;
+
   // Frame size required for this method.
   uint32_t frame_size_;
   uint32_t core_spill_mask_;
@@ -309,6 +388,9 @@ class CodeGenerator : public ArenaObject {
   GrowableArray<Label> block_labels_;
   GrowableArray<PcInfo> pc_infos_;
 
+  // Temporary data structure used when doing register allocation.
+  bool* const blocked_registers_;
+
   DISALLOW_COPY_AND_ASSIGN(CodeGenerator);
 };
 
diff --git a/compiler/optimizing/code_generator_arm.cc b/compiler/optimizing/code_generator_arm.cc
index 6e528f93d7..a446701b39 100644
--- a/compiler/optimizing/code_generator_arm.cc
+++ b/compiler/optimizing/code_generator_arm.cc
@@ -35,6 +35,81 @@ namespace arm {
 static constexpr int kNumberOfPushedRegistersAtEntry = 1;
 static constexpr int kCurrentMethodStackOffset = 0;
 
+CodeGeneratorARM::CodeGeneratorARM(HGraph* graph)
+    : CodeGenerator(graph, kNumberOfRegIds),
+      location_builder_(graph, this),
+      instruction_visitor_(graph, this) {}
+
+static bool* GetBlockedRegisterPairs(bool* blocked_registers) {
+  return blocked_registers + kNumberOfAllocIds;
+}
+
+ManagedRegister CodeGeneratorARM::AllocateFreeRegister(Primitive::Type type,
+                                                       bool* blocked_registers) const {
+  switch (type) {
+    case Primitive::kPrimLong: {
+      size_t reg = AllocateFreeRegisterInternal(
+          GetBlockedRegisterPairs(blocked_registers), kNumberOfRegisterPairs);
+      ArmManagedRegister pair =
+          ArmManagedRegister::FromRegisterPair(static_cast<RegisterPair>(reg));
+      blocked_registers[pair.AsRegisterPairLow()] = true;
+      blocked_registers[pair.AsRegisterPairHigh()] = true;
+      return pair;
+    }
+
+    case Primitive::kPrimByte:
+    case Primitive::kPrimBoolean:
+    case Primitive::kPrimChar:
+    case Primitive::kPrimShort:
+    case Primitive::kPrimInt:
+    case Primitive::kPrimNot: {
+      size_t reg = AllocateFreeRegisterInternal(blocked_registers, kNumberOfCoreRegisters);
+      return ArmManagedRegister::FromCoreRegister(static_cast<Register>(reg));
+    }
+
+    case Primitive::kPrimFloat:
+    case Primitive::kPrimDouble:
+      LOG(FATAL) << "Unimplemented register type " << type;
+
+    case Primitive::kPrimVoid:
+      LOG(FATAL) << "Unreachable type " << type;
+  }
+
+  return ManagedRegister::NoRegister();
+}
+
+void CodeGeneratorARM::SetupBlockedRegisters(bool* blocked_registers) const {
+  bool* blocked_register_pairs = GetBlockedRegisterPairs(blocked_registers);
+
+  // Don't allocate the dalvik style register pair passing.
+  blocked_register_pairs[R1_R2] = true;
+
+  // Stack register, LR and PC are always reserved.
+  blocked_registers[SP] = true;
+  blocked_registers[LR] = true;
+  blocked_registers[PC] = true;
+
+  // Reserve R4 for suspend check.
+  blocked_registers[R4] = true;
+  blocked_register_pairs[R4_R5] = true;
+
+  // Reserve thread register.
+  blocked_registers[TR] = true;
+
+  // TODO: We currently don't use Quick's callee saved registers.
+  blocked_registers[R5] = true;
+  blocked_registers[R6] = true;
+  blocked_registers[R7] = true;
+  blocked_registers[R8] = true;
+  blocked_registers[R10] = true;
+  blocked_registers[R11] = true;
+  blocked_register_pairs[R6_R7] = true;
+}
+
+size_t CodeGeneratorARM::GetNumberOfRegisters() const {
+  return kNumberOfRegIds;
+}
+
 static Location ArmCoreLocation(Register reg) {
   return Location::RegisterLocation(ArmManagedRegister::FromCoreRegister(reg));
 }
@@ -85,6 +160,32 @@ int32_t CodeGeneratorARM::GetStackSlot(HLocal* local) const {
   }
 }
 
+Location CodeGeneratorARM::GetStackLocation(HLoadLocal* load) const {
+  switch (load->GetType()) {
+    case Primitive::kPrimLong:
+      return Location::DoubleStackSlot(GetStackSlot(load->GetLocal()));
+      break;
+
+    case Primitive::kPrimInt:
+    case Primitive::kPrimNot:
+      return Location::StackSlot(GetStackSlot(load->GetLocal()));
+
+    case Primitive::kPrimFloat:
+    case Primitive::kPrimDouble:
+      LOG(FATAL) << "Unimplemented type " << load->GetType();
+
+    case Primitive::kPrimBoolean:
+    case Primitive::kPrimByte:
+    case Primitive::kPrimChar:
+    case Primitive::kPrimShort:
+    case Primitive::kPrimVoid:
+      LOG(FATAL) << "Unexpected type " << load->GetType();
+  }
+
+  LOG(FATAL) << "Unreachable";
+  return Location();
+}
+
 Location InvokeDexCallingConventionVisitor::GetNextLocation(Primitive::Type type) {
   switch (type) {
     case Primitive::kPrimBoolean:
@@ -302,7 +403,7 @@ void InstructionCodeGeneratorARM::VisitExit(HExit* exit) {
 
 void LocationsBuilderARM::VisitIf(HIf* if_instr) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(if_instr);
-  locations->SetInAt(0, ArmCoreLocation(R0));
+  locations->SetInAt(0, Location::RequiresRegister());
   if_instr->SetLocations(locations);
 }
 
@@ -317,9 +418,9 @@ void InstructionCodeGeneratorARM::VisitIf(HIf* if_instr) {
 
 void LocationsBuilderARM::VisitEqual(HEqual* equal) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(equal);
-  locations->SetInAt(0, ArmCoreLocation(R0));
-  locations->SetInAt(1, ArmCoreLocation(R1));
-  locations->SetOut(ArmCoreLocation(R0));
+  locations->SetInAt(0, Location::RequiresRegister());
+  locations->SetInAt(1, Location::RequiresRegister());
+  locations->SetOut(Location::RequiresRegister());
   equal->SetLocations(locations);
 }
 
@@ -409,7 +510,8 @@ void LocationsBuilderARM::VisitReturn(HReturn* ret) {
       break;
 
     case Primitive::kPrimLong:
-      locations->SetInAt(0, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
+      locations->SetInAt(
+          0, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
       break;
 
     default:
@@ -444,7 +546,7 @@ void InstructionCodeGeneratorARM::VisitReturn(HReturn* ret) {
 
 void LocationsBuilderARM::VisitInvokeStatic(HInvokeStatic* invoke) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(invoke);
-  locations->AddTemp(ArmCoreLocation(R0));
+  locations->AddTemp(Location::RequiresRegister());
 
   InvokeDexCallingConventionVisitor calling_convention_visitor;
   for (size_t i = 0; i < invoke->InputCount(); i++) {
@@ -512,19 +614,11 @@ void InstructionCodeGeneratorARM::VisitInvokeStatic(HInvokeStatic* invoke) {
 void LocationsBuilderARM::VisitAdd(HAdd* add) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(add);
   switch (add->GetResultType()) {
-    case Primitive::kPrimInt: {
-      locations->SetInAt(0, ArmCoreLocation(R0));
-      locations->SetInAt(1, ArmCoreLocation(R1));
-      locations->SetOut(ArmCoreLocation(R0));
-      break;
-    }
-
+    case Primitive::kPrimInt:
     case Primitive::kPrimLong: {
-      locations->SetInAt(
-          0, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
-      locations->SetInAt(
-          1, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R2_R3)));
-      locations->SetOut(Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
+      locations->SetInAt(0, Location::RequiresRegister());
+      locations->SetInAt(1, Location::RequiresRegister());
+      locations->SetOut(Location::RequiresRegister());
       break;
     }
 
@@ -574,19 +668,11 @@ void InstructionCodeGeneratorARM::VisitAdd(HAdd* add) {
 void LocationsBuilderARM::VisitSub(HSub* sub) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(sub);
   switch (sub->GetResultType()) {
-    case Primitive::kPrimInt: {
-      locations->SetInAt(0, ArmCoreLocation(R0));
-      locations->SetInAt(1, ArmCoreLocation(R1));
-      locations->SetOut(ArmCoreLocation(R0));
-      break;
-    }
-
+    case Primitive::kPrimInt:
     case Primitive::kPrimLong: {
-      locations->SetInAt(
-          0, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
-      locations->SetInAt(
-          1, Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R2_R3)));
-      locations->SetOut(Location::RegisterLocation(ArmManagedRegister::FromRegisterPair(R0_R1)));
+      locations->SetInAt(0, Location::RequiresRegister());
+      locations->SetInAt(1, Location::RequiresRegister());
+      locations->SetOut(Location::RequiresRegister());
       break;
     }
 
@@ -649,6 +735,9 @@ class InvokeRuntimeCallingConvention : public CallingConvention<Register> {
 
 void LocationsBuilderARM::VisitNewInstance(HNewInstance* instruction) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(instruction);
+  InvokeRuntimeCallingConvention calling_convention;
+  locations->AddTemp(ArmCoreLocation(calling_convention.GetRegisterAt(0)));
+  locations->AddTemp(ArmCoreLocation(calling_convention.GetRegisterAt(1)));
   locations->SetOut(ArmCoreLocation(R0));
   instruction->SetLocations(locations);
 }
@@ -683,8 +772,8 @@ void InstructionCodeGeneratorARM::VisitParameterValue(HParameterValue* instructi
 
 void LocationsBuilderARM::VisitNot(HNot* instruction) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(instruction);
-  locations->SetInAt(0, ArmCoreLocation(R0));
-  locations->SetOut(ArmCoreLocation(R0));
+  locations->SetInAt(0, Location::RequiresRegister());
+  locations->SetOut(Location::RequiresRegister());
   instruction->SetLocations(locations);
 }
 
diff --git a/compiler/optimizing/code_generator_arm.h b/compiler/optimizing/code_generator_arm.h
index ed35f94e2b..2405d4b5a6 100644
--- a/compiler/optimizing/code_generator_arm.h
+++ b/compiler/optimizing/code_generator_arm.h
@@ -101,10 +101,7 @@ class InstructionCodeGeneratorARM : public HGraphVisitor {
 
 class CodeGeneratorARM : public CodeGenerator {
  public:
-  explicit CodeGeneratorARM(HGraph* graph)
-      : CodeGenerator(graph),
-        location_builder_(graph, this),
-        instruction_visitor_(graph, this) { }
+  explicit CodeGeneratorARM(HGraph* graph);
   virtual ~CodeGeneratorARM() { }
 
   virtual void GenerateFrameEntry() OVERRIDE;
@@ -128,7 +125,13 @@ class CodeGeneratorARM : public CodeGenerator {
     return &assembler_;
   }
 
+  virtual void SetupBlockedRegisters(bool* blocked_registers) const OVERRIDE;
+  virtual ManagedRegister AllocateFreeRegister(
+      Primitive::Type type, bool* blocked_registers) const OVERRIDE;
+  virtual size_t GetNumberOfRegisters() const OVERRIDE;
+
   int32_t GetStackSlot(HLocal* local) const;
+  virtual Location GetStackLocation(HLoadLocal* load) const OVERRIDE;
 
  private:
   // Helper method to move a 32bits value between two locations.
diff --git a/compiler/optimizing/code_generator_x86.cc b/compiler/optimizing/code_generator_x86.cc
index dc10830e7f..fbb054ae88 100644
--- a/compiler/optimizing/code_generator_x86.cc
+++ b/compiler/optimizing/code_generator_x86.cc
@@ -35,6 +35,72 @@ namespace x86 {
 static constexpr int kNumberOfPushedRegistersAtEntry = 1;
 static constexpr int kCurrentMethodStackOffset = 0;
 
+CodeGeneratorX86::CodeGeneratorX86(HGraph* graph)
+    : CodeGenerator(graph, kNumberOfRegIds),
+      location_builder_(graph, this),
+      instruction_visitor_(graph, this) {}
+
+static bool* GetBlockedRegisterPairs(bool* blocked_registers) {
+  return blocked_registers + kNumberOfAllocIds;
+}
+
+ManagedRegister CodeGeneratorX86::AllocateFreeRegister(Primitive::Type type,
+                                                       bool* blocked_registers) const {
+  switch (type) {
+    case Primitive::kPrimLong: {
+      size_t reg = AllocateFreeRegisterInternal(
+          GetBlockedRegisterPairs(blocked_registers), kNumberOfRegisterPairs);
+      X86ManagedRegister pair =
+          X86ManagedRegister::FromRegisterPair(static_cast<RegisterPair>(reg));
+      blocked_registers[pair.AsRegisterPairLow()] = true;
+      blocked_registers[pair.AsRegisterPairHigh()] = true;
+      return pair;
+    }
+
+    case Primitive::kPrimByte:
+    case Primitive::kPrimBoolean:
+    case Primitive::kPrimChar:
+    case Primitive::kPrimShort:
+    case Primitive::kPrimInt:
+    case Primitive::kPrimNot: {
+      size_t reg = AllocateFreeRegisterInternal(blocked_registers, kNumberOfCpuRegisters);
+      return X86ManagedRegister::FromCpuRegister(static_cast<Register>(reg));
+    }
+
+    case Primitive::kPrimFloat:
+    case Primitive::kPrimDouble:
+      LOG(FATAL) << "Unimplemented register type " << type;
+
+    case Primitive::kPrimVoid:
+      LOG(FATAL) << "Unreachable type " << type;
+  }
+
+  return ManagedRegister::NoRegister();
+}
+
+void CodeGeneratorX86::SetupBlockedRegisters(bool* blocked_registers) const {
+  bool* blocked_register_pairs = GetBlockedRegisterPairs(blocked_registers);
+
+  // Don't allocate the dalvik style register pair passing.
+  blocked_register_pairs[ECX_EDX] = true;
+
+  // Stack register is always reserved.
+  blocked_registers[ESP] = true;
+
+  // TODO: We currently don't use Quick's callee saved registers.
+  blocked_registers[EBP] = true;
+  blocked_registers[ESI] = true;
+  blocked_registers[EDI] = true;
+  blocked_register_pairs[EAX_EDI] = true;
+  blocked_register_pairs[EDX_EDI] = true;
+  blocked_register_pairs[ECX_EDI] = true;
+  blocked_register_pairs[EBX_EDI] = true;
+}
+
+size_t CodeGeneratorX86::GetNumberOfRegisters() const {
+  return kNumberOfRegIds;
+}
+
 static Location X86CpuLocation(Register reg) {
   return Location::RegisterLocation(X86ManagedRegister::FromCpuRegister(reg));
 }
@@ -90,6 +156,33 @@ int32_t CodeGeneratorX86::GetStackSlot(HLocal* local) const {
   }
 }
 
+
+Location CodeGeneratorX86::GetStackLocation(HLoadLocal* load) const {
+  switch (load->GetType()) {
+    case Primitive::kPrimLong:
+      return Location::DoubleStackSlot(GetStackSlot(load->GetLocal()));
+      break;
+
+    case Primitive::kPrimInt:
+    case Primitive::kPrimNot:
+      return Location::StackSlot(GetStackSlot(load->GetLocal()));
+
+    case Primitive::kPrimFloat:
+    case Primitive::kPrimDouble:
+      LOG(FATAL) << "Unimplemented type " << load->GetType();
+
+    case Primitive::kPrimBoolean:
+    case Primitive::kPrimByte:
+    case Primitive::kPrimChar:
+    case Primitive::kPrimShort:
+    case Primitive::kPrimVoid:
+      LOG(FATAL) << "Unexpected type " << load->GetType();
+  }
+
+  LOG(FATAL) << "Unreachable";
+  return Location();
+}
+
 static constexpr Register kRuntimeParameterCoreRegisters[] = { EAX, ECX, EDX };
 static constexpr size_t kRuntimeParameterCoreRegistersLength =
     arraysize(kRuntimeParameterCoreRegisters);
@@ -311,13 +404,18 @@ void InstructionCodeGeneratorX86::VisitExit(HExit* exit) {
 
 void LocationsBuilderX86::VisitIf(HIf* if_instr) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(if_instr);
-  locations->SetInAt(0, X86CpuLocation(EAX));
+  locations->SetInAt(0, Location::Any());
   if_instr->SetLocations(locations);
 }
 
 void InstructionCodeGeneratorX86::VisitIf(HIf* if_instr) {
   // TODO: Generate the input as a condition, instead of materializing in a register.
-  __ cmpl(if_instr->GetLocations()->InAt(0).AsX86().AsCpuRegister(), Immediate(0));
+  Location location = if_instr->GetLocations()->InAt(0);
+  if (location.IsRegister()) {
+    __ cmpl(location.AsX86().AsCpuRegister(), Immediate(0));
+  } else {
+    __ cmpl(Address(ESP, location.GetStackIndex()), Immediate(0));
+  }
   __ j(kEqual, codegen_->GetLabelOf(if_instr->IfFalseSuccessor()));
   if (!codegen_->GoesToNextBlock(if_instr->GetBlock(), if_instr->IfTrueSuccessor())) {
     __ jmp(codegen_->GetLabelOf(if_instr->IfTrueSuccessor()));
@@ -367,16 +465,22 @@ void InstructionCodeGeneratorX86::VisitStoreLocal(HStoreLocal* store) {
 
 void LocationsBuilderX86::VisitEqual(HEqual* equal) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(equal);
-  locations->SetInAt(0, X86CpuLocation(EAX));
-  locations->SetInAt(1, X86CpuLocation(ECX));
-  locations->SetOut(X86CpuLocation(EAX));
+  locations->SetInAt(0, Location::RequiresRegister());
+  locations->SetInAt(1, Location::Any());
+  locations->SetOut(Location::SameAsFirstInput());
   equal->SetLocations(locations);
 }
 
 void InstructionCodeGeneratorX86::VisitEqual(HEqual* equal) {
-  __ cmpl(equal->GetLocations()->InAt(0).AsX86().AsCpuRegister(),
-          equal->GetLocations()->InAt(1).AsX86().AsCpuRegister());
-  __ setb(kEqual, equal->GetLocations()->Out().AsX86().AsCpuRegister());
+  LocationSummary* locations = equal->GetLocations();
+  if (locations->InAt(1).IsRegister()) {
+    __ cmpl(locations->InAt(0).AsX86().AsCpuRegister(),
+            locations->InAt(1).AsX86().AsCpuRegister());
+  } else {
+    __ cmpl(locations->InAt(0).AsX86().AsCpuRegister(),
+            Address(ESP, locations->InAt(1).GetStackIndex()));
+  }
+  __ setb(kEqual, locations->Out().AsX86().AsCpuRegister());
 }
 
 void LocationsBuilderX86::VisitIntConstant(HIntConstant* constant) {
@@ -453,7 +557,7 @@ void InstructionCodeGeneratorX86::VisitReturn(HReturn* ret) {
 
 void LocationsBuilderX86::VisitInvokeStatic(HInvokeStatic* invoke) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(invoke);
-  locations->AddTemp(X86CpuLocation(EAX));
+  locations->AddTemp(Location::RequiresRegister());
 
   InvokeDexCallingConventionVisitor calling_convention_visitor;
   for (size_t i = 0; i < invoke->InputCount(); i++) {
@@ -514,18 +618,11 @@ void InstructionCodeGeneratorX86::VisitInvokeStatic(HInvokeStatic* invoke) {
 void LocationsBuilderX86::VisitAdd(HAdd* add) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(add);
   switch (add->GetResultType()) {
-    case Primitive::kPrimInt: {
-      locations->SetInAt(0, X86CpuLocation(EAX));
-      locations->SetInAt(1, X86CpuLocation(ECX));
-      locations->SetOut(X86CpuLocation(EAX));
-      break;
-    }
+    case Primitive::kPrimInt:
     case Primitive::kPrimLong: {
-      locations->SetInAt(
-          0, Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(EAX_EDX)));
-      locations->SetInAt(
-          1, Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(ECX_EBX)));
-      locations->SetOut(Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(EAX_EDX)));
+      locations->SetInAt(0, Location::RequiresRegister());
+      locations->SetInAt(1, Location::Any());
+      locations->SetOut(Location::SameAsFirstInput());
       break;
     }
 
@@ -548,18 +645,30 @@ void InstructionCodeGeneratorX86::VisitAdd(HAdd* add) {
     case Primitive::kPrimInt: {
       DCHECK_EQ(locations->InAt(0).AsX86().AsCpuRegister(),
                 locations->Out().AsX86().AsCpuRegister());
-      __ addl(locations->InAt(0).AsX86().AsCpuRegister(),
-              locations->InAt(1).AsX86().AsCpuRegister());
+      if (locations->InAt(1).IsRegister()) {
+        __ addl(locations->InAt(0).AsX86().AsCpuRegister(),
+                locations->InAt(1).AsX86().AsCpuRegister());
+      } else {
+        __ addl(locations->InAt(0).AsX86().AsCpuRegister(),
+                Address(ESP, locations->InAt(1).GetStackIndex()));
+      }
       break;
     }
 
     case Primitive::kPrimLong: {
       DCHECK_EQ(locations->InAt(0).AsX86().AsRegisterPair(),
                 locations->Out().AsX86().AsRegisterPair());
-      __ addl(locations->InAt(0).AsX86().AsRegisterPairLow(),
-              locations->InAt(1).AsX86().AsRegisterPairLow());
-      __ adcl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
-              locations->InAt(1).AsX86().AsRegisterPairHigh());
+      if (locations->InAt(1).IsRegister()) {
+        __ addl(locations->InAt(0).AsX86().AsRegisterPairLow(),
+                locations->InAt(1).AsX86().AsRegisterPairLow());
+        __ adcl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
+                locations->InAt(1).AsX86().AsRegisterPairHigh());
+      } else {
+        __ addl(locations->InAt(0).AsX86().AsRegisterPairLow(),
+                Address(ESP, locations->InAt(1).GetStackIndex()));
+        __ adcl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
+                Address(ESP, locations->InAt(1).GetHighStackIndex(kX86WordSize)));
+      }
       break;
     }
 
@@ -578,19 +687,11 @@ void InstructionCodeGeneratorX86::VisitAdd(HAdd* add) {
 void LocationsBuilderX86::VisitSub(HSub* sub) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(sub);
   switch (sub->GetResultType()) {
-    case Primitive::kPrimInt: {
-      locations->SetInAt(0, X86CpuLocation(EAX));
-      locations->SetInAt(1, X86CpuLocation(ECX));
-      locations->SetOut(X86CpuLocation(EAX));
-      break;
-    }
-
+    case Primitive::kPrimInt:
     case Primitive::kPrimLong: {
-      locations->SetInAt(
-          0, Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(EAX_EDX)));
-      locations->SetInAt(
-          1, Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(ECX_EBX)));
-      locations->SetOut(Location::RegisterLocation(X86ManagedRegister::FromRegisterPair(EAX_EDX)));
+      locations->SetInAt(0, Location::RequiresRegister());
+      locations->SetInAt(1, Location::Any());
+      locations->SetOut(Location::SameAsFirstInput());
       break;
     }
 
@@ -613,18 +714,30 @@ void InstructionCodeGeneratorX86::VisitSub(HSub* sub) {
     case Primitive::kPrimInt: {
       DCHECK_EQ(locations->InAt(0).AsX86().AsCpuRegister(),
                 locations->Out().AsX86().AsCpuRegister());
-      __ subl(locations->InAt(0).AsX86().AsCpuRegister(),
-              locations->InAt(1).AsX86().AsCpuRegister());
+      if (locations->InAt(1).IsRegister()) {
+        __ subl(locations->InAt(0).AsX86().AsCpuRegister(),
+                locations->InAt(1).AsX86().AsCpuRegister());
+      } else {
+        __ subl(locations->InAt(0).AsX86().AsCpuRegister(),
+                Address(ESP, locations->InAt(1).GetStackIndex()));
+      }
       break;
     }
 
     case Primitive::kPrimLong: {
       DCHECK_EQ(locations->InAt(0).AsX86().AsRegisterPair(),
                 locations->Out().AsX86().AsRegisterPair());
-      __ subl(locations->InAt(0).AsX86().AsRegisterPairLow(),
-              locations->InAt(1).AsX86().AsRegisterPairLow());
-      __ sbbl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
-              locations->InAt(1).AsX86().AsRegisterPairHigh());
+      if (locations->InAt(1).IsRegister()) {
+        __ subl(locations->InAt(0).AsX86().AsRegisterPairLow(),
+                locations->InAt(1).AsX86().AsRegisterPairLow());
+        __ sbbl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
+                locations->InAt(1).AsX86().AsRegisterPairHigh());
+      } else {
+        __ subl(locations->InAt(0).AsX86().AsRegisterPairLow(),
+                Address(ESP, locations->InAt(1).GetStackIndex()));
+        __ sbbl(locations->InAt(0).AsX86().AsRegisterPairHigh(),
+                Address(ESP, locations->InAt(1).GetHighStackIndex(kX86WordSize)));
+      }
       break;
     }
 
@@ -643,14 +756,16 @@ void InstructionCodeGeneratorX86::VisitSub(HSub* sub) {
 void LocationsBuilderX86::VisitNewInstance(HNewInstance* instruction) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(instruction);
   locations->SetOut(X86CpuLocation(EAX));
+  InvokeRuntimeCallingConvention calling_convention;
+  locations->AddTemp(X86CpuLocation(calling_convention.GetRegisterAt(0)));
+  locations->AddTemp(X86CpuLocation(calling_convention.GetRegisterAt(1)));
   instruction->SetLocations(locations);
 }
 
 void InstructionCodeGeneratorX86::VisitNewInstance(HNewInstance* instruction) {
   InvokeRuntimeCallingConvention calling_convention;
   LoadCurrentMethod(calling_convention.GetRegisterAt(1));
-  __ movl(calling_convention.GetRegisterAt(0),
-          Immediate(instruction->GetTypeIndex()));
+  __ movl(calling_convention.GetRegisterAt(0), Immediate(instruction->GetTypeIndex()));
 
   __ fs()->call(
       Address::Absolute(QUICK_ENTRYPOINT_OFFSET(kX86WordSize, pAllocObjectWithAccessCheck)));
@@ -676,15 +791,16 @@ void InstructionCodeGeneratorX86::VisitParameterValue(HParameterValue* instructi
 
 void LocationsBuilderX86::VisitNot(HNot* instruction) {
   LocationSummary* locations = new (GetGraph()->GetArena()) LocationSummary(instruction);
-  locations->SetInAt(0, X86CpuLocation(EAX));
-  locations->SetOut(X86CpuLocation(EAX));
+  locations->SetInAt(0, Location::RequiresRegister());
+  locations->SetOut(Location::SameAsFirstInput());
   instruction->SetLocations(locations);
 }
 
 void InstructionCodeGeneratorX86::VisitNot(HNot* instruction) {
   LocationSummary* locations = instruction->GetLocations();
-  DCHECK_EQ(locations->InAt(0).AsX86().AsCpuRegister(), locations->Out().AsX86().AsCpuRegister());
-  __ xorl(locations->Out().AsX86().AsCpuRegister(), Immediate(1));
+  Location out = locations->Out();
+  DCHECK_EQ(locations->InAt(0).AsX86().AsCpuRegister(), out.AsX86().AsCpuRegister());
+  __ xorl(out.AsX86().AsCpuRegister(), Immediate(1));
 }
 
 void LocationsBuilderX86::VisitPhi(HPhi* instruction) {
diff --git a/compiler/optimizing/code_generator_x86.h b/compiler/optimizing/code_generator_x86.h
index f22890e708..1ee11bf0e8 100644
--- a/compiler/optimizing/code_generator_x86.h
+++ b/compiler/optimizing/code_generator_x86.h
@@ -102,10 +102,7 @@ class InstructionCodeGeneratorX86 : public HGraphVisitor {
 
 class CodeGeneratorX86 : public CodeGenerator {
  public:
-  explicit CodeGeneratorX86(HGraph* graph)
-      : CodeGenerator(graph),
-        location_builder_(graph, this),
-        instruction_visitor_(graph, this) { }
+  explicit CodeGeneratorX86(HGraph* graph);
   virtual ~CodeGeneratorX86() { }
 
   virtual void GenerateFrameEntry() OVERRIDE;
@@ -129,7 +126,13 @@ class CodeGeneratorX86 : public CodeGenerator {
     return &assembler_;
   }
 
+  virtual size_t GetNumberOfRegisters() const OVERRIDE;
+  virtual void SetupBlockedRegisters(bool* blocked_registers) const OVERRIDE;
+  virtual ManagedRegister AllocateFreeRegister(
+      Primitive::Type type, bool* blocked_registers) const OVERRIDE;
+
   int32_t GetStackSlot(HLocal* local) const;
+  virtual Location GetStackLocation(HLoadLocal* load) const OVERRIDE;
 
  private:
   // Helper method to move a 32bits value between two locations.