[SelectionDAG] Split very large token factors for loads into 64k chunks

[aemerson]

Summary

There's a 64k limit on the number of SDNode operands, and some very large functions with 64k or more loads can cause crashes due to this limit being hit when a TokenFactor with this many operands is created. To fix this, create sub-tokenfactors if we've exceeded the limit. No test case as it requires a very large function, however, the test is just this:

define void @foo() {                                                                                                                                                                                                                                                                                                          
  %r1 = load i8, i8* undef                                                                                                                                                                                                                                                                                                    
  %r2 = load i8, i8* undef                                                                                                                                                                                                                                                                                                    
  %r3 = load i8, i8* undef
  ... etc etc 2^16 times
  call void @llvm.trap()                                                                                                                                                                                                                                                                                                      
  unreachable

rdar://45196621

[efriedma]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1052 ↗

(On Diff #175944)

Instead of making a tree of TokenFactors, could you make a list? It seems a little simpler (less code, and you don't have to worry about the length of TokenFactors itself).

I'm a little worried that other code dealing with TokenFactors might end up violating the limit if we're very close... any idea if there's other code that could be affected, like DAGCombine? Do we have an assertion somewhere that will reliably catch this issue?

[aemerson]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1052 ↗

(On Diff #175944)

Could you clarify what you mean by list?

[efriedma]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1052 ↗

(On Diff #175944)

I mean, each TokenFactor should only have one TokenFactor operand. So the algorithm would be something like this: remove "Limit" loads from end of PendingLoads, make a TokenFactor, and append the resulting TokenFactor to PendingLoads. Repeat as necessary.

[RKSimon]

Does this fix PR7250 and PR37000 ? Those test cases look usable.

[aemerson]

In D55073#1313703, @RKSimon wrote:

Does this fix PR7250 and PR37000 ? Those test cases look usable.

PR7250 asserts in a different place, for the number of result values of a node. This issue is very similar, but instead exceeds the number of operands.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1052 ↗

(On Diff #175944)

Ok, I could do that. As for the other places, since SDAG will CSE nodes we should be able to catch this case in getNode() somewhere and assert.

[aemerson]

Changed the approach to replace values in PendingLoads in place. Also added an assert in SelectionDAG::createOperands() to check if we can store the requested number of operands.

[efriedma]

LGTM

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9027 ↗

(On Diff #176189)

Maybe this would be more readable using std::numeric_limits? Your choice.

[aemerson]

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

9027 ↗

(On Diff #176189)

Sure.

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

SDNode *SelectionDAG::isConstantFPBuildVectorOrConstantFP(SDValue N) {
   if (ISD::isBuildVectorOfConstantFPSDNodes(N.getNode()))
     return N.getNode();
 
   return nullptr;
 }
 
 void SelectionDAG::createOperands(SDNode *Node, ArrayRef<SDValue> Vals) {
   assert(!Node->OperandList && "Node already has operands");
+  assert(std::numeric_limits<decltype(SDNode::NumOperands)>::max() >
+             Vals.size() &&
+         "too many operands to fit into SDNode");
   SDUse *Ops = OperandRecycler.allocate(
       ArrayRecycler<SDUse>::Capacity::get(Vals.size()), OperandAllocator);
 
   bool IsDivergent = false;
   for (unsigned I = 0; I != Vals.size(); ++I) {
     Ops[I].setUser(Node);
     Ops[I].setInitial(Vals[I]);
     if (Ops[I].Val.getValueType() != MVT::Other) // Skip Chain. It does not carry divergence.
       IsDivergent = IsDivergent || Ops[I].getNode()->isDivergent();
   }

llvm/trunk/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.
-  SDValue Root = DAG.getNode(ISD::TokenFactor, getCurSDLoc(), MVT::Other,
-                             PendingLoads);
+  // 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);
   PendingLoads.clear();
   DAG.setRoot(Root);
   return Root;
 }
 
 SDValue SelectionDAGBuilder::getControlRoot() {
   SDValue Root = DAG.getRoot();