Remove entry SuspendCheck for methods which only call on the slow path

We can skip SuspendCheck if the method is e.g. a leaf method with a read barrier call in the slow path.

Bug: 135477345
Test: ART tests
Change-Id: I6e04f10544ec61b46bb5763a88c28248e88193bf
diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index a353d07..cabf949 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -982,6 +982,9 @@
     if (locations != nullptr) {
       if (locations->CanCall()) {
         MarkNotLeaf();
+        if (locations->NeedsSuspendCheckEntry()) {
+          MarkNeedsSuspendCheckEntry();
+        }
       } else if (locations->Intrinsified() &&
                  instruction->IsInvokeStaticOrDirect() &&
                  !instruction->AsInvokeStaticOrDirect()->HasCurrentMethodInput()) {
@@ -1061,6 +1064,7 @@
       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 99de61d..cf59be8 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -403,6 +403,14 @@
     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 @@
   // 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 7bb754c..acaea71 100644
--- a/compiler/optimizing/locations.h
+++ b/compiler/optimizing/locations.h
@@ -605,13 +605,19 @@
   }
 
   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 42e6498..1e56928 100644
--- a/compiler/optimizing/register_allocator_graph_color.cc
+++ b/compiler/optimizing/register_allocator_graph_color.cc
@@ -810,10 +810,9 @@
   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 0d6c5a3..35ac2d6 100644
--- a/compiler/optimizing/register_allocator_linear_scan.cc
+++ b/compiler/optimizing/register_allocator_linear_scan.cc
@@ -225,7 +225,19 @@
   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 @@
   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);
   }
diff --git a/test/2234-checker-remove-entry-suspendcheck/expected-stderr.txt b/test/2234-checker-remove-entry-suspendcheck/expected-stderr.txt
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test/2234-checker-remove-entry-suspendcheck/expected-stderr.txt
diff --git a/test/2234-checker-remove-entry-suspendcheck/expected-stdout.txt b/test/2234-checker-remove-entry-suspendcheck/expected-stdout.txt
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test/2234-checker-remove-entry-suspendcheck/expected-stdout.txt
diff --git a/test/2234-checker-remove-entry-suspendcheck/info.txt b/test/2234-checker-remove-entry-suspendcheck/info.txt
new file mode 100644
index 0000000..7887589
--- /dev/null
+++ b/test/2234-checker-remove-entry-suspendcheck/info.txt
@@ -0,0 +1 @@
+Checks that we can skip the entry SuspendCheck for some leaf methods.
diff --git a/test/2234-checker-remove-entry-suspendcheck/src/Main.java b/test/2234-checker-remove-entry-suspendcheck/src/Main.java
new file mode 100644
index 0000000..bdb87d6
--- /dev/null
+++ b/test/2234-checker-remove-entry-suspendcheck/src/Main.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Main {
+  public static class Inner {
+    String str;
+    int[] arr = {1, 2, 3};
+
+    // Test 1: This test checks whether the SuspendCheck is removed from a simple field get.
+
+    /// CHECK-START: java.lang.String Main$Inner.$noinline$removeSuspendCheckFieldGet() register (before)
+    /// CHECK: SuspendCheck
+    /// CHECK: InstanceFieldGet
+    /// CHECK-START: java.lang.String Main$Inner.$noinline$removeSuspendCheckFieldGet() register (after)
+    /// CHECK-NOT: SuspendCheck
+    /// CHECK: InstanceFieldGet
+    public String $noinline$removeSuspendCheckFieldGet() {
+      return this.str;
+    }
+
+    // Test 2: This test checks whether the SuspendCheck is removed from a simple array get.
+
+    /// CHECK-START: int Main$Inner.$noinline$removeSuspendCheckArrayGet() register (before)
+    /// CHECK: SuspendCheck
+    /// CHECK: ArrayGet
+    /// CHECK-START: int Main$Inner.$noinline$removeSuspendCheckArrayGet() register (after)
+    /// CHECK-NOT: SuspendCheck
+    /// CHECK: ArrayGet
+    public int $noinline$removeSuspendCheckArrayGet() {
+      return this.arr[0];
+    }
+
+    // Test 3: This test checks whether the SuspendCheck is removed from a simple array set.
+
+    /// CHECK-START: int Main$Inner.$noinline$removeSuspendCheckArraySet() register (before)
+    /// CHECK: SuspendCheck
+    /// CHECK: ArraySet
+    /// CHECK-START: int Main$Inner.$noinline$removeSuspendCheckArraySet() register (after)
+    /// CHECK-NOT: SuspendCheck
+    /// CHECK: ArraySet
+    public int $noinline$removeSuspendCheckArraySet() {
+      return this.arr[0] = 2;
+    }
+  }
+
+  public static void main(String[] args) throws Exception {
+    Inner i = new Inner();
+    i.$noinline$removeSuspendCheckFieldGet();
+    i.$noinline$removeSuspendCheckArrayGet();
+    i.$noinline$removeSuspendCheckArraySet();
+  }
+}
\ No newline at end of file