Reland "LSE fix to allow different operation kinds in the same heap location"

This reverts commit 9ec5df7d57ed7d007cc8e3f03803184850753c82.

Reason for reland: Update CHECKs to reflect that the
crash only consistently crashes for Arm64.

Bug: 243136064
Change-Id: Ic2229d6f9963f0d3c0cef26c79b01e9a55876c1c

7 files changed, 118 insertions, 27 deletions

diff --git a/compiler/optimizing/load_store_analysis.h b/compiler/optimizing/load_store_analysis.h
index c0799dbe9e..c46a5b9cc1 100644
--- a/compiler/optimizing/load_store_analysis.h
+++ b/compiler/optimizing/load_store_analysis.h
@@ -171,14 +171,16 @@ class HeapLocation : public ArenaObject<kArenaAllocLSA> {
                size_t offset,
                HInstruction* index,
                size_t vector_length,
-               int16_t declaring_class_def_index)
+               int16_t declaring_class_def_index,
+               bool is_vec_op)
       : ref_info_(ref_info),
         type_(DataType::ToSigned(type)),
         offset_(offset),
         index_(index),
         vector_length_(vector_length),
         declaring_class_def_index_(declaring_class_def_index),
-        has_aliased_locations_(false) {
+        has_aliased_locations_(false),
+        is_vec_op_(is_vec_op) {
     DCHECK(ref_info != nullptr);
     DCHECK((offset == kInvalidFieldOffset && index != nullptr) ||
            (offset != kInvalidFieldOffset && index == nullptr));
@@ -189,6 +191,7 @@ class HeapLocation : public ArenaObject<kArenaAllocLSA> {
   size_t GetOffset() const { return offset_; }
   HInstruction* GetIndex() const { return index_; }
   size_t GetVectorLength() const { return vector_length_; }
+  bool IsVecOp() const { return is_vec_op_; }
 
   // Returns the definition of declaring class' dex index.
   // It's kDeclaringClassDefIndexForArrays for an array element.
@@ -227,11 +230,12 @@ class HeapLocation : public ArenaObject<kArenaAllocLSA> {
   // Declaring class's def's dex index.
   // Invalid when this HeapLocation is not field access.
   const int16_t declaring_class_def_index_;
-
   // Has aliased heap locations in the method, due to either the
   // reference is aliased or the array element is aliased via different
   // index names.
   bool has_aliased_locations_;
+  // Whether this HeapLocation represents a vector operation.
+  bool is_vec_op_;
 
   DISALLOW_COPY_AND_ASSIGN(HeapLocation);
 };
@@ -318,7 +322,8 @@ class HeapLocationCollector : public HGraphVisitor {
                                  field->GetFieldOffset().SizeValue(),
                                  nullptr,
                                  HeapLocation::kScalar,
-                                 field->GetDeclaringClassDefIndex());
+                                 field->GetDeclaringClassDefIndex(),
+                                 /*is_vec_op=*/false);
   }
 
   size_t GetArrayHeapLocation(HInstruction* instruction) const {
@@ -327,10 +332,10 @@ class HeapLocationCollector : public HGraphVisitor {
     HInstruction* index = instruction->InputAt(1);
     DataType::Type type = instruction->GetType();
     size_t vector_length = HeapLocation::kScalar;
+    const bool is_vec_op = instruction->IsVecStore() || instruction->IsVecLoad();
     if (instruction->IsArraySet()) {
       type = instruction->AsArraySet()->GetComponentType();
-    } else if (instruction->IsVecStore() ||
-               instruction->IsVecLoad()) {
+    } else if (is_vec_op) {
       HVecOperation* vec_op = instruction->AsVecOperation();
       type = vec_op->GetPackedType();
       vector_length = vec_op->GetVectorLength();
@@ -342,7 +347,8 @@ class HeapLocationCollector : public HGraphVisitor {
                                  HeapLocation::kInvalidFieldOffset,
                                  index,
                                  vector_length,
-                                 HeapLocation::kDeclaringClassDefIndexForArrays);
+                                 HeapLocation::kDeclaringClassDefIndexForArrays,
+                                 is_vec_op);
   }
 
   bool HasHeapStores() const {
@@ -364,7 +370,8 @@ class HeapLocationCollector : public HGraphVisitor {
                                size_t offset,
                                HInstruction* index,
                                size_t vector_length,
-                               int16_t declaring_class_def_index) const {
+                               int16_t declaring_class_def_index,
+                               bool is_vec_op) const {
     DataType::Type lookup_type = DataType::ToSigned(type);
     for (size_t i = 0; i < heap_locations_.size(); i++) {
       HeapLocation* loc = heap_locations_[i];
@@ -373,7 +380,8 @@ class HeapLocationCollector : public HGraphVisitor {
           loc->GetOffset() == offset &&
           loc->GetIndex() == index &&
           loc->GetVectorLength() == vector_length &&
-          loc->GetDeclaringClassDefIndex() == declaring_class_def_index) {
+          loc->GetDeclaringClassDefIndex() == declaring_class_def_index &&
+          loc->IsVecOp() == is_vec_op) {
         return i;
       }
     }
@@ -518,14 +526,15 @@ class HeapLocationCollector : public HGraphVisitor {
                                size_t offset,
                                HInstruction* index,
                                size_t vector_length,
-                               int16_t declaring_class_def_index) {
+                               int16_t declaring_class_def_index,
+                               bool is_vec_op) {
     HInstruction* original_ref = HuntForOriginalReference(ref);
     ReferenceInfo* ref_info = GetOrCreateReferenceInfo(original_ref);
     size_t heap_location_idx = FindHeapLocationIndex(
-        ref_info, type, offset, index, vector_length, declaring_class_def_index);
+        ref_info, type, offset, index, vector_length, declaring_class_def_index, is_vec_op);
     if (heap_location_idx == kHeapLocationNotFound) {
-      HeapLocation* heap_loc = new (allocator_)
-          HeapLocation(ref_info, type, offset, index, vector_length, declaring_class_def_index);
+      HeapLocation* heap_loc = new (allocator_) HeapLocation(
+          ref_info, type, offset, index, vector_length, declaring_class_def_index, is_vec_op);
       heap_locations_.push_back(heap_loc);
     }
   }
@@ -539,19 +548,22 @@ class HeapLocationCollector : public HGraphVisitor {
                             offset,
                             nullptr,
                             HeapLocation::kScalar,
-                            declaring_class_def_index);
+                            declaring_class_def_index,
+                            /*is_vec_op=*/false);
   }
 
   void VisitArrayAccess(HInstruction* array,
                         HInstruction* index,
                         DataType::Type type,
-                        size_t vector_length) {
+                        size_t vector_length,
+                        bool is_vec_op) {
     MaybeCreateHeapLocation(array,
                             type,
                             HeapLocation::kInvalidFieldOffset,
                             index,
                             vector_length,
-                            HeapLocation::kDeclaringClassDefIndexForArrays);
+                            HeapLocation::kDeclaringClassDefIndexForArrays,
+                            is_vec_op);
   }
 
   void VisitPredicatedInstanceFieldGet(HPredicatedInstanceFieldGet* instruction) override {
@@ -585,7 +597,7 @@ class HeapLocationCollector : public HGraphVisitor {
     HInstruction* array = instruction->InputAt(0);
     HInstruction* index = instruction->InputAt(1);
     DataType::Type type = instruction->GetType();
-    VisitArrayAccess(array, index, type, HeapLocation::kScalar);
+    VisitArrayAccess(array, index, type, HeapLocation::kScalar, /*is_vec_op=*/false);
     CreateReferenceInfoForReferenceType(instruction);
   }
 
@@ -593,7 +605,7 @@ class HeapLocationCollector : public HGraphVisitor {
     HInstruction* array = instruction->InputAt(0);
     HInstruction* index = instruction->InputAt(1);
     DataType::Type type = instruction->GetComponentType();
-    VisitArrayAccess(array, index, type, HeapLocation::kScalar);
+    VisitArrayAccess(array, index, type, HeapLocation::kScalar, /*is_vec_op=*/false);
     has_heap_stores_ = true;
   }
 
@@ -601,7 +613,7 @@ class HeapLocationCollector : public HGraphVisitor {
     HInstruction* array = instruction->InputAt(0);
     HInstruction* index = instruction->InputAt(1);
     DataType::Type type = instruction->GetPackedType();
-    VisitArrayAccess(array, index, type, instruction->GetVectorLength());
+    VisitArrayAccess(array, index, type, instruction->GetVectorLength(), /*is_vec_op=*/true);
     CreateReferenceInfoForReferenceType(instruction);
   }
 
@@ -609,7 +621,7 @@ class HeapLocationCollector : public HGraphVisitor {
     HInstruction* array = instruction->InputAt(0);
     HInstruction* index = instruction->InputAt(1);
     DataType::Type type = instruction->GetPackedType();
-    VisitArrayAccess(array, index, type, instruction->GetVectorLength());
+    VisitArrayAccess(array, index, type, instruction->GetVectorLength(), /*is_vec_op=*/true);
     has_heap_stores_ = true;
   }
 
diff --git a/compiler/optimizing/load_store_analysis_test.cc b/compiler/optimizing/load_store_analysis_test.cc
index 7930c0d23c..b6cf8b4abb 100644
--- a/compiler/optimizing/load_store_analysis_test.cc
+++ b/compiler/optimizing/load_store_analysis_test.cc
@@ -118,12 +118,13 @@ TEST_F(LoadStoreAnalysisTest, ArrayHeapLocations) {
   size_t field = HeapLocation::kInvalidFieldOffset;
   size_t vec = HeapLocation::kScalar;
   size_t class_def = HeapLocation::kDeclaringClassDefIndexForArrays;
+  const bool is_vec_op = false;
   size_t loc1 = heap_location_collector.FindHeapLocationIndex(
-      ref, type, field, c1, vec, class_def);
+      ref, type, field, c1, vec, class_def, is_vec_op);
   size_t loc2 = heap_location_collector.FindHeapLocationIndex(
-      ref, type, field, c2, vec, class_def);
+      ref, type, field, c2, vec, class_def, is_vec_op);
   size_t loc3 = heap_location_collector.FindHeapLocationIndex(
-      ref, type, field, index, vec, class_def);
+      ref, type, field, index, vec, class_def, is_vec_op);
   // must find this reference info for array in HeapLocationCollector.
   ASSERT_TRUE(ref != nullptr);
   // must find these heap locations;
diff --git a/compiler/optimizing/load_store_elimination.cc b/compiler/optimizing/load_store_elimination.cc
index b385e035d3..cd44082ba4 100644
--- a/compiler/optimizing/load_store_elimination.cc
+++ b/compiler/optimizing/load_store_elimination.cc
@@ -2113,9 +2113,15 @@ bool LSEVisitor::TryReplacingLoopPhiPlaceholderWithDefault(
   HInstruction* replacement = GetDefaultValue(type);
   for (uint32_t phi_placeholder_index : visited.Indexes()) {
     DCHECK(phi_placeholder_replacements_[phi_placeholder_index].IsInvalid());
-    phi_placeholder_replacements_[phi_placeholder_index] = Value::ForInstruction(replacement);
+    PhiPlaceholder curr = GetPhiPlaceholderAt(phi_placeholder_index);
+    HeapLocation* hl = heap_location_collector_.GetHeapLocation(curr.GetHeapLocation());
+    // We use both vector and non vector operations to analyze the information. However, we replace
+    // only non vector operations in this code path.
+    if (!hl->IsVecOp()) {
+      phi_placeholder_replacements_[phi_placeholder_index] = Value::ForInstruction(replacement);
+      phi_placeholders_to_materialize->ClearBit(phi_placeholder_index);
+    }
   }
-  phi_placeholders_to_materialize->Subtract(&visited);
   return true;
 }
 
@@ -2170,9 +2176,15 @@ bool LSEVisitor::TryReplacingLoopPhiPlaceholderWithSingleInput(
   DCHECK(replacement != nullptr);
   for (uint32_t phi_placeholder_index : visited.Indexes()) {
     DCHECK(phi_placeholder_replacements_[phi_placeholder_index].IsInvalid());
-    phi_placeholder_replacements_[phi_placeholder_index] = Value::ForInstruction(replacement);
+    PhiPlaceholder curr = GetPhiPlaceholderAt(phi_placeholder_index);
+    HeapLocation* hl = heap_location_collector_.GetHeapLocation(curr.GetHeapLocation());
+    // We use both vector and non vector operations to analyze the information. However, we replace
+    // only vector operations in this code path.
+    if (hl->IsVecOp()) {
+      phi_placeholder_replacements_[phi_placeholder_index] = Value::ForInstruction(replacement);
+      phi_placeholders_to_materialize->ClearBit(phi_placeholder_index);
+    }
   }
-  phi_placeholders_to_materialize->Subtract(&visited);
   return true;
 }
 
diff --git a/test/2256-checker-vector-replacement/expected-stderr.txt b/test/2256-checker-vector-replacement/expected-stderr.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/2256-checker-vector-replacement/expected-stderr.txt
diff --git a/test/2256-checker-vector-replacement/expected-stdout.txt b/test/2256-checker-vector-replacement/expected-stdout.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/2256-checker-vector-replacement/expected-stdout.txt
diff --git a/test/2256-checker-vector-replacement/info.txt b/test/2256-checker-vector-replacement/info.txt
new file mode 100644
index 0000000000..374d7eb13c
--- /dev/null
+++ b/test/2256-checker-vector-replacement/info.txt
@@ -0,0 +1,2 @@
+Regression test to check we can perform LSE even when vector and
+non-vector operations reference the same heap location.
diff --git a/test/2256-checker-vector-replacement/src/Main.java b/test/2256-checker-vector-replacement/src/Main.java
new file mode 100644
index 0000000000..fca80fbeda
--- /dev/null
+++ b/test/2256-checker-vector-replacement/src/Main.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2023 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 void main(String[] args) {
+        $noinline$testVectorAndNonVector();
+    }
+
+    // Before loop optimization we only had an array get. After it, we optimized to also have
+    // VecLoad operations. This happens consistently only for Arm64. Arm32 vectorizes consistently
+    // but it also removes the ArrayGet. X86/X86_64 doesn't vectorize consistently (other
+    // vectorization tests also ignore x86/x86_64).
+
+    /// CHECK-START-ARM64: void Main.$noinline$testVectorAndNonVector() loop_optimization (before)
+    /// CHECK:     ArrayGet
+
+    /// CHECK-START-ARM64: void Main.$noinline$testVectorAndNonVector() loop_optimization (after)
+    /// CHECK:     ArrayGet
+
+    /// CHECK-START-ARM64: void Main.$noinline$testVectorAndNonVector() loop_optimization (before)
+    /// CHECK-NOT: VecLoad
+
+    /// CHECK-START-ARM64: void Main.$noinline$testVectorAndNonVector() loop_optimization (after)
+    /// CHECK:     VecLoad
+
+    // In LoadStoreElimination both ArrayGet and VecLoad have the same heap location. We will try to
+    // replace the ArrayGet with the constant 0. The crash happens when we want to do the same with
+    // the vector operation, changing the vector operation to a scalar.
+
+    /// CHECK-START-ARM64: void Main.$noinline$testVectorAndNonVector() load_store_elimination (before)
+    /// CHECK-DAG:     VecLoad outer_loop:<<VecBlock:B\d+>>
+    /// CHECK-DAG:     ArrayGet outer_loop:<<ScalarBlock:B\d+>>
+    /// CHECK-EVAL:    "<<VecBlock>>" == "<<ScalarBlock>>"
+
+    private static void $noinline$testVectorAndNonVector() {
+        int[] result = new int[2];
+        int[] source = new int[12];
+
+        // This will get vectorized.
+        for (int i = 0; i < result.length; ++i) {
+            int value = 0;
+            // Always true but needed to repro a crash since we need Phis.
+            if (i + 10 < source.length) {
+                for (int j = 0; j < 10; j++) {
+                    value += Math.abs(source[i + j]);
+                }
+            }
+            result[i] = value;
+        }
+    }
+}