Deduplicate register masks

Data is commonly shared between different stack maps. The register
masks are stored after the stack masks.

Oat size for a large app:
96722288 -> 94485872 (-2.31%)

Average oat size reduction according to golem -3.193%.

Bug: 34621054

Test: test-art-host

Change-Id: I5eacf668992e866d11ddba0c01675038a16cdfb4
diff --git a/compiler/optimizing/stack_map_stream.cc b/compiler/optimizing/stack_map_stream.cc
index 46497e3..668108d 100644
--- a/compiler/optimizing/stack_map_stream.cc
+++ b/compiler/optimizing/stack_map_stream.cc
@@ -157,37 +157,43 @@
 }
 
 size_t StackMapStream::PrepareForFillIn() {
-  size_t stack_mask_size_in_bits = stack_mask_max_ + 1;  // Need room for max element too.
-  size_t number_of_stack_masks = PrepareStackMasks(stack_mask_size_in_bits);
+  const size_t stack_mask_size_in_bits = stack_mask_max_ + 1;  // Need room for max element too.
+  const size_t number_of_stack_masks = PrepareStackMasks(stack_mask_size_in_bits);
+  const size_t register_mask_size_in_bits = MinimumBitsToStore(register_mask_max_);
+  const size_t number_of_register_masks = PrepareRegisterMasks();
   dex_register_maps_size_ = ComputeDexRegisterMapsSize();
   ComputeInlineInfoEncoding();  // needs dex_register_maps_size_.
   inline_info_size_ = inline_infos_.size() * inline_info_encoding_.GetEntrySize();
   CodeOffset max_native_pc_offset = ComputeMaxNativePcCodeOffset();
   // The stack map contains compressed native PC offsets.
-  size_t stack_map_size = stack_map_encoding_.SetFromSizes(max_native_pc_offset.CompressedValue(),
-                                                           dex_pc_max_,
-                                                           dex_register_maps_size_,
-                                                           inline_info_size_,
-                                                           register_mask_max_,
-                                                           number_of_stack_masks);
+  const size_t stack_map_size = stack_map_encoding_.SetFromSizes(
+      max_native_pc_offset.CompressedValue(),
+      dex_pc_max_,
+      dex_register_maps_size_,
+      inline_info_size_,
+      number_of_register_masks,
+      number_of_stack_masks);
   stack_maps_size_ = RoundUp(stack_maps_.size() * stack_map_size, kBitsPerByte) / kBitsPerByte;
   dex_register_location_catalog_size_ = ComputeDexRegisterLocationCatalogSize();
-  size_t stack_masks_bytes =
-      RoundUp(number_of_stack_masks * stack_mask_size_in_bits, kBitsPerByte) / kBitsPerByte;
-
-  size_t non_header_size =
+  const size_t stack_masks_bits = number_of_stack_masks * stack_mask_size_in_bits;
+  const size_t register_masks_bits = number_of_register_masks * register_mask_size_in_bits;
+  // Register masks are last, stack masks are right before that last.
+  // They are both bit packed / aligned.
+  const size_t non_header_size =
       stack_maps_size_ +
       dex_register_location_catalog_size_ +
       dex_register_maps_size_ +
       inline_info_size_ +
-      stack_masks_bytes;
+      RoundUp(stack_masks_bits + register_masks_bits, kBitsPerByte) / kBitsPerByte;
 
   // Prepare the CodeInfo variable-sized encoding.
   CodeInfoEncoding code_info_encoding;
   code_info_encoding.non_header_size = non_header_size;
   code_info_encoding.number_of_stack_maps = stack_maps_.size();
   code_info_encoding.number_of_stack_masks = number_of_stack_masks;
+  code_info_encoding.number_of_register_masks = number_of_register_masks;
   code_info_encoding.stack_mask_size_in_bits = stack_mask_size_in_bits;
+  code_info_encoding.register_mask_size_in_bits = register_mask_size_in_bits;
   code_info_encoding.stack_map_encoding = stack_map_encoding_;
   code_info_encoding.inline_info_encoding = inline_info_encoding_;
   code_info_encoding.number_of_location_catalog_entries = location_catalog_entries_.size();
@@ -330,7 +336,7 @@
 
     stack_map.SetDexPc(stack_map_encoding_, entry.dex_pc);
     stack_map.SetNativePcCodeOffset(stack_map_encoding_, entry.native_pc_code_offset);
-    stack_map.SetRegisterMask(stack_map_encoding_, entry.register_mask);
+    stack_map.SetRegisterMaskIndex(stack_map_encoding_, entry.register_mask_index);
     stack_map.SetStackMaskIndex(stack_map_encoding_, entry.stack_mask_index);
 
     if (entry.num_dex_registers == 0 || (entry.live_dex_registers_mask->NumSetBits() == 0)) {
@@ -422,7 +428,7 @@
     }
   }
 
-  // Write stack masks at the end.
+  // Write stack masks table.
   size_t stack_mask_bits = encoding.stack_mask_size_in_bits;
   if (stack_mask_bits > 0) {
     size_t stack_mask_bytes = RoundUp(stack_mask_bits, kBitsPerByte) / kBitsPerByte;
@@ -435,6 +441,12 @@
     }
   }
 
+  // Write register masks table.
+  for (size_t i = 0; i < encoding.number_of_register_masks; ++i) {
+    BitMemoryRegion register_mask = code_info.GetRegisterMask(encoding, i);
+    register_mask.StoreBits(0, register_masks_[i], encoding.register_mask_size_in_bits);
+  }
+
   // Verify all written data in debug build.
   if (kIsDebugBuild) {
     CheckCodeInfo(region);
@@ -548,6 +560,17 @@
   }
 }
 
+size_t StackMapStream::PrepareRegisterMasks() {
+  register_masks_.resize(stack_maps_.size(), 0u);
+  std::unordered_map<uint32_t, size_t> dedupe;
+  for (StackMapEntry& stack_map : stack_maps_) {
+    const size_t index = dedupe.size();
+    stack_map.register_mask_index = dedupe.emplace(stack_map.register_mask, index).first->second;
+    register_masks_[index] = stack_map.register_mask;
+  }
+  return dedupe.size();
+}
+
 size_t StackMapStream::PrepareStackMasks(size_t entry_size_in_bits) {
   // Preallocate memory since we do not want it to move (the dedup map will point into it).
   const size_t byte_entry_size = RoundUp(entry_size_in_bits, kBitsPerByte) / kBitsPerByte;
@@ -583,7 +606,8 @@
     DCHECK_EQ(stack_map.GetNativePcOffset(stack_map_encoding, instruction_set_),
               entry.native_pc_code_offset.Uint32Value(instruction_set_));
     DCHECK_EQ(stack_map.GetDexPc(stack_map_encoding), entry.dex_pc);
-    DCHECK_EQ(stack_map.GetRegisterMask(stack_map_encoding), entry.register_mask);
+    DCHECK_EQ(stack_map.GetRegisterMaskIndex(stack_map_encoding), entry.register_mask_index);
+    DCHECK_EQ(code_info.GetRegisterMaskOf(encoding, stack_map), entry.register_mask);
     const size_t num_stack_mask_bits = code_info.GetNumberOfStackMaskBits(encoding);
     DCHECK_EQ(stack_map.GetStackMaskIndex(stack_map_encoding), entry.stack_mask_index);
     BitMemoryRegion stack_mask = code_info.GetStackMaskOf(encoding, stack_map);
diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index e2e16e8..b1069a1 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -69,6 +69,7 @@
         dex_register_locations_(allocator->Adapter(kArenaAllocStackMapStream)),
         inline_infos_(allocator->Adapter(kArenaAllocStackMapStream)),
         stack_masks_(allocator->Adapter(kArenaAllocStackMapStream)),
+        register_masks_(allocator->Adapter(kArenaAllocStackMapStream)),
         stack_mask_max_(-1),
         dex_pc_max_(0),
         register_mask_max_(0),
@@ -109,6 +110,7 @@
     uint32_t dex_register_map_hash;
     size_t same_dex_register_map_as_;
     uint32_t stack_mask_index;
+    uint32_t register_mask_index;
   };
 
   struct InlineInfoEntry {
@@ -165,6 +167,9 @@
   // Returns the number of unique stack masks.
   size_t PrepareStackMasks(size_t entry_size_in_bits);
 
+  // Returns the number of unique register masks.
+  size_t PrepareRegisterMasks();
+
   // Returns the index of an entry with the same dex register map as the current_entry,
   // or kNoSameDexMapFound if no such entry exists.
   size_t FindEntryWithTheSameDexMap();
@@ -199,6 +204,7 @@
   ArenaVector<size_t> dex_register_locations_;
   ArenaVector<InlineInfoEntry> inline_infos_;
   ArenaVector<uint8_t> stack_masks_;
+  ArenaVector<uint32_t> register_masks_;
   int stack_mask_max_;
   uint32_t dex_pc_max_;
   uint32_t register_mask_max_;
diff --git a/compiler/optimizing/stack_map_test.cc b/compiler/optimizing/stack_map_test.cc
index da68b60..ce6d5c2 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -80,7 +80,7 @@
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64, encoding)));
   ASSERT_EQ(0u, stack_map.GetDexPc(encoding.stack_map_encoding));
   ASSERT_EQ(64u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-  ASSERT_EQ(0x3u, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+  ASSERT_EQ(0x3u, code_info.GetRegisterMaskOf(encoding, stack_map));
 
   ASSERT_TRUE(CheckStackMask(code_info, encoding, stack_map, sp_mask));
 
@@ -195,7 +195,7 @@
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64, encoding)));
     ASSERT_EQ(0u, stack_map.GetDexPc(encoding.stack_map_encoding));
     ASSERT_EQ(64u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-    ASSERT_EQ(0x3u, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+    ASSERT_EQ(0x3u, code_info.GetRegisterMaskOf(encoding, stack_map));
 
     ASSERT_TRUE(CheckStackMask(code_info, encoding, stack_map, sp_mask1));
 
@@ -254,7 +254,7 @@
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(128u, encoding)));
     ASSERT_EQ(1u, stack_map.GetDexPc(encoding.stack_map_encoding));
     ASSERT_EQ(128u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-    ASSERT_EQ(0xFFu, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+    ASSERT_EQ(0xFFu, code_info.GetRegisterMaskOf(encoding, stack_map));
 
     ASSERT_TRUE(CheckStackMask(code_info, encoding, stack_map, sp_mask2));
 
@@ -308,7 +308,7 @@
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(192u, encoding)));
     ASSERT_EQ(2u, stack_map.GetDexPc(encoding.stack_map_encoding));
     ASSERT_EQ(192u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-    ASSERT_EQ(0xABu, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+    ASSERT_EQ(0xABu, code_info.GetRegisterMaskOf(encoding, stack_map));
 
     ASSERT_TRUE(CheckStackMask(code_info, encoding, stack_map, sp_mask3));
 
@@ -362,7 +362,7 @@
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(256u, encoding)));
     ASSERT_EQ(3u, stack_map.GetDexPc(encoding.stack_map_encoding));
     ASSERT_EQ(256u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-    ASSERT_EQ(0xCDu, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+    ASSERT_EQ(0xCDu, code_info.GetRegisterMaskOf(encoding, stack_map));
 
     ASSERT_TRUE(CheckStackMask(code_info, encoding, stack_map, sp_mask4));
 
@@ -444,7 +444,7 @@
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64, encoding)));
   ASSERT_EQ(0u, stack_map.GetDexPc(encoding.stack_map_encoding));
   ASSERT_EQ(64u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-  ASSERT_EQ(0x3u, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+  ASSERT_EQ(0x3u, code_info.GetRegisterMaskOf(encoding, stack_map));
 
   ASSERT_TRUE(stack_map.HasDexRegisterMap(encoding.stack_map_encoding));
   DexRegisterMap dex_register_map =
@@ -643,7 +643,7 @@
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64, encoding)));
   ASSERT_EQ(0u, stack_map.GetDexPc(encoding.stack_map_encoding));
   ASSERT_EQ(64u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-  ASSERT_EQ(0x3u, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+  ASSERT_EQ(0x3u, code_info.GetRegisterMaskOf(encoding, stack_map));
 
   ASSERT_FALSE(stack_map.HasDexRegisterMap(encoding.stack_map_encoding));
   ASSERT_FALSE(stack_map.HasInlineInfo(encoding.stack_map_encoding));
@@ -653,7 +653,7 @@
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(68, encoding)));
   ASSERT_EQ(1u, stack_map.GetDexPc(encoding.stack_map_encoding));
   ASSERT_EQ(68u, stack_map.GetNativePcOffset(encoding.stack_map_encoding, kRuntimeISA));
-  ASSERT_EQ(0x4u, stack_map.GetRegisterMask(encoding.stack_map_encoding));
+  ASSERT_EQ(0x4u, code_info.GetRegisterMaskOf(encoding, stack_map));
 
   ASSERT_FALSE(stack_map.HasDexRegisterMap(encoding.stack_map_encoding));
   ASSERT_FALSE(stack_map.HasInlineInfo(encoding.stack_map_encoding));
diff --git a/oatdump/oatdump.cc b/oatdump/oatdump.cc
index ae175eb..0f02da7 100644
--- a/oatdump/oatdump.cc
+++ b/oatdump/oatdump.cc
@@ -590,11 +590,12 @@
       kByteKindCodeInfoEncoding,
       kByteKindCodeInfoOther,
       kByteKindCodeInfoStackMasks,
+      kByteKindCodeInfoRegisterMasks,
       kByteKindStackMapNativePc,
       kByteKindStackMapDexPc,
       kByteKindStackMapDexRegisterMap,
       kByteKindStackMapInlineInfo,
-      kByteKindStackMapRegisterMask,
+      kByteKindStackMapRegisterMaskIndex,
       kByteKindStackMapStackMaskIndex,
       kByteKindCount,
       kByteKindStackMapFirst = kByteKindCodeInfoOther,
@@ -626,43 +627,44 @@
         const int64_t stack_map_bits = std::accumulate(bits + kByteKindStackMapFirst,
                                                        bits + kByteKindStackMapLast + 1,
                                                        0u);
-        Dump(os, "Code                           ", bits[kByteKindCode], sum);
-        Dump(os, "QuickMethodHeader              ", bits[kByteKindQuickMethodHeader], sum);
-        Dump(os, "CodeInfoEncoding               ", bits[kByteKindCodeInfoEncoding], sum);
-        Dump(os, "CodeInfoLocationCatalog        ", bits[kByteKindCodeInfoLocationCatalog], sum);
-        Dump(os, "CodeInfoDexRegisterMap         ", bits[kByteKindCodeInfoDexRegisterMap], sum);
-        Dump(os, "CodeInfoInlineInfo             ", bits[kByteKindCodeInfoInlineInfo], sum);
-        Dump(os, "CodeInfoStackMasks             ", bits[kByteKindCodeInfoStackMasks], sum);
-        Dump(os, "CodeInfoStackMap               ", stack_map_bits, sum);
+        Dump(os, "Code                            ", bits[kByteKindCode], sum);
+        Dump(os, "QuickMethodHeader               ", bits[kByteKindQuickMethodHeader], sum);
+        Dump(os, "CodeInfoEncoding                ", bits[kByteKindCodeInfoEncoding], sum);
+        Dump(os, "CodeInfoLocationCatalog         ", bits[kByteKindCodeInfoLocationCatalog], sum);
+        Dump(os, "CodeInfoDexRegisterMap          ", bits[kByteKindCodeInfoDexRegisterMap], sum);
+        Dump(os, "CodeInfoInlineInfo              ", bits[kByteKindCodeInfoInlineInfo], sum);
+        Dump(os, "CodeInfoStackMasks              ", bits[kByteKindCodeInfoStackMasks], sum);
+        Dump(os, "CodeInfoRegisterMasks           ", bits[kByteKindCodeInfoRegisterMasks], sum);
+        Dump(os, "CodeInfoStackMap                ", stack_map_bits, sum);
         {
           ScopedIndentation indent1(&os);
           Dump(os,
-               "StackMapNativePc             ",
+               "StackMapNativePc              ",
                bits[kByteKindStackMapNativePc],
                stack_map_bits,
                "stack map");
           Dump(os,
-               "StackMapDexPcEncoding        ",
+               "StackMapDexPcEncoding         ",
                bits[kByteKindStackMapDexPc],
                stack_map_bits,
                "stack map");
           Dump(os,
-               "StackMapDexRegisterMap       ",
+               "StackMapDexRegisterMap        ",
                bits[kByteKindStackMapDexRegisterMap],
                stack_map_bits,
                "stack map");
           Dump(os,
-               "StackMapInlineInfo           ",
+               "StackMapInlineInfo            ",
                bits[kByteKindStackMapInlineInfo],
                stack_map_bits,
                "stack map");
           Dump(os,
-               "StackMapRegisterMaskEncoding ",
-               bits[kByteKindStackMapRegisterMask],
+               "StackMapRegisterMaskIndex     ",
+               bits[kByteKindStackMapRegisterMaskIndex],
                stack_map_bits,
                "stack map");
           Dump(os,
-               "StackMapStackMaskIndex       ",
+               "StackMapStackMaskIndex        ",
                bits[kByteKindStackMapStackMaskIndex],
                stack_map_bits,
                "stack map");
@@ -1569,8 +1571,8 @@
               Stats::kByteKindStackMapInlineInfo,
               stack_map_encoding.GetInlineInfoEncoding().BitSize() * num_stack_maps);
           stats_.AddBits(
-              Stats::kByteKindStackMapRegisterMask,
-              stack_map_encoding.GetRegisterMaskEncoding().BitSize() * num_stack_maps);
+              Stats::kByteKindStackMapRegisterMaskIndex,
+              stack_map_encoding.GetRegisterMaskIndexEncoding().BitSize() * num_stack_maps);
           stats_.AddBits(
               Stats::kByteKindStackMapStackMaskIndex,
               stack_map_encoding.GetStackMaskIndexEncoding().BitSize() * num_stack_maps);
@@ -1578,6 +1580,9 @@
               Stats::kByteKindCodeInfoStackMasks,
               helper.GetCodeInfo().GetNumberOfStackMaskBits(encoding) *
                   encoding.number_of_stack_masks);
+          stats_.AddBits(
+              Stats::kByteKindCodeInfoRegisterMasks,
+              encoding.register_mask_size_in_bits * encoding.number_of_stack_masks);
           const size_t stack_map_bytes = helper.GetCodeInfo().GetStackMapsSize(encoding);
           const size_t location_catalog_bytes =
               helper.GetCodeInfo().GetDexRegisterLocationCatalogSize(encoding);
diff --git a/runtime/check_reference_map_visitor.h b/runtime/check_reference_map_visitor.h
index 2252fe7..a955cb5 100644
--- a/runtime/check_reference_map_visitor.h
+++ b/runtime/check_reference_map_visitor.h
@@ -67,7 +67,7 @@
     uint16_t number_of_dex_registers = m->GetCodeItem()->registers_size_;
     DexRegisterMap dex_register_map =
         code_info.GetDexRegisterMapOf(stack_map, encoding, number_of_dex_registers);
-    uint32_t register_mask = stack_map.GetRegisterMask(encoding.stack_map_encoding);
+    uint32_t register_mask = code_info.GetRegisterMaskOf(encoding, stack_map);
     BitMemoryRegion stack_mask = code_info.GetStackMaskOf(encoding, stack_map);
     for (int i = 0; i < number_of_references; ++i) {
       int reg = registers[i];
diff --git a/runtime/oat.h b/runtime/oat.h
index e7f5c7f..532c968 100644
--- a/runtime/oat.h
+++ b/runtime/oat.h
@@ -32,7 +32,7 @@
 class PACKED(4) OatHeader {
  public:
   static constexpr uint8_t kOatMagic[] = { 'o', 'a', 't', '\n' };
-  static constexpr uint8_t kOatVersion[] = { '1', '0', '8', '\0' };  // revert: hash DexCache types
+  static constexpr uint8_t kOatVersion[] = { '1', '0', '9', '\0' };  // Register mask change.
 
   static constexpr const char* kImageLocationKey = "image-location";
   static constexpr const char* kDex2OatCmdLineKey = "dex2oat-cmdline";
diff --git a/runtime/quick_exception_handler.cc b/runtime/quick_exception_handler.cc
index 3ba3011..bf99509 100644
--- a/runtime/quick_exception_handler.cc
+++ b/runtime/quick_exception_handler.cc
@@ -407,7 +407,7 @@
     CodeInfoEncoding encoding = code_info.ExtractEncoding();
     StackMap stack_map = code_info.GetStackMapForNativePcOffset(native_pc_offset, encoding);
     const size_t number_of_vregs = m->GetCodeItem()->registers_size_;
-    uint32_t register_mask = stack_map.GetRegisterMask(encoding.stack_map_encoding);
+    uint32_t register_mask = code_info.GetRegisterMaskOf(encoding, stack_map);
     BitMemoryRegion stack_mask = code_info.GetStackMaskOf(encoding, stack_map);
     DexRegisterMap vreg_map = IsInInlinedFrame()
         ? code_info.GetDexRegisterMapAtDepth(GetCurrentInliningDepth() - 1,
diff --git a/runtime/stack_map.cc b/runtime/stack_map.cc
index f470ae9..4e7c3f4 100644
--- a/runtime/stack_map.cc
+++ b/runtime/stack_map.cc
@@ -97,7 +97,7 @@
       << ", dex_pc_bit_offset=" << static_cast<uint32_t>(dex_pc_bit_offset_)
       << ", dex_register_map_bit_offset=" << static_cast<uint32_t>(dex_register_map_bit_offset_)
       << ", inline_info_bit_offset=" << static_cast<uint32_t>(inline_info_bit_offset_)
-      << ", register_mask_bit_offset=" << static_cast<uint32_t>(register_mask_bit_offset_)
+      << ", register_mask_bit_offset=" << static_cast<uint32_t>(register_mask_index_bit_offset_)
       << ", stack_mask_index_bit_offset=" << static_cast<uint32_t>(stack_mask_index_bit_offset_)
       << ", total_bit_size=" << static_cast<uint32_t>(total_bit_size_)
       << ")\n";
@@ -204,7 +204,7 @@
       << ", native_pc_offset=0x" << pc_offset
       << ", dex_register_map_offset=0x" << GetDexRegisterMapOffset(stack_map_encoding)
       << ", inline_info_offset=0x" << GetInlineDescriptorOffset(stack_map_encoding)
-      << ", register_mask=0x" << GetRegisterMask(stack_map_encoding)
+      << ", register_mask=0x" << code_info.GetRegisterMaskOf(encoding, *this)
       << std::dec
       << ", stack_mask=0b";
   BitMemoryRegion stack_mask = code_info.GetStackMaskOf(encoding, *this);
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 83ba457..062404d 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -694,7 +694,7 @@
                       size_t dex_pc_max,
                       size_t dex_register_map_size,
                       size_t inline_info_size,
-                      size_t register_mask_max,
+                      size_t number_of_register_masks,
                       size_t number_of_stack_masks) {
     total_bit_size_ = 0;
     DCHECK_EQ(kNativePcBitOffset, total_bit_size_);
@@ -716,8 +716,8 @@
       total_bit_size_ += MinimumBitsToStore(dex_register_map_size + inline_info_size);
     }
 
-    register_mask_bit_offset_ = total_bit_size_;
-    total_bit_size_ += MinimumBitsToStore(register_mask_max);
+    register_mask_index_bit_offset_ = total_bit_size_;
+    total_bit_size_ += MinimumBitsToStore(number_of_register_masks);
 
     stack_mask_index_bit_offset_ = total_bit_size_;
     total_bit_size_ += MinimumBitsToStore(number_of_stack_masks);
@@ -735,10 +735,12 @@
     return FieldEncoding(dex_register_map_bit_offset_, inline_info_bit_offset_, -1 /* min_value */);
   }
   ALWAYS_INLINE FieldEncoding GetInlineInfoEncoding() const {
-    return FieldEncoding(inline_info_bit_offset_, register_mask_bit_offset_, -1 /* min_value */);
+    return FieldEncoding(inline_info_bit_offset_,
+                         register_mask_index_bit_offset_,
+                         -1 /* min_value */);
   }
-  ALWAYS_INLINE FieldEncoding GetRegisterMaskEncoding() const {
-    return FieldEncoding(register_mask_bit_offset_, stack_mask_index_bit_offset_);
+  ALWAYS_INLINE FieldEncoding GetRegisterMaskIndexEncoding() const {
+    return FieldEncoding(register_mask_index_bit_offset_, stack_mask_index_bit_offset_);
   }
   ALWAYS_INLINE FieldEncoding GetStackMaskIndexEncoding() const {
     return FieldEncoding(stack_mask_index_bit_offset_, total_bit_size_);
@@ -754,7 +756,7 @@
   uint8_t dex_pc_bit_offset_;
   uint8_t dex_register_map_bit_offset_;
   uint8_t inline_info_bit_offset_;
-  uint8_t register_mask_bit_offset_;
+  uint8_t register_mask_index_bit_offset_;
   uint8_t stack_mask_index_bit_offset_;
   uint8_t total_bit_size_;
 };
@@ -769,7 +771,7 @@
  *
  * The information is of the form:
  *
- *   [native_pc_offset, dex_pc, dex_register_map_offset, inlining_info_offset, register_mask,
+ *   [native_pc_offset, dex_pc, dex_register_map_offset, inlining_info_offset, register_mask_index,
  *   stack_mask_index].
  */
 class StackMap {
@@ -815,12 +817,12 @@
     encoding.GetInlineInfoEncoding().Store(region_, offset);
   }
 
-  ALWAYS_INLINE uint32_t GetRegisterMask(const StackMapEncoding& encoding) const {
-    return encoding.GetRegisterMaskEncoding().Load(region_);
+  ALWAYS_INLINE uint32_t GetRegisterMaskIndex(const StackMapEncoding& encoding) const {
+    return encoding.GetRegisterMaskIndexEncoding().Load(region_);
   }
 
-  ALWAYS_INLINE void SetRegisterMask(const StackMapEncoding& encoding, uint32_t mask) {
-    encoding.GetRegisterMaskEncoding().Store(region_, mask);
+  ALWAYS_INLINE void SetRegisterMaskIndex(const StackMapEncoding& encoding, uint32_t mask) {
+    encoding.GetRegisterMaskIndexEncoding().Store(region_, mask);
   }
 
   ALWAYS_INLINE uint32_t GetStackMaskIndex(const StackMapEncoding& encoding) const {
@@ -1031,7 +1033,9 @@
   uint32_t non_header_size;
   uint32_t number_of_stack_maps;
   uint32_t number_of_stack_masks;
+  uint32_t number_of_register_masks;
   uint32_t stack_mask_size_in_bits;
+  uint32_t register_mask_size_in_bits;
   uint32_t number_of_location_catalog_entries;
   StackMapEncoding stack_map_encoding;
   InlineInfoEncoding inline_info_encoding;
@@ -1044,7 +1048,9 @@
     non_header_size = DecodeUnsignedLeb128(&ptr);
     number_of_stack_maps = DecodeUnsignedLeb128(&ptr);
     number_of_stack_masks = DecodeUnsignedLeb128(&ptr);
+    number_of_register_masks = DecodeUnsignedLeb128(&ptr);
     stack_mask_size_in_bits = DecodeUnsignedLeb128(&ptr);
+    register_mask_size_in_bits = DecodeUnsignedLeb128(&ptr);
     number_of_location_catalog_entries = DecodeUnsignedLeb128(&ptr);
     static_assert(alignof(StackMapEncoding) == 1,
                   "StackMapEncoding should not require alignment");
@@ -1066,7 +1072,9 @@
     EncodeUnsignedLeb128(dest, non_header_size);
     EncodeUnsignedLeb128(dest, number_of_stack_maps);
     EncodeUnsignedLeb128(dest, number_of_stack_masks);
+    EncodeUnsignedLeb128(dest, number_of_register_masks);
     EncodeUnsignedLeb128(dest, stack_mask_size_in_bits);
+    EncodeUnsignedLeb128(dest, register_mask_size_in_bits);
     EncodeUnsignedLeb128(dest, number_of_location_catalog_entries);
     const uint8_t* stack_map_ptr = reinterpret_cast<const uint8_t*>(&stack_map_encoding);
     dest->insert(dest->end(), stack_map_ptr, stack_map_ptr + sizeof(StackMapEncoding));
@@ -1125,10 +1133,13 @@
   }
 
   BitMemoryRegion GetStackMask(const CodeInfoEncoding& encoding, size_t stack_mask_index) const {
-    // All stack mask data is stored at the very end.
+    // All stack mask data is stored before register map data (which is at the very end).
     const size_t entry_size = GetNumberOfStackMaskBits(encoding);
+    const size_t register_mask_bits =
+        encoding.register_mask_size_in_bits * encoding.number_of_register_masks;
     return BitMemoryRegion(region_,
-                           region_.size_in_bits() - entry_size * (stack_mask_index + 1),
+                           region_.size_in_bits() - register_mask_bits -
+                               entry_size * (stack_mask_index + 1),
                            entry_size);
   }
 
@@ -1137,6 +1148,18 @@
     return GetStackMask(encoding, stack_map.GetStackMaskIndex(encoding.stack_map_encoding));
   }
 
+  BitMemoryRegion GetRegisterMask(const CodeInfoEncoding& encoding, size_t index) const {
+    const size_t entry_size = encoding.register_mask_size_in_bits;
+    return BitMemoryRegion(region_,
+                           region_.size_in_bits() - entry_size * (index + 1),
+                           entry_size);
+  }
+
+  uint32_t GetRegisterMaskOf(const CodeInfoEncoding& encoding, const StackMap& stack_map) const {
+    size_t index = stack_map.GetRegisterMaskIndex(encoding.stack_map_encoding);
+    return GetRegisterMask(encoding, index).LoadBits(0u, encoding.register_mask_size_in_bits);
+  }
+
   uint32_t GetNumberOfLocationCatalogEntries(const CodeInfoEncoding& encoding) const {
     return encoding.number_of_location_catalog_entries;
   }
diff --git a/runtime/thread.cc b/runtime/thread.cc
index b36b741..6843e31 100644
--- a/runtime/thread.cc
+++ b/runtime/thread.cc
@@ -3056,7 +3056,7 @@
         }
       }
       // Visit callee-save registers that hold pointers.
-      uint32_t register_mask = map.GetRegisterMask(encoding.stack_map_encoding);
+      uint32_t register_mask = code_info.GetRegisterMaskOf(encoding, map);
       for (size_t i = 0; i < BitSizeOf<uint32_t>(); ++i) {
         if (register_mask & (1 << i)) {
           mirror::Object** ref_addr = reinterpret_cast<mirror::Object**>(GetGPRAddress(i));