[CodeGen] Allow negative frame indicies in Register class. (#164459)

The register values between `2 << 30` (inclusive) and `2 << 31`
(exclusive) correspond to frame indices. To obtain the frame index from
the given register value we interpret first 30 bits as an unsigned
integer. Thus, currently only non-negative frame indices can be
represented.

However, we should also be able to represent negative frame indices as
register values as well. This is used by reaching definitions analysis
for example.

In order to do that, we interpret the first 30 bits of the register
value as a signed integer.

---------

Co-authored-by: Mikhail Gudim <[email protected]>
Co-authored-by: Petr Penzin <[email protected]>

Author: 3 people

llvm/include/llvm/CodeGen/Register.h

@@ -10,6 +10,7 @@
 #define LLVM_CODEGEN_REGISTER_H
 
 #include "llvm/MC/MCRegister.h"
+#include "llvm/Support/MathExtras.h"
 #include <cassert>
 
 namespace llvm {
@@ -35,19 +36,23 @@ class Register {
   // DenseMapInfo<unsigned> uses -1u and -2u.
   static_assert(std::numeric_limits<decltype(Reg)>::max() >= 0xFFFFFFFF,
                 "Reg isn't large enough to hold full range.");
-  static constexpr unsigned FirstStackSlot = 1u << 30;
-  static_assert(FirstStackSlot >= MCRegister::LastPhysicalReg);
+  static constexpr unsigned MaxFrameIndexBitwidth = 30;
+  static constexpr unsigned StackSlotZero = 1u << MaxFrameIndexBitwidth;
+  static constexpr const unsigned StackSlotMask = StackSlotZero - 1;
+  static_assert(StackSlotZero >= MCRegister::LastPhysicalReg);
   static constexpr unsigned VirtualRegFlag = 1u << 31;
 
   /// Return true if this is a stack slot.
   constexpr bool isStack() const {
-    return Register::FirstStackSlot <= Reg && Reg < Register::VirtualRegFlag;
+    return Register::StackSlotZero <= Reg && Reg < Register::VirtualRegFlag;
   }
 
   /// Convert a non-negative frame index to a stack slot register value.
   static Register index2StackSlot(int FI) {
-    assert(FI >= 0 && "Cannot hold a negative frame index.");
-    return Register(FI + Register::FirstStackSlot);
+    assert(isInt<MaxFrameIndexBitwidth>(FI) &&
+           "Frame index must be at most 30 bits.");
+    unsigned FIMasked = FI & Register::StackSlotMask;
+    return Register(FIMasked | Register::StackSlotZero);
   }
 
   /// Return true if the specified register number is in
@@ -87,7 +92,7 @@ class Register {
   /// Compute the frame index from a register value representing a stack slot.
   int stackSlotIndex() const {
     assert(isStack() && "Not a stack slot");
-    return static_cast<int>(Reg - Register::FirstStackSlot);
+    return SignExtend32<MaxFrameIndexBitwidth>(Reg & Register::StackSlotMask);
   }
 
   constexpr operator unsigned() const { return Reg; }

llvm/lib/CodeGen/ReachingDefAnalysis.cpp

@@ -193,7 +193,6 @@ void ReachingDefInfo::processDefs(MachineInstr *MI) {
   for (auto &MO : MI->operands()) {
     if (MO.isFI()) {
       int FrameIndex = MO.getIndex();
-      assert(FrameIndex >= 0 && "Can't handle negative frame indicies yet!");
       if (!isFIDef(*MI, FrameIndex, TII))
         continue;
       MBBFrameObjsReachingDefs[{MBBNumber, FrameIndex}].push_back(CurInstr);
@@ -302,8 +301,6 @@ void ReachingDefInfo::print(raw_ostream &OS) {
         Register Reg;
         if (MO.isFI()) {
           int FrameIndex = MO.getIndex();
-          assert(FrameIndex >= 0 &&
-                 "Can't handle negative frame indicies yet!");
           Reg = Register::index2StackSlot(FrameIndex);
         } else if (MO.isReg()) {
           if (MO.isDef())

llvm/unittests/CodeGen/CMakeLists.txt

@@ -39,6 +39,7 @@ add_llvm_unittest(CodeGenTests
   MachineOperandTest.cpp
   MIR2VecTest.cpp
   RegAllocScoreTest.cpp
+  RegisterTest.cpp
   PassManagerTest.cpp
   ScalableVectorMVTsTest.cpp
   SchedBoundary.cpp

llvm/unittests/CodeGen/RegisterTest.cpp

@@ -0,0 +1,38 @@
+//===- RegisterTest.cpp -----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/Register.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+TEST(RegisterTest, Idx2StackSlot) {
+  EXPECT_EQ(Register::index2StackSlot(0), Register::StackSlotZero);
+  EXPECT_EQ(Register::index2StackSlot(1), Register::StackSlotZero | 1);
+  EXPECT_EQ(Register::index2StackSlot(-1),
+            Register::StackSlotZero | Register::StackSlotMask);
+  int MaxPowOf2 = 1 << (Register::MaxFrameIndexBitwidth - 1);
+  // Check the highest possible value of frame index
+  EXPECT_EQ(Register::index2StackSlot(MaxPowOf2 - 1),
+            Register::StackSlotZero | (MaxPowOf2 - 1));
+  // Check the lowest possible value of frame index
+  EXPECT_EQ(Register::index2StackSlot(-MaxPowOf2),
+            Register::StackSlotZero | (-MaxPowOf2 & Register::StackSlotMask));
+}
+
+TEST(RegisterTest, StackSlotIndex) {
+  int MaxPowOf2 = 1 << (Register::MaxFrameIndexBitwidth - 1);
+  std::vector<int> FIs = {0, 1 - 1, MaxPowOf2 - 1, -MaxPowOf2};
+
+  for (int FI : FIs) {
+    Register Reg = Register::index2StackSlot(FI);
+    EXPECT_EQ(Reg.stackSlotIndex(), FI);
+  }
+}
+} // end namespace