[Statepoint] Optimize Location structure size

[danilaml]

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

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Danila Malyutin (danilaml)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/78600.diff

3 Files Affected:

• (modified) llvm/include/llvm/CodeGen/StackMaps.h (+7-6)
• (modified) llvm/lib/CodeGen/StackMaps.cpp (+18-23)
• (modified) llvm/test/CodeGen/X86/statepoint-fixup-undef.mir (+4-5)

diff --git a/llvm/include/llvm/CodeGen/StackMaps.h b/llvm/include/llvm/CodeGen/StackMaps.h
index 467e31f17bc82b..ca5dfa5666003c 100644
--- a/llvm/include/llvm/CodeGen/StackMaps.h
+++ b/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.
diff --git a/llvm/lib/CodeGen/StackMaps.cpp b/llvm/lib/CodeGen/StackMaps.cpp
index f9115e43487844..32453091fb6b5c 100644
--- a/llvm/lib/CodeGen/StackMaps.cpp
+++ b/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,23 @@ 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 +513,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(
diff --git a/llvm/test/CodeGen/X86/statepoint-fixup-undef.mir b/llvm/test/CodeGen/X86/statepoint-fixup-undef.mir
index 0adccba88a3271..30a68e6c2efd2a 100644
--- a/llvm/test/CodeGen/X86/statepoint-fixup-undef.mir
+++ b/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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[danilaml]

Do this eagerly, that way we can reduce Offset size from 8 down to 4 bytes.

[danilaml]

Not sure about sizeof(int64_t). I think the size field is not really used for constants, since they are always 4 bytes, so I just kept it as 8 to reduce tests diff.

[danilaml]

0xFEFEFEFE is no longer sign-extended to int64_t and as such doesn't require a constant pool entry. Since, as I understand, this was added mostly to not crash and it doesn't havy specific meaning, I think we are fine with this change.

@arsenm I saw that you had a stack that removed this special handling completely, although it didn't go anywhere, even though some patches were approved https://reviews.llvm.org/D122605

[danilaml]

This struct can be made even smaller (8 bytes) if we borrow some bits from size (or Reg), but that'd probably require changing StackMap format, since it allows for full 16 bits for size, even though it's not really useful in practice for most cases.

[dantrushin]

Looks good.
Suggestion for future improvement:
Change return type of getDwarfRegNumber (in StackMaps.cpp) to uint16_t and add assert that number it gets from TRI really fits in uint16_t

[dantrushin]

Does it makes sense to use uint16_t instead of unsigned short?

[danilaml]

Just followed the types used in LiveOutReg struct below for consistency. Can change both to uint16_t.