[Statepoint] Optimize Location structure size (#78600)

Reduce its size from 24 to 12 bytes. Improves memory consumption when
dealing with statepoint-heavy code.

Author: danilaml

llvm/include/llvm/CodeGen/StackMaps.h

@@ -259,7 +259,7 @@ class StatepointOpers {
 class StackMaps {
 public:
   struct Location {
-    enum LocationType {
+    enum LocationType : unsigned short {
       Unprocessed,
       Register,
       Direct,
@@ -268,12 +268,13 @@ class StackMaps {
       ConstantIndex
     };
     LocationType Type = Unprocessed;
-    unsigned Size = 0;
-    unsigned Reg = 0;
-    int64_t Offset = 0;
+    unsigned short Size = 0;
+    unsigned short Reg = 0;
+    int32_t Offset = 0;
 
     Location() = default;
-    Location(LocationType Type, unsigned Size, unsigned Reg, int64_t Offset)
+    Location(LocationType Type, unsigned short Size, unsigned short Reg,
+             int32_t Offset)
         : Type(Type), Size(Size), Reg(Reg), Offset(Offset) {}
   };
 
@@ -369,7 +370,7 @@ class StackMaps {
   MachineInstr::const_mop_iterator
   parseOperand(MachineInstr::const_mop_iterator MOI,
                MachineInstr::const_mop_iterator MOE, LocationVec &Locs,
-               LiveOutVec &LiveOuts) const;
+               LiveOutVec &LiveOuts);
 
   /// Specialized parser of statepoint operands.
   /// They do not directly correspond to StackMap record entries.

llvm/lib/CodeGen/StackMaps.cpp

@@ -207,7 +207,7 @@ static unsigned getDwarfRegNum(unsigned Reg, const TargetRegisterInfo *TRI) {
 MachineInstr::const_mop_iterator
 StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
                         MachineInstr::const_mop_iterator MOE, LocationVec &Locs,
-                        LiveOutVec &LiveOuts) const {
+                        LiveOutVec &LiveOuts) {
   const TargetRegisterInfo *TRI = AP.MF->getSubtarget().getRegisterInfo();
   if (MOI->isImm()) {
     switch (MOI->getImm()) {
@@ -238,7 +238,22 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
       ++MOI;
       assert(MOI->isImm() && "Expected constant operand.");
       int64_t Imm = MOI->getImm();
-      Locs.emplace_back(Location::Constant, sizeof(int64_t), 0, Imm);
+      if (isInt<32>(Imm)) {
+        Locs.emplace_back(Location::Constant, sizeof(int64_t), 0, Imm);
+      } else {
+        // ConstPool is intentionally a MapVector of 'uint64_t's (as
+        // opposed to 'int64_t's).  We should never be in a situation
+        // where we have to insert either the tombstone or the empty
+        // keys into a map, and for a DenseMap<uint64_t, T> these are
+        // (uint64_t)0 and (uint64_t)-1.  They can be and are
+        // represented using 32 bit integers.
+        assert((uint64_t)Imm != DenseMapInfo<uint64_t>::getEmptyKey() &&
+               (uint64_t)Imm != DenseMapInfo<uint64_t>::getTombstoneKey() &&
+               "empty and tombstone keys should fit in 32 bits!");
+        auto Result = ConstPool.insert(std::make_pair(Imm, Imm));
+        Locs.emplace_back(Location::ConstantIndex, sizeof(int64_t), 0,
+                          Result.first - ConstPool.begin());
+      }
       break;
     }
     }
@@ -497,27 +512,6 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
     while (MOI != MOE)
       MOI = parseOperand(MOI, MOE, Locations, LiveOuts);
 
-  // Move large constants into the constant pool.
-  for (auto &Loc : Locations) {
-    // Constants are encoded as sign-extended integers.
-    // -1 is directly encoded as .long 0xFFFFFFFF with no constant pool.
-    if (Loc.Type == Location::Constant && !isInt<32>(Loc.Offset)) {
-      Loc.Type = Location::ConstantIndex;
-      // ConstPool is intentionally a MapVector of 'uint64_t's (as
-      // opposed to 'int64_t's).  We should never be in a situation
-      // where we have to insert either the tombstone or the empty
-      // keys into a map, and for a DenseMap<uint64_t, T> these are
-      // (uint64_t)0 and (uint64_t)-1.  They can be and are
-      // represented using 32 bit integers.
-      assert((uint64_t)Loc.Offset != DenseMapInfo<uint64_t>::getEmptyKey() &&
-             (uint64_t)Loc.Offset !=
-                 DenseMapInfo<uint64_t>::getTombstoneKey() &&
-             "empty and tombstone keys should fit in 32 bits!");
-      auto Result = ConstPool.insert(std::make_pair(Loc.Offset, Loc.Offset));
-      Loc.Offset = Result.first - ConstPool.begin();
-    }
-  }
-
   // Create an expression to calculate the offset of the callsite from function
   // entry.
   const MCExpr *CSOffsetExpr = MCBinaryExpr::createSub(

llvm/test/CodeGen/X86/statepoint-fixup-undef.mir

@@ -124,12 +124,11 @@ body:             |
   ; STACKMAP-NEXT:	.byte	0
   ; STACKMAP-NEXT:	.short	0
   ; STACKMAP-NEXT:	.long	1
-  ; STACKMAP-NEXT:	.long	1
+  ; STACKMAP-NEXT:	.long	0
   ; STACKMAP-NEXT:	.long	1
   ; STACKMAP-NEXT:	.quad	test_undef
   ; STACKMAP-NEXT:	.quad	88
   ; STACKMAP-NEXT:	.quad	1
-  ; STACKMAP-NEXT:	.quad	4278124286
   ; STACKMAP-NEXT:	.quad	2
   ; STACKMAP-NEXT:	.long	.Ltmp0-test_undef
   ; STACKMAP-NEXT:	.short	0
@@ -182,13 +181,13 @@ body:             |
   ; STACKMAP-NEXT:	.short	3
   ; STACKMAP-NEXT:	.short	0
   ; STACKMAP-NEXT:	.long	0
-  ;      This is entry we're looking for, reference to constant pool entry 0xFEFEFEFE
-  ; STACKMAP-NEXT:	.byte	5
+  ;      This is a constant 0xFEFEFEFE we are looking for
+  ; STACKMAP-NEXT:	.byte	4
   ; STACKMAP-NEXT:	.byte	0
   ; STACKMAP-NEXT:	.short	8
   ; STACKMAP-NEXT:	.short	0
   ; STACKMAP-NEXT:	.short	0
-  ; STACKMAP-NEXT:	.long	0
+  ; STACKMAP-NEXT:	.long	-16843010
   ; STACKMAP-NEXT:	.byte	4
   ; STACKMAP-NEXT:	.byte	0
   ; STACKMAP-NEXT:	.short	8