5 files changed, 35 insertions, 15 deletions

diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index a353d0758e..cabf94991d 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -982,6 +982,9 @@ void CodeGenerator::AllocateLocations(HInstruction* instruction) {
     if (locations != nullptr) {
       if (locations->CanCall()) {
         MarkNotLeaf();
+        if (locations->NeedsSuspendCheckEntry()) {
+          MarkNeedsSuspendCheckEntry();
+        }
       } else if (locations->Intrinsified() &&
                  instruction->IsInvokeStaticOrDirect() &&
                  !instruction->AsInvokeStaticOrDirect()->HasCurrentMethodInput()) {
@@ -1061,6 +1064,7 @@ CodeGenerator::CodeGenerator(HGraph* graph,
       current_slow_path_(nullptr),
       current_block_index_(0),
       is_leaf_(true),
+      needs_suspend_check_entry_(false),
       requires_current_method_(false),
       code_generation_data_() {
   if (GetGraph()->IsCompilingOsr()) {
diff --git a/compiler/optimizing/code_generator.h b/compiler/optimizing/code_generator.h
index 99de61ddce..cf59be86bf 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -403,6 +403,14 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
     requires_current_method_ = true;
   }
 
+  bool NeedsSuspendCheckEntry() const {
+    return needs_suspend_check_entry_;
+  }
+
+  void MarkNeedsSuspendCheckEntry() {
+    needs_suspend_check_entry_ = true;
+  }
+
   void SetRequiresCurrentMethod() {
     requires_current_method_ = true;
   }
@@ -855,6 +863,9 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
   // Whether the method is a leaf method.
   bool is_leaf_;
 
+  // Whether the method has to emit a SuspendCheck at entry.
+  bool needs_suspend_check_entry_;
+
   // Whether an instruction in the graph accesses the current method.
   // TODO: Rename: this actually indicates that some instruction in the method
   // needs the environment including a valid stack frame.
diff --git a/compiler/optimizing/locations.h b/compiler/optimizing/locations.h
index 7bb754cfec..acaea71a49 100644
--- a/compiler/optimizing/locations.h
+++ b/compiler/optimizing/locations.h
@@ -605,13 +605,19 @@ class LocationSummary : public ArenaObject<kArenaAllocLocationSummary> {
   }
 
   bool CallsOnSlowPath() const {
-    return call_kind_ == kCallOnSlowPath || call_kind_ == kCallOnMainAndSlowPath;
+    return OnlyCallsOnSlowPath() || CallsOnMainAndSlowPath();
   }
 
   bool OnlyCallsOnSlowPath() const {
     return call_kind_ == kCallOnSlowPath;
   }
 
+  bool NeedsSuspendCheckEntry() const {
+    // Slow path calls do not need a SuspendCheck at method entry since they go into the runtime,
+    // which we expect to either do a suspend check or return quickly.
+    return WillCall();
+  }
+
   bool CallsOnMainAndSlowPath() const {
     return call_kind_ == kCallOnMainAndSlowPath;
   }
diff --git a/compiler/optimizing/register_allocator_graph_color.cc b/compiler/optimizing/register_allocator_graph_color.cc
index 42e6498148..1e5692855f 100644
--- a/compiler/optimizing/register_allocator_graph_color.cc
+++ b/compiler/optimizing/register_allocator_graph_color.cc
@@ -810,10 +810,9 @@ void RegisterAllocatorGraphColor::ProcessInstruction(HInstruction* instruction)
   if (locations == nullptr) {
     return;
   }
-  if (locations->NeedsSafepoint() && codegen_->IsLeafMethod()) {
+  if (instruction->IsSuspendCheckEntry() && !codegen_->NeedsSuspendCheckEntry()) {
     // We do this here because we do not want the suspend check to artificially
     // create live registers.
-    DCHECK(instruction->IsSuspendCheckEntry());
     DCHECK_EQ(locations->GetTempCount(), 0u);
     instruction->GetBlock()->RemoveInstruction(instruction);
     return;
diff --git a/compiler/optimizing/register_allocator_linear_scan.cc b/compiler/optimizing/register_allocator_linear_scan.cc
index 0d6c5a3eff..35ac2d649e 100644
--- a/compiler/optimizing/register_allocator_linear_scan.cc
+++ b/compiler/optimizing/register_allocator_linear_scan.cc
@@ -225,7 +225,19 @@ void RegisterAllocatorLinearScan::ProcessInstruction(HInstruction* instruction)
   LocationSummary* locations = instruction->GetLocations();
   size_t position = instruction->GetLifetimePosition();
 
+  // Check for early returns.
   if (locations == nullptr) return;
+  if (locations->NeedsSafepoint()) {
+    if (instruction->IsSuspendCheckEntry() && !codegen_->NeedsSuspendCheckEntry()) {
+      // TODO: We do this here because we do not want the suspend check to artificially
+      // create live registers. We should find another place, but this is currently the
+      // simplest.
+      DCHECK_EQ(locations->GetTempCount(), 0u);
+      instruction->GetBlock()->RemoveInstruction(instruction);
+      return;
+    }
+    safepoints_.push_back(instruction);
+  }
 
   // Create synthesized intervals for temporaries.
   for (size_t i = 0; i < locations->GetTempCount(); ++i) {
@@ -271,18 +283,6 @@ void RegisterAllocatorLinearScan::ProcessInstruction(HInstruction* instruction)
   bool core_register = (instruction->GetType() != DataType::Type::kFloat64)
       && (instruction->GetType() != DataType::Type::kFloat32);
 
-  if (locations->NeedsSafepoint()) {
-    if (codegen_->IsLeafMethod()) {
-      // TODO: We do this here because we do not want the suspend check to artificially
-      // create live registers. We should find another place, but this is currently the
-      // simplest.
-      DCHECK(instruction->IsSuspendCheckEntry());
-      instruction->GetBlock()->RemoveInstruction(instruction);
-      return;
-    }
-    safepoints_.push_back(instruction);
-  }
-
   if (locations->WillCall()) {
     BlockRegisters(position, position + 1, /* caller_save_only= */ true);
   }