AArch64: Implement Register Association

Implement Register Association for aarch64.

- Introduce `assocreg` pseudo instruction.
- Keep track of the real register to virtual register mappings
  which are set up by register dependencies.
- Emit `assocreg` instructions at the end of every extended basic block.
- Emit `assocreg` instructions when setting up register dependencies.
- In register assignment phase, update virtual register to real register mappings
  when we encounter an `assocreg` instruction.

Signed-off-by: Akira Saitoh <saiaki@jp.ibm.com>

compiler/aarch64/codegen/ARM64Debug.cpp

@@ -499,7 +499,8 @@ static const char *opCodeToNameMap[] =
    "return",
    "dd",
    "label",
-   "vgdnop"
+   "vgdnop",
+   "assocreg"
    };
 
 const char *
@@ -796,6 +797,12 @@ TR_Debug::print(TR::FILE *pOutFile, TR::ARM64RegBranchInstruction *instr)
 void
 TR_Debug::print(TR::FILE *pOutFile, TR::ARM64AdminInstruction *instr)
    {
+   if (instr->getOpCodeValue() == TR::InstOpCode::assocreg)
+      {
+      printAssocRegDirective(pOutFile, instr);
+      return;
+      }
+
    printPrefix(pOutFile, instr);
    trfprintf(pOutFile, "%s ", getOpCodeName(&instr->getOpCode()));
 
@@ -1489,6 +1496,30 @@ TR_Debug::print(TR::FILE *pOutFile, TR::RegisterDependencyConditions *conditions
       }
    }
 
+void
+TR_Debug::printAssocRegDirective(TR::FILE *pOutFile, TR::Instruction *instr)
+   {
+   TR_ARM64RegisterDependencyGroup * depGroup = instr->getDependencyConditions()->getPostConditions();
+
+   printPrefix(pOutFile, instr);
+   trfprintf(pOutFile, "%s", getOpCodeName(&instr->getOpCode()));
+   trfflush(pOutFile);
+
+   auto numPostConditions = instr->getDependencyConditions()->getAddCursorForPost();
+   for (int i = 0; i < numPostConditions; i++)
+      {
+      TR::RegisterDependency *dependency = depGroup->getRegisterDependency(i);
+      TR::Register *virtReg = dependency->getRegister();
+
+      if (virtReg)
+         {
+         print(pOutFile, dependency);
+         }
+      }
+
+   trfflush(pOutFile);
+   }
+
 void
 TR_Debug::print(TR::FILE *pOutFile, TR::MemoryReference *mr)
    {

compiler/aarch64/codegen/ARM64Instruction.cpp

@@ -91,6 +91,49 @@ void TR::ARM64LabelInstruction::assignRegisters(TR_RegisterKinds kindToBeAssigne
    assignRegistersForOutOfLineCodeSection(kindToBeAssigned);
    }
 
+// TR::ARM64AdminInstruction member functions
+
+void TR::ARM64AdminInstruction::assignRegisters(TR_RegisterKinds kindToBeAssigned)
+   {
+   if (getOpCodeValue() != TR::InstOpCode::assocreg)
+      {
+      OMR::ARM64::Instruction::assignRegisters(kindToBeAssigned);
+      }
+   else if (self()->cg()->enableRegisterAssociations())
+      {
+      TR::Machine *machine = self()->cg()->machine();
+
+      int32_t first = TR::RealRegister::FirstGPR;
+      int32_t last  = TR::RealRegister::LastAssignableFPR;
+      // Step 1 : First traverse the existing associations and remove them so that they don't interfere with the new ones
+      for (int32_t i = first; i <= last; ++i)
+         {
+         TR::Register *virtReg = machine->getVirtualAssociatedWithReal(static_cast<TR::RealRegister::RegNum>(i));
+         if (virtReg)
+            {
+            virtReg->setAssociation(TR::RealRegister::NoReg);
+            }
+         }
+
+      // Step 2 : loop through and set up the new associations (both on the machine and by associating the virtual
+      // registers with their real dependencies)
+      auto depCond = self()->getDependencyConditions();
+      auto depGroup = depCond->getPostConditions();
+      auto numPostCond = depCond->getNumPostConditions();
+      for (int32_t j = 0; j < numPostCond; ++j)
+         {
+         auto registerDependency = depGroup->getRegisterDependency(j);
+         TR::RealRegister::RegNum regNum = registerDependency->getRealRegister();
+         TR::Register *virtReg = registerDependency->getRegister();
+
+         machine->setVirtualAssociatedWithReal(regNum, virtReg);
+         }
+
+      // Set up virtual to real associations and set register weights
+      machine->setRegisterWeightsFromAssociations();
+      }
+   }
+
 // TR::ARM64Trg1Instruction:: member functions
 
 bool TR::ARM64Trg1Instruction::refsRegister(TR::Register *reg)

compiler/aarch64/codegen/ARM64Instruction.hpp

@@ -1066,6 +1066,12 @@ class ARM64AdminInstruction : public TR::Instruction
     */
    TR::Node *getFenceNode() { return _fenceNode; }
 
+   /**
+    * @brief Assigns registers
+    * @param[in] kindToBeAssigned : register kind
+    */
+   virtual void assignRegisters(TR_RegisterKinds kindToBeAssigned);
+
    /**
     * @brief Generates binary encoding of the instruction
     * @return instruction cursor

compiler/aarch64/codegen/OMRCodeGenerator.hpp

@@ -458,6 +458,12 @@ class OMR_EXTENSIBLE CodeGenerator : public OMR::CodeGenerator
     */
    TR::Instruction *generateNop(TR::Node *node, TR::Instruction *preced = 0);
 
+   /**
+    * @brief Determines whether register association is enabled
+    * @return true if register association is enabled
+    */
+   bool enableRegisterAssociations() { return true; }
+
    /**
     * @brief Returns bit mask for real register
     * @param[in] reg: real register number

compiler/aarch64/codegen/OMRInstOpCodeEnum.hpp

@@ -481,5 +481,6 @@
 		dd,    // Define word
 		label, // Destination of a jump
 		vgdnop, // Virtual Guard NOP instruction
-		ARM64LastOp = vgdnop,
+		assocreg, // Associate real registers with Virtual registers.
+		ARM64LastOp = assocreg,
 		ARM64NumOpCodes = ARM64LastOp+1,

compiler/aarch64/codegen/OMRMachine.cpp

@@ -26,17 +26,26 @@
 #include "codegen/BackingStore.hpp"
 #include "codegen/CodeGenerator.hpp"
 #include "codegen/GenerateInstructions.hpp"
+#include "codegen/Linkage.hpp"
 #include "codegen/Machine.hpp"
 #include "codegen/Machine_inlines.hpp"
 #include "codegen/RealRegister.hpp"
 #include "il/Node.hpp"
 #include "il/Node_inlines.hpp"
 #include "infra/Assert.hpp"
 
+// Register Association
+#define ARM64_REGISTER_HEAVIEST_WEIGHT               0x0000ffff
+#define ARM64_REGISTER_INITIAL_PRESERVED_WEIGHT      0x00001000
+#define ARM64_REGISTER_ASSOCIATED_WEIGHT             0x00000800
+#define ARM64_REGISTER_PLACEHOLDER_WEIGHT            0x00000100
+#define ARM64_REGISTER_BASIC_WEIGHT                  0x00000080
+
 OMR::ARM64::Machine::Machine(TR::CodeGenerator *cg) :
       OMR::Machine(cg)
    {
    self()->initializeRegisterFile();
+   self()->clearRegisterAssociations();
    }
 
 TR::RealRegister *OMR::ARM64::Machine::findBestFreeRegister(TR_RegisterKinds rk,
@@ -1148,6 +1157,91 @@ void OMR::ARM64::Machine::initializeRegisterFile()
                                                  self()->cg());
    }
 
+
+// Register Association ////////////////////////////////////////////
+void
+OMR::ARM64::Machine::setRegisterWeightsFromAssociations()
+   {
+   TR::ARM64LinkageProperties linkageProperties = self()->cg()->getProperties();
+   int32_t first = TR::RealRegister::FirstGPR;
+   TR::Compilation *comp = self()->cg()->comp();
+   int32_t last = TR::RealRegister::LastAssignableFPR;
+
+   for (int32_t i = first; i <= last; ++i)
+      {
+      TR::Register * assocReg = getVirtualAssociatedWithReal(static_cast<TR::RealRegister::RegNum>(i));
+      if (linkageProperties.getPreserved(static_cast<TR::RealRegister::RegNum>(i)) && _registerFile[i]->getHasBeenAssignedInMethod() == false)
+         {
+         if (assocReg)
+            {
+            assocReg->setAssociation(i);
+            }
+         _registerFile[i]->setWeight(ARM64_REGISTER_INITIAL_PRESERVED_WEIGHT);
+         }
+      else if (assocReg == NULL)
+         {
+         _registerFile[i]->setWeight(ARM64_REGISTER_BASIC_WEIGHT);
+         }
+      else
+         {
+         assocReg->setAssociation(i);
+         if (assocReg->isPlaceholderReg())
+            {
+            // placeholder register and is only needed at the specific dependency
+            // site (usually a killed register on a call)
+            // so defer this register's weight to that of registers
+            // where the associated register has a longer life
+            _registerFile[i]->setWeight(ARM64_REGISTER_PLACEHOLDER_WEIGHT);
+            }
+         else
+            {
+            _registerFile[i]->setWeight(ARM64_REGISTER_ASSOCIATED_WEIGHT);
+            }
+         }
+      }
+   }
+
+void
+OMR::ARM64::Machine::createRegisterAssociationDirective(TR::Instruction *cursor)
+   {
+   TR::Compilation *comp = self()->cg()->comp();
+   int32_t first = TR::RealRegister::FirstGPR;
+   int32_t last = TR::RealRegister::LastAssignableFPR;
+   TR::RegisterDependencyConditions *associations  = new (self()->cg()->trHeapMemory()) TR::RegisterDependencyConditions(0, last, self()->cg()->trMemory());
+
+   // Go through the current associations held in the machine and put a copy of
+   // that state out into the stream after the cursor
+   // so that when the register assigner goes backwards through this point
+   // it updates the machine and register association states properly
+   //
+   for (int32_t i = first; i <= last; i++)
+      {
+      TR::RealRegister::RegNum regNum = static_cast<TR::RealRegister::RegNum>(i);
+      associations->addPostCondition(self()->getVirtualAssociatedWithReal(regNum), regNum);
+      }
+
+   generateAdminInstruction(self()->cg(), TR::InstOpCode::assocreg, cursor->getNode(), associations, NULL, cursor);
+   }
+
+TR::Register *
+OMR::ARM64::Machine::setVirtualAssociatedWithReal(TR::RealRegister::RegNum regNum, TR::Register *virtReg)
+   {
+   if (virtReg)
+      {
+      // disable previous association
+      if (virtReg->getAssociation() && (_registerAssociations[virtReg->getAssociation()] == virtReg))
+         _registerAssociations[virtReg->getAssociation()] = NULL;
+
+      if (regNum == TR::RealRegister::NoReg || regNum == TR::RealRegister::xzr || regNum >= TR::RealRegister::NumRegisters)
+         {
+         virtReg->setAssociation(TR::RealRegister::NoReg);
+         return NULL;
+         }
+      }
+
+   return _registerAssociations[regNum] = virtReg;
+   }
+
 void
 OMR::ARM64::Machine::takeRegisterStateSnapShot()
    {
@@ -1157,6 +1251,8 @@ OMR::ARM64::Machine::takeRegisterStateSnapShot()
       _registerStatesSnapShot[i] = _registerFile[i]->getState();
       _assignedRegisterSnapShot[i] = _registerFile[i]->getAssignedRegister();
       _registerFlagsSnapShot[i] = _registerFile[i]->getFlags();
+      _registerAssociationsSnapShot[i] = self()->getVirtualAssociatedWithReal(static_cast<TR::RealRegister::RegNum>(i));
+      _registerWeightSnapShot[i] = _registerFile[i]->getWeight();
       }
    }
 
@@ -1166,8 +1262,11 @@ OMR::ARM64::Machine::restoreRegisterStateFromSnapShot()
    int32_t i;
    for (i = TR::RealRegister::FirstGPR; i < TR::RealRegister::NumRegisters - 1; i++) // Skipping SpilledReg
       {
+      _registerFile[i]->setWeight(_registerWeightSnapShot[i]);
       _registerFile[i]->setFlags(_registerFlagsSnapShot[i]);
       _registerFile[i]->setState(_registerStatesSnapShot[i]);
+      self()->setVirtualAssociatedWithReal(static_cast<TR::RealRegister::RegNum>(i), _registerAssociationsSnapShot[i]);
+
       if (_registerFile[i]->getState() == TR::RealRegister::Free)
          {
          if (_registerFile[i]->getAssignedRegister() != NULL)

compiler/aarch64/codegen/OMRMachine.hpp

@@ -157,12 +157,60 @@ class OMR_EXTENSIBLE Machine : public OMR::Machine
     */
    void decFutureUseCountAndUnlatch(TR::Instruction *currentInstruction, TR::Register *virtualRegister);
 
+   // Register Association Stuff ///////////////
+
+   /**
+    * @brief Sets register weights from register associations
+    *
+    */
+   void setRegisterWeightsFromAssociations();
+
+   /**
+    * @brief Creates a regassoc pseudo instruction
+    *
+    */
+   void createRegisterAssociationDirective(TR::Instruction *cursor);
+
+   /**
+    * @brief Returns a virtual register associted to the passed real register.
+    *
+    * @param regNum                 : real register number
+    */
+   TR::Register *getVirtualAssociatedWithReal(TR::RealRegister::RegNum regNum)
+      {
+      return _registerAssociations[regNum];
+      }
+
+   /**
+    * @brief Associates the virtual register to the real register.
+    *
+    * @param regNum                 : real register number
+    * @param virtReg                : virtual register
+    * @returns virtual register
+    */
+   TR::Register *setVirtualAssociatedWithReal(TR::RealRegister::RegNum regNum, TR::Register *virtReg);
+
+
+   /**
+    * @brief Clears register association
+    *
+    */
+   void clearRegisterAssociations()
+      {
+      memset(_registerAssociations, 0, sizeof(TR::Register *) * (TR::RealRegister::NumRegisters));
+      }
+
 private:
 
    // For register snap shot
    uint16_t                   _registerFlagsSnapShot[TR::RealRegister::NumRegisters];
    TR::RealRegister::RegState _registerStatesSnapShot[TR::RealRegister::NumRegisters];
    TR::Register               *_assignedRegisterSnapShot[TR::RealRegister::NumRegisters];
+   TR::Register               *_registerAssociationsSnapShot[TR::RealRegister::NumRegisters];
+   uint16_t                   _registerWeightSnapShot[TR::RealRegister::NumRegisters];
+
+   // register association
+   TR::Register               *_registerAssociations[TR::RealRegister::NumRegisters];
 
    void initializeRegisterFile();
    };

compiler/aarch64/codegen/OMRRegisterDependency.cpp

@@ -222,6 +222,40 @@ bool OMR::ARM64::RegisterDependencyConditions::usesRegister(TR::Register *r)
 
 void OMR::ARM64::RegisterDependencyConditions::bookKeepingRegisterUses(TR::Instruction *instr, TR::CodeGenerator *cg)
    {
+   if (instr->getOpCodeValue() == TR::InstOpCode::assocreg)
+      return;
+
+   // We create a register association directive for each dependency
+
+   if (cg->enableRegisterAssociations() && !cg->isOutOfLineColdPath())
+      {
+      TR::Machine *machine = cg->machine();
+
+      machine->createRegisterAssociationDirective(instr->getPrev());
+
+      // Now set up the new register associations as required by the current
+      // dependent register instruction onto the machine.
+      // Only the registers that this instruction interferes with are modified.
+      //
+      for (int i = 0; i < _addCursorForPre; i++)
+         {
+         TR::RegisterDependency *dependency = _preConditions->getRegisterDependency(i);
+         if (dependency->getRegister())
+            {
+            machine->setVirtualAssociatedWithReal(dependency->getRealRegister(), dependency->getRegister());
+            }
+         }
+
+      for (int j = 0; j < _addCursorForPost; j++)
+         {
+         TR::RegisterDependency *dependency = _postConditions->getRegisterDependency(j);
+         if (dependency->getRegister())
+            {
+            machine->setVirtualAssociatedWithReal(dependency->getRealRegister(), dependency->getRegister());
+            }
+         }
+      }
+
    for (int i = 0; i < _addCursorForPre; i++)
       {
       instr->useRegister(_preConditions->getRegisterDependency(i)->getRegister());

compiler/aarch64/codegen/OMRRegisterDependency.hpp

@@ -157,7 +157,8 @@ class TR_ARM64RegisterDependencyGroup
       {
       for (uint32_t i = 0; i < numberOfRegisters; i++)
          {
-         _dependencies[i].getRegister()->block();
+         if (_dependencies[i].getRegister())
+            _dependencies[i].getRegister()->block();
          }
       }
 
@@ -169,7 +170,8 @@ class TR_ARM64RegisterDependencyGroup
       {
       for (uint32_t i = 0; i < numberOfRegisters; i++)
          {
-         _dependencies[i].getRegister()->unblock();
+         if (_dependencies[i].getRegister())
+            _dependencies[i].getRegister()->unblock();
          }
       }
    };