Store one register inside GenTree

jakobbotsch · Apr 27, 2024 · 426fdee · 426fdee
1 parent 17486f5
commit 426fdee
Show file tree

Hide file tree

Showing 8 changed files with 66 additions and 14 deletions.
diff --git a/src/coreclr/jit/codegenarmarch.cpp b/src/coreclr/jit/codegenarmarch.cpp
@@ -3376,7 +3376,7 @@ void CodeGen::genCall(GenTreeCall* call)
                 call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
             GetEmitter()->emitIns_R_R(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg);
             // We will use this again when emitting the jump in genCallInstruction in the epilog
-            internalRegisters.Add(call, genRegMask(tmpReg));
+            internalRegisters.Add(call, tmpReg);
         }
 #endif
 

diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
@@ -67,25 +67,36 @@ NodeInternalRegisters::NodeInternalRegisters(Compiler* comp)
 {
 }
 
+void NodeInternalRegisters::Clear(GenTree* tree)
+{
+    tree->m_internalRegTag = 0;
+}
+
 //------------------------------------------------------------------------
 // Add: Add internal allocated registers for the specified node.
 //
 // Parameters:
 //   tree - IR node to add internal allocated registers to
 //   regs - Registers to add
 //
-void NodeInternalRegisters::Add(GenTree* tree, regMaskTP regs)
+void NodeInternalRegisters::Add(GenTree* tree, regNumber reg)
 {
     assert(regs != RBM_NONE);
 
-    regMaskTP* result = m_table.LookupPointer(tree);
-    if (result == nullptr)
+    if (tree->m_internalRegTag == 0)
     {
-        m_table.Set(tree, regs);
+        tree->m_internalRegTag = (uint8_t)(1 + reg);
+    }
+    else if (tree->m_internalRegTag != 0xFF)
+    {
+        regMaskTP newRegs = genRegMask(reg) | genRegMask((regNumber)(tree->m_internalRegTag - 1));
+        m_table.Set(tree, newRegs);
+        tree->m_internalRegTag = 0xFF;
     }
     else
     {
-        *result |= regs;
+        regMaskTP* oldRegs = m_table.LookupPointer(tree);
+        *oldRegs |= genRegMask(reg);
     }
 }
 
@@ -104,6 +115,15 @@ void NodeInternalRegisters::Add(GenTree* tree, regMaskTP regs)
 //
 regNumber NodeInternalRegisters::Extract(GenTree* tree, regMaskTP mask)
 {
+    assert(tree->m_internalRegTag != 0);
+    if (tree->m_internalRegTag != 0xFF)
+    {
+        regNumber result = (regNumber)(tree->m_internalRegTag - 1);
+        assert((mask & genRegMask(result)) != RBM_NONE);
+        tree->m_internalRegTag = 0;
+        return result;
+    }
+
     regMaskTP* regs = m_table.LookupPointer(tree);
     assert(regs != nullptr);
 
@@ -131,6 +151,14 @@ regNumber NodeInternalRegisters::Extract(GenTree* tree, regMaskTP mask)
 //
 regNumber NodeInternalRegisters::GetSingle(GenTree* tree, regMaskTP mask)
 {
+    assert(tree->m_internalRegTag != 0);
+    if (tree->m_internalRegTag != 0xFF)
+    {
+        regNumber result = (regNumber)(tree->m_internalRegTag - 1);
+        assert((mask & genRegMask(result)) != RBM_NONE);
+        return result;
+    }
+
     regMaskTP* regs = m_table.LookupPointer(tree);
     assert(regs != nullptr);
 
@@ -154,8 +182,17 @@ regNumber NodeInternalRegisters::GetSingle(GenTree* tree, regMaskTP mask)
 //
 regMaskTP NodeInternalRegisters::GetAll(GenTree* tree)
 {
-    regMaskTP regs;
-    return m_table.Lookup(tree, &regs) ? regs : RBM_NONE;
+    if (tree->m_internalRegTag == 0)
+    {
+        return RBM_NONE;
+    }
+
+    if (tree->m_internalRegTag != 0xFF)
+    {
+        return genRegMask((regNumber)(tree->m_internalRegTag - 1));
+    }
+
+    return m_table[tree];
 }
 
 //------------------------------------------------------------------------
@@ -171,8 +208,17 @@ regMaskTP NodeInternalRegisters::GetAll(GenTree* tree)
 //
 unsigned NodeInternalRegisters::Count(GenTree* tree, regMaskTP mask)
 {
-    regMaskTP regs;
-    return m_table.Lookup(tree, &regs) ? genCountBits(regs & mask) : 0;
+    if (tree->m_internalRegTag == 0)
+    {
+        return 0;
+    }
+
+    if (tree->m_internalRegTag != 0xFF)
+    {
+        return (mask & genRegMask((regNumber)(tree->m_internalRegTag - 1))) != RBM_NONE ? 1 : 0;
+    }
+
+    return genCountBits(m_table[tree] & mask);
 }
 
 // CodeGen constructor

diff --git a/src/coreclr/jit/codegeninterface.h b/src/coreclr/jit/codegeninterface.h
@@ -54,7 +54,8 @@ class NodeInternalRegisters
 public:
     NodeInternalRegisters(Compiler* comp);
 
-    void      Add(GenTree* tree, regMaskTP reg);
+    void      Clear(GenTree* tree);
+    void      Add(GenTree* tree, regNumber reg);
     regNumber Extract(GenTree* tree, regMaskTP mask = static_cast<regMaskTP>(-1));
     regNumber GetSingle(GenTree* tree, regMaskTP mask = static_cast<regMaskTP>(-1));
     regMaskTP GetAll(GenTree* tree);

diff --git a/src/coreclr/jit/codegenloongarch64.cpp b/src/coreclr/jit/codegenloongarch64.cpp
@@ -6255,7 +6255,7 @@ void CodeGen::genCall(GenTreeCall* call)
                 call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
             GetEmitter()->emitIns_R_R_I(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg, 0);
             // We will use this again when emitting the jump in genCallInstruction in the epilog
-            internalRegisters.Add(call, genRegMask(tmpReg));
+            internalRegisters.Add(call, tmpReg);
         }
 #endif
 

diff --git a/src/coreclr/jit/codegenriscv64.cpp b/src/coreclr/jit/codegenriscv64.cpp
@@ -6399,7 +6399,7 @@ void CodeGen::genCall(GenTreeCall* call)
                 call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
             GetEmitter()->emitIns_R_R_I(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg, 0);
             // We will use this again when emitting the jump in genCallInstruction in the epilog
-            internalRegisters.Add(call, genRegMask(tmpReg));
+            internalRegisters.Add(call, tmpReg);
         }
 #endif
 

diff --git a/src/coreclr/jit/gentree.h b/src/coreclr/jit/gentree.h
@@ -853,6 +853,9 @@ struct GenTree
     friend struct GenTreeMultiOp;
     uint8_t m_operandCount;
 
+    friend class NodeInternalRegisters;
+    uint8_t m_internalRegTag;
+
 public:
     // The register number is stored in a small format (8 bits), but the getters return and the setters take
     // a full-size (unsigned) format, to localize the casts here.

diff --git a/src/coreclr/jit/lsra.cpp b/src/coreclr/jit/lsra.cpp
@@ -8106,7 +8106,7 @@ void LinearScan::resolveRegisters()
             assert(currentRefPosition->isIntervalRef());
             if (currentRefPosition->getInterval()->isInternal)
             {
-                compiler->codeGen->internalRegisters.Add(treeNode, currentRefPosition->registerAssignment);
+                compiler->codeGen->internalRegisters.Add(treeNode, currentRefPosition->assignedReg());
             }
             else
             {

diff --git a/src/coreclr/jit/lsrabuild.cpp b/src/coreclr/jit/lsrabuild.cpp
@@ -1730,6 +1730,8 @@ int LinearScan::ComputeAvailableSrcCount(GenTree* node)
 //
 void LinearScan::buildRefPositionsForNode(GenTree* tree, LsraLocation currentLoc)
 {
+    compiler->codeGen->internalRegisters.Clear(tree);
+
 #ifdef DEBUG
     if (VERBOSE)
     {