[llvm] [CodeGen][SDAG] Remove Combiner WorklistMap (PR #92900)

[via llvm-commits]

https://github.com/aengelke updated https://github.com/llvm/llvm-project/pull/92900

>From 13eff8c73d11595da9c01a0bb6580fea698aa3c9 Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Tue, 21 May 2024 11:52:14 +0200
Subject: [PATCH 1/2] [CodeGen][SDAG] Remove Combiner WorklistMap

DenseMap for pointer lookup is expensive, and this is only used for
deduplication and index lookup. Instead, store the worklist index in the
node itself.
---
 llvm/include/llvm/CodeGen/SelectionDAGNodes.h |  9 ++++++
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 28 ++++++++-----------
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
index ac94c6099d080..7adda059647cb 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
@@ -619,6 +619,9 @@ END_TWO_BYTE_PACK()
   /// Unique id per SDNode in the DAG.
   int NodeId = -1;
 
+  /// Index in worklist of DAGCombiner, or -1.
+  int CombinerWorklistIndex = -1;
+
   /// The values that are used by this operation.
   SDUse *OperandList = nullptr;
 
@@ -747,6 +750,12 @@ END_TWO_BYTE_PACK()
   /// Set unique node id.
   void setNodeId(int Id) { NodeId = Id; }
 
+  /// Get worklist index for DAGCombiner
+  int getCombinerWorklistIndex() const { return CombinerWorklistIndex; }
+
+  /// Set worklist index for DAGCombiner
+  void setCombinerWorklistIndex(int Index) { CombinerWorklistIndex = Index; }
+
   /// Return the node ordering.
   unsigned getIROrder() const { return IROrder; }
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 8607b50175359..90c940e991835 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -170,16 +170,11 @@ namespace {
     /// back and when processing we pop off of the back.
     ///
     /// The worklist will not contain duplicates but may contain null entries
-    /// due to nodes being deleted from the underlying DAG.
+    /// due to nodes being deleted from the underlying DAG. For fast lookup and
+    /// deduplication, the index of the node in this vector is stored in the
+    /// node in SDNode::CombinerWorklistIndex.
     SmallVector<SDNode *, 64> Worklist;
 
-    /// Mapping from an SDNode to its position on the worklist.
-    ///
-    /// This is used to find and remove nodes from the worklist (by nulling
-    /// them) when they are deleted from the underlying DAG. It relies on
-    /// stable indices of nodes within the worklist.
-    DenseMap<SDNode *, unsigned> WorklistMap;
-
     /// This records all nodes attempted to be added to the worklist since we
     /// considered a new worklist entry. As we keep do not add duplicate nodes
     /// in the worklist, this is different from the tail of the worklist.
@@ -238,10 +233,9 @@ namespace {
       }
 
       if (N) {
-        bool GoodWorklistEntry = WorklistMap.erase(N);
-        (void)GoodWorklistEntry;
-        assert(GoodWorklistEntry &&
+        assert(N->getCombinerWorklistIndex() >= 0 &&
                "Found a worklist entry without a corresponding map entry!");
+        N->setCombinerWorklistIndex(-1);
       }
       return N;
     }
@@ -285,8 +279,10 @@ namespace {
       if (IsCandidateForPruning)
         ConsiderForPruning(N);
 
-      if (WorklistMap.insert(std::make_pair(N, Worklist.size())).second)
+      if (N->getCombinerWorklistIndex() == -1) {
+        N->setCombinerWorklistIndex(Worklist.size());
         Worklist.push_back(N);
+      }
     }
 
     /// Remove all instances of N from the worklist.
@@ -295,13 +291,13 @@ namespace {
       PruningList.remove(N);
       StoreRootCountMap.erase(N);
 
-      auto It = WorklistMap.find(N);
-      if (It == WorklistMap.end())
+      int WorklistIndex = N->getCombinerWorklistIndex();
+      if (WorklistIndex == -1)
         return; // Not in the worklist.
 
       // Null out the entry rather than erasing it to avoid a linear operation.
-      Worklist[It->second] = nullptr;
-      WorklistMap.erase(It);
+      Worklist[WorklistIndex] = nullptr;
+      N->setCombinerWorklistIndex(-1);
     }
 
     void deleteAndRecombine(SDNode *N);

>From e640c9726a12938a15c98467683b35e251c2787a Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Mon, 3 Jun 2024 14:49:13 +0200
Subject: [PATCH 2/2] Shuffle SDNode fields

This feels a bit pointless at this stage, but keeps the four padding
bytes in the middle of the struct for future use.
---
 llvm/include/llvm/CodeGen/SelectionDAGNodes.h | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
index 7adda059647cb..fb0bb34ef2e71 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
@@ -480,16 +480,12 @@ class SDNode : public FoldingSetNode, public ilist_node<SDNode> {
   int32_t NodeType;
 
 public:
-  /// Unique and persistent id per SDNode in the DAG. Used for debug printing.
-  /// We do not place that under `#if LLVM_ENABLE_ABI_BREAKING_CHECKS`
-  /// intentionally because it adds unneeded complexity without noticeable
-  /// benefits (see discussion with @thakis in D120714).
-  uint16_t PersistentId = 0xffff;
+  SDNodeFlags Flags;
 
 protected:
   // We define a set of mini-helper classes to help us interpret the bits in our
   // SubclassData.  These are designed to fit within a uint16_t so they pack
-  // with PersistentId.
+  // with SDNodeFlags.
 
 #if defined(_AIX) && (!defined(__GNUC__) || defined(__clang__))
 // Except for GCC; by default, AIX compilers store bit-fields in 4-byte words
@@ -619,8 +615,12 @@ END_TWO_BYTE_PACK()
   /// Unique id per SDNode in the DAG.
   int NodeId = -1;
 
-  /// Index in worklist of DAGCombiner, or -1.
-  int CombinerWorklistIndex = -1;
+  /// Unique and persistent id per SDNode in the DAG. Used for debug printing.
+  /// We do not place that under `#if LLVM_ENABLE_ABI_BREAKING_CHECKS`
+  /// intentionally because it adds unneeded complexity without noticeable
+  /// benefits (see discussion with @thakis in D120714). Currently, there are
+  /// two padding bytes after this field.
+  uint16_t PersistentId = 0xffff;
 
   /// The values that are used by this operation.
   SDUse *OperandList = nullptr;
@@ -649,7 +649,8 @@ END_TWO_BYTE_PACK()
   /// Return a pointer to the specified value type.
   static const EVT *getValueTypeList(EVT VT);
 
-  SDNodeFlags Flags;
+  /// Index in worklist of DAGCombiner, or -1.
+  int CombinerWorklistIndex = -1;
 
   uint32_t CFIType = 0;