[SelectionDAG] Add getTokenFactor, which splits nodes with > 64k operands.

[fhahn]

Summary

This functionality is required at multiple places which potentially
create large operand lists, like SelectionDAGBuilder or DAGCombiner.

Also updates the limit computation to use SDNode::getMaxNumOperands.

[efriedma]

llvm/include/llvm/CodeGen/SelectionDAG.h

1161

This API seems strange... could we just add a "getTokenFactor" API instead?

[RKSimon]

@fhahn Are you intending to look at the similar SDNode::NumValues limit? PR7250 + PR37000 ?

[fhahn]

Thanks Eli, changed to getTokenFactor.

[fhahn]

In D56739#1359828, @RKSimon wrote:

@fhahn Are you intending to look at the similar SDNode::NumValues limit? PR7250 + PR37000 ?

I am primarily trying to fix SelectionDAG to deal with programs with a very large number of stores at the moment. But I could have a look at the mentioned issues in a bit and see if the same solution is also applicable.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9291

Maybe better to use an ArrayRef<SDValue> here instead of modifying an external value?

9292

Would it be better to create a SDNode::getMaxNumOperands() static/constexpr helper to do this?

[RKSimon]

In D56739#1360490, @fhahn wrote:

In D56739#1359828, @RKSimon wrote:

@fhahn Are you intending to look at the similar SDNode::NumValues limit? PR7250 + PR37000 ?

I am primarily trying to fix SelectionDAG to deal with programs with a very large number of stores at the moment. But I could have a look at the mentioned issues in a bit and see if the same solution is also applicable.

No problem, just trying to ensure those bugs stay on the radar. Thanks.

[fhahn]

updated to use SDNode::getMaxNumOperands

[fhahn]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9291

The reason it modifies an external value is that it allows us to nicely collapse the incoming values into TokenFactors, without needing to allocate an additional array. I changed it to take a SmallVectorImpl, to make it more generic. Does that make sense?

[RKSimon]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9291

I understand, I'm still a little worried but given how infrequent this is likely to be called it should be OK.

llvm/include/llvm/CodeGen/SelectionDAG.h

#endif
   SDNode *UpdateNodeOperands(SDNode *N, SDValue Op1, SDValue Op2,
                                SDValue Op3);
   SDNode *UpdateNodeOperands(SDNode *N, SDValue Op1, SDValue Op2,
                                SDValue Op3, SDValue Op4);
   SDNode *UpdateNodeOperands(SDNode *N, SDValue Op1, SDValue Op2,
                                SDValue Op3, SDValue Op4, SDValue Op5);
   SDNode *UpdateNodeOperands(SDNode *N, ArrayRef<SDValue> Ops);
 
+  /// Creates a new TokenFactor containing \p Vals. If \p Vals contains 64k
+  /// values or more, move values into new TokenFactors in 64k-1 blocks, until
+  /// the final TokenFactor has less than 64k operands.
+  SDValue getTokenFactor(const SDLoc &DL, SmallVectorImpl<SDValue> &Vals);
+
   /// *Mutate* the specified machine node's memory references to the provided
   /// list.
   void setNodeMemRefs(MachineSDNode *N,
                       ArrayRef<MachineMemOperand *> NewMemRefs);
 
   // Propagates the change in divergence to users
   void updateDivergence(SDNode * N);
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

void SelectionDAG::createOperands(SDNode *Node, ArrayRef<SDValue> Vals) {
   Node->NumOperands = Vals.size();
   Node->OperandList = Ops;
   IsDivergent |= TLI->isSDNodeSourceOfDivergence(Node, FLI, DA);
   if (!TLI->isSDNodeAlwaysUniform(Node))
     Node->SDNodeBits.IsDivergent = IsDivergent;
   checkForCycles(Node);
 }
 
+SDValue SelectionDAG::getTokenFactor(const SDLoc &DL,
+                                     SmallVectorImpl<SDValue> &Vals) {
+  size_t Limit = SDNode::getMaxNumOperands();
+  while (Vals.size() > Limit) {
+    unsigned SliceIdx = Vals.size() - Limit;
+    auto ExtractedTFs = ArrayRef<SDValue>(Vals).slice(SliceIdx, Limit);
+    SDValue NewTF = getNode(ISD::TokenFactor, DL, MVT::Other, ExtractedTFs);
+    Vals.erase(Vals.begin() + SliceIdx, Vals.end());
+    Vals.emplace_back(NewTF);
+  }
+  return getNode(ISD::TokenFactor, DL, MVT::Other, Vals);
+}
+
 #ifndef NDEBUG
 static void checkForCyclesHelper(const SDNode *N,
                                  SmallPtrSetImpl<const SDNode*> &Visited,
                                  SmallPtrSetImpl<const SDNode*> &Checked,
                                  const llvm::SelectionDAG *DAG) {
   // If this node has already been checked, don't check it again.
   if (Checked.count(N))
     return;

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

SDValue SelectionDAGBuilder::getRoot() {
   if (PendingLoads.size() == 1) {
     SDValue Root = PendingLoads[0];
     DAG.setRoot(Root);
     PendingLoads.clear();
     return Root;
   }
 
   // Otherwise, we have to make a token factor node.
-  // If we have >= 2^16 loads then split across multiple token factors as
-  // there's a 64k limit on the number of SDNode operands.
-  SDValue Root;
-  size_t Limit = (1 << 16) - 1;
-  while (PendingLoads.size() > Limit) {
-    unsigned SliceIdx = PendingLoads.size() - Limit;
-    auto ExtractedTFs = ArrayRef<SDValue>(PendingLoads).slice(SliceIdx, Limit);
-    SDValue NewTF =
-        DAG.getNode(ISD::TokenFactor, getCurSDLoc(), MVT::Other, ExtractedTFs);
-    PendingLoads.erase(PendingLoads.begin() + SliceIdx, PendingLoads.end());
-    PendingLoads.emplace_back(NewTF);
-  }
-  Root = DAG.getNode(ISD::TokenFactor, getCurSDLoc(), MVT::Other, PendingLoads);
+  SDValue Root = DAG.getTokenFactor(getCurSDLoc(), PendingLoads);
   PendingLoads.clear();
   DAG.setRoot(Root);
   return Root;
 }
 
 SDValue SelectionDAGBuilder::getControlRoot() {
   SDValue Root = DAG.getRoot();