Merge "Share dex register maps between stack maps when possible."

2 files changed, 182 insertions, 37 deletions

diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index 9914ef49c3..5818a37a46 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -17,7 +17,8 @@
 #ifndef ART_COMPILER_OPTIMIZING_STACK_MAP_STREAM_H_
 #define ART_COMPILER_OPTIMIZING_STACK_MAP_STREAM_H_
 
-#include "base/bit_vector.h"
+#include "base/arena_containers.h"
+#include "base/bit_vector-inl.h"
 #include "base/value_object.h"
 #include "memory_region.h"
 #include "nodes.h"
@@ -40,7 +41,8 @@ class StackMapStream : public ValueObject {
         stack_mask_max_(-1),
         dex_pc_max_(0),
         native_pc_offset_max_(0),
-        number_of_stack_maps_with_inline_info_(0) {}
+        number_of_stack_maps_with_inline_info_(0),
+        dex_map_hash_to_stack_map_indices_(std::less<uint32_t>(), allocator->Adapter()) {}
 
   // Compute bytes needed to encode a mask with the given maximum element.
   static uint32_t StackMaskEncodingSize(int max_element) {
@@ -59,6 +61,7 @@ class StackMapStream : public ValueObject {
     size_t dex_register_locations_start_index;
     size_t inline_infos_start_index;
     BitVector* live_dex_registers_mask;
+    uint32_t dex_register_map_hash;
   };
 
   struct InlineInfoEntry {
@@ -80,6 +83,7 @@ class StackMapStream : public ValueObject {
     entry.inlining_depth = inlining_depth;
     entry.dex_register_locations_start_index = dex_register_locations_.Size();
     entry.inline_infos_start_index = inline_infos_.Size();
+    entry.dex_register_map_hash = 0;
     if (num_dex_registers != 0) {
       entry.live_dex_registers_mask =
           new (allocator_) ArenaBitVector(allocator_, num_dex_registers, true);
@@ -105,7 +109,7 @@ class StackMapStream : public ValueObject {
     inline_infos_.Add(entry);
   }
 
-  size_t ComputeNeededSize() const {
+  size_t ComputeNeededSize() {
     size_t size = CodeInfo::kFixedSize
         + ComputeStackMapsSize()
         + ComputeDexRegisterMapsSize()
@@ -118,7 +122,7 @@ class StackMapStream : public ValueObject {
     return StackMaskEncodingSize(stack_mask_max_);
   }
 
-  size_t ComputeStackMapsSize() const {
+  size_t ComputeStackMapsSize() {
     return stack_maps_.Size() * StackMap::ComputeStackMapSize(
         ComputeStackMaskSize(),
         ComputeInlineInfoSize(),
@@ -146,10 +150,13 @@ class StackMapStream : public ValueObject {
   }
 
   // Compute the size of all the Dex register maps.
-  size_t ComputeDexRegisterMapsSize() const {
+  size_t ComputeDexRegisterMapsSize() {
     size_t size = 0;
     for (size_t i = 0; i < stack_maps_.Size(); ++i) {
-      size += ComputeDexRegisterMapSize(stack_maps_.Get(i));
+      if (FindEntryWithTheSameDexMap(i) == kNoSameDexMapFound) {
+        // Entries with the same dex map will have the same offset.
+        size += ComputeDexRegisterMapSize(stack_maps_.Get(i));
+      }
     }
     return size;
   }
@@ -161,11 +168,11 @@ class StackMapStream : public ValueObject {
       + (number_of_stack_maps_with_inline_info_ * InlineInfo::kFixedSize);
   }
 
-  size_t ComputeDexRegisterMapsStart() const {
+  size_t ComputeDexRegisterMapsStart() {
     return CodeInfo::kFixedSize + ComputeStackMapsSize();
   }
 
-  size_t ComputeInlineInfoStart() const {
+  size_t ComputeInlineInfoStart() {
     return ComputeDexRegisterMapsStart() + ComputeDexRegisterMapsSize();
   }
 
@@ -206,38 +213,47 @@ class StackMapStream : public ValueObject {
         stack_map.SetStackMask(code_info, *entry.sp_mask);
       }
 
-      if (entry.num_dex_registers != 0) {
-        // Set the Dex register map.
-        MemoryRegion register_region =
-            dex_register_locations_region.Subregion(
-                next_dex_register_map_offset,
-                ComputeDexRegisterMapSize(entry));
-        next_dex_register_map_offset += register_region.size();
-        DexRegisterMap dex_register_map(register_region);
-        stack_map.SetDexRegisterMapOffset(
+      if (entry.num_dex_registers == 0) {
+        // No dex map available.
+        stack_map.SetDexRegisterMapOffset(code_info, StackMap::kNoDexRegisterMap);
+      } else {
+        // Search for an entry with the same dex map.
+        size_t entry_with_same_map = FindEntryWithTheSameDexMap(i);
+        if (entry_with_same_map != kNoSameDexMapFound) {
+          // If we have a hit reuse the offset.
+          stack_map.SetDexRegisterMapOffset(code_info,
+              code_info.GetStackMapAt(entry_with_same_map).GetDexRegisterMapOffset(code_info));
+        } else {
+          // New dex registers maps should be added to the stack map.
+          MemoryRegion register_region =
+              dex_register_locations_region.Subregion(
+                  next_dex_register_map_offset,
+                  ComputeDexRegisterMapSize(entry));
+          next_dex_register_map_offset += register_region.size();
+          DexRegisterMap dex_register_map(register_region);
+          stack_map.SetDexRegisterMapOffset(
             code_info, register_region.start() - dex_register_locations_region.start());
 
-        // Offset in `dex_register_map` where to store the next register entry.
-        size_t offset = DexRegisterMap::kFixedSize;
-        dex_register_map.SetLiveBitMask(offset,
-                                        entry.num_dex_registers,
-                                        *entry.live_dex_registers_mask);
-        offset += DexRegisterMap::LiveBitMaskSize(entry.num_dex_registers);
-        for (size_t dex_register_number = 0, index_in_dex_register_locations = 0;
-             dex_register_number < entry.num_dex_registers;
-             ++dex_register_number) {
-          if (entry.live_dex_registers_mask->IsBitSet(dex_register_number)) {
-            DexRegisterLocation dex_register_location = dex_register_locations_.Get(
-                entry.dex_register_locations_start_index + index_in_dex_register_locations);
-            dex_register_map.SetRegisterInfo(offset, dex_register_location);
-            offset += DexRegisterMap::EntrySize(dex_register_location);
-            ++index_in_dex_register_locations;
+          // Offset in `dex_register_map` where to store the next register entry.
+          size_t offset = DexRegisterMap::kFixedSize;
+          dex_register_map.SetLiveBitMask(offset,
+                                          entry.num_dex_registers,
+                                          *entry.live_dex_registers_mask);
+          offset += DexRegisterMap::LiveBitMaskSize(entry.num_dex_registers);
+          for (size_t dex_register_number = 0, index_in_dex_register_locations = 0;
+               dex_register_number < entry.num_dex_registers;
+               ++dex_register_number) {
+            if (entry.live_dex_registers_mask->IsBitSet(dex_register_number)) {
+              DexRegisterLocation dex_register_location = dex_register_locations_.Get(
+                  entry.dex_register_locations_start_index + index_in_dex_register_locations);
+              dex_register_map.SetRegisterInfo(offset, dex_register_location);
+              offset += DexRegisterMap::EntrySize(dex_register_location);
+              ++index_in_dex_register_locations;
+            }
           }
+          // Ensure we reached the end of the Dex registers region.
+          DCHECK_EQ(offset, register_region.size());
         }
-        // Ensure we reached the end of the Dex registers region.
-        DCHECK_EQ(offset, register_region.size());
-      } else {
-        stack_map.SetDexRegisterMapOffset(code_info, StackMap::kNoDexRegisterMap);
       }
 
       // Set the inlining info.
@@ -271,11 +287,86 @@ class StackMapStream : public ValueObject {
       DCHECK(DexRegisterLocation::IsShortLocationKind(kind))
           << DexRegisterLocation::PrettyDescriptor(kind);
       dex_register_locations_.Add(DexRegisterLocation(kind, value));
-      stack_maps_.Get(stack_maps_.Size() - 1).live_dex_registers_mask->SetBit(dex_register);
+      StackMapEntry entry = stack_maps_.Get(stack_maps_.Size() - 1);
+      entry.live_dex_registers_mask->SetBit(dex_register);
+      entry.dex_register_map_hash += (1 << dex_register);
+      entry.dex_register_map_hash += static_cast<uint32_t>(value);
+      entry.dex_register_map_hash += static_cast<uint32_t>(kind);
+      stack_maps_.Put(stack_maps_.Size() - 1, entry);
     }
   }
 
  private:
+  // Returns the index of an entry with the same dex register map
+  // or kNoSameDexMapFound if no such entry exists.
+  size_t FindEntryWithTheSameDexMap(size_t entry_index) {
+    StackMapEntry entry = stack_maps_.Get(entry_index);
+    auto entries_it = dex_map_hash_to_stack_map_indices_.find(entry.dex_register_map_hash);
+    if (entries_it == dex_map_hash_to_stack_map_indices_.end()) {
+      // We don't have a perfect hash functions so we need a list to collect all stack maps
+      // which might have the same dex register map.
+      GrowableArray<uint32_t> stack_map_indices(allocator_, 1);
+      stack_map_indices.Add(entry_index);
+      dex_map_hash_to_stack_map_indices_.Put(entry.dex_register_map_hash, stack_map_indices);
+      return kNoSameDexMapFound;
+    }
+
+    // TODO: We don't need to add ourselves to the map if we can guarantee that
+    // FindEntryWithTheSameDexMap is called just once per stack map entry.
+    // A good way to do this is to cache the offset in the stack map entry. This
+    // is easier to do if we add markers when the stack map constructions begins
+    // and when it ends.
+
+    // We might have collisions, so we need to check whether or not we should
+    // add the entry to the map. `needs_to_be_added` keeps track of this.
+    bool needs_to_be_added = true;
+    size_t result = kNoSameDexMapFound;
+    for (size_t i = 0; i < entries_it->second.Size(); i++) {
+      size_t test_entry_index = entries_it->second.Get(i);
+      if (test_entry_index == entry_index) {
+        needs_to_be_added = false;
+      } else if (HaveTheSameDexMaps(stack_maps_.Get(test_entry_index), entry)) {
+        result = test_entry_index;
+        needs_to_be_added = false;
+        break;
+      }
+    }
+    if (needs_to_be_added) {
+      entries_it->second.Add(entry_index);
+    }
+    return result;
+  }
+
+  bool HaveTheSameDexMaps(const StackMapEntry& a, const StackMapEntry& b) const {
+    if (a.live_dex_registers_mask == nullptr && b.live_dex_registers_mask == nullptr) {
+      return true;
+    }
+    if (a.live_dex_registers_mask == nullptr || b.live_dex_registers_mask == nullptr) {
+      return false;
+    }
+    if (a.num_dex_registers != b.num_dex_registers) {
+      return false;
+    }
+
+    int index_in_dex_register_locations = 0;
+    for (uint32_t i = 0; i < a.num_dex_registers; i++) {
+      if (a.live_dex_registers_mask->IsBitSet(i) != b.live_dex_registers_mask->IsBitSet(i)) {
+        return false;
+      }
+      if (a.live_dex_registers_mask->IsBitSet(i)) {
+        DexRegisterLocation a_loc = dex_register_locations_.Get(
+            a.dex_register_locations_start_index + index_in_dex_register_locations);
+        DexRegisterLocation b_loc = dex_register_locations_.Get(
+            b.dex_register_locations_start_index + index_in_dex_register_locations);
+        if (a_loc != b_loc) {
+          return false;
+        }
+        ++index_in_dex_register_locations;
+      }
+    }
+    return true;
+  }
+
   ArenaAllocator* allocator_;
   GrowableArray<StackMapEntry> stack_maps_;
   GrowableArray<DexRegisterLocation> dex_register_locations_;
@@ -285,6 +376,10 @@ class StackMapStream : public ValueObject {
   uint32_t native_pc_offset_max_;
   size_t number_of_stack_maps_with_inline_info_;
 
+  ArenaSafeMap<uint32_t, GrowableArray<uint32_t>> dex_map_hash_to_stack_map_indices_;
+
+  static constexpr uint32_t kNoSameDexMapFound = -1;
+
   ART_FRIEND_TEST(StackMapTest, Test1);
   ART_FRIEND_TEST(StackMapTest, Test2);
   ART_FRIEND_TEST(StackMapTest, TestNonLiveDexRegisters);
diff --git a/compiler/optimizing/stack_map_test.cc b/compiler/optimizing/stack_map_test.cc
index e7075c0aef..e5a9790254 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -231,4 +231,54 @@ TEST(StackMapTest, DexRegisterMapOffsetOverflow) {
   ASSERT_EQ(stack_map.GetDexRegisterMapOffset(code_info), StackMap::kNoDexRegisterMapSmallEncoding);
 }
 
+TEST(StackMapTest, TestShareDexRegisterMap) {
+  ArenaPool pool;
+  ArenaAllocator arena(&pool);
+  StackMapStream stream(&arena);
+
+  ArenaBitVector sp_mask(&arena, 0, false);
+  uint32_t number_of_dex_registers = 2;
+  // First stack map.
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 0);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+  // Second stack map, which should share the same dex register map.
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 0);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+  // Third stack map (doesn't share the dex register map).
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 2);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+
+  size_t size = stream.ComputeNeededSize();
+  void* memory = arena.Alloc(size, kArenaAllocMisc);
+  MemoryRegion region(memory, size);
+  stream.FillIn(region);
+
+  CodeInfo ci(region);
+  // Verify first stack map.
+  StackMap sm0 = ci.GetStackMapAt(0);
+  DexRegisterMap dex_registers0 = ci.GetDexRegisterMapOf(sm0, number_of_dex_registers);
+  ASSERT_EQ(0, dex_registers0.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers0.GetConstant(1, number_of_dex_registers));
+
+  // Verify second stack map.
+  StackMap sm1 = ci.GetStackMapAt(1);
+  DexRegisterMap dex_registers1 = ci.GetDexRegisterMapOf(sm1, number_of_dex_registers);
+  ASSERT_EQ(0, dex_registers1.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers1.GetConstant(1, number_of_dex_registers));
+
+  // Verify third stack map.
+  StackMap sm2 = ci.GetStackMapAt(2);
+  DexRegisterMap dex_registers2 = ci.GetDexRegisterMapOf(sm2, number_of_dex_registers);
+  ASSERT_EQ(2, dex_registers2.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers2.GetConstant(1, number_of_dex_registers));
+
+  // Verify dex register map offsets.
+  ASSERT_EQ(sm0.GetDexRegisterMapOffset(ci), sm1.GetDexRegisterMapOffset(ci));
+  ASSERT_NE(sm0.GetDexRegisterMapOffset(ci), sm2.GetDexRegisterMapOffset(ci));
+  ASSERT_NE(sm1.GetDexRegisterMapOffset(ci), sm2.GetDexRegisterMapOffset(ci));
+}
+
 }  // namespace art