[CodeGen] Improve speed of ComplexDeinterleaving pass

[igor.kirillov]

Summary

Cache all results of running identifyNode, even those that do not identify
potential complex operations. This patch prevents ComplexDeinterleaving pass
from repeatedly trying to identify Nodes for the same pair of instructions.

Fixes https://github.com/llvm/llvm-project/issues/64379

[igor.kirillov]

Added link to GitHub issue

[NickGuy]

LGTM.

And after adding one extra u[i] = u[i] * (a[i] * b[i] + c[i]); line it takes 23 seconds.

Out of interest; How long does that sample take with this change? And any idea why the time it takes grows so quickly with a single statement?

[igor.kirillov]

@NickGuy, it looks like we had O(2**N) complexity there where N is the depth of identifyNode call stack :)

After turning off complex deinterleaving, the file compiles for 0.5 seconds on my machine and with this patch and the complex deinterleaving enabled, it is still the same 0.5 seconds no matter how many times I copy/paste those lines.

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

public:
   explicit ComplexDeinterleavingGraph(const TargetLowering *TL,
                                       const TargetLibraryInfo *TLI)
       : TL(TL), TLI(TLI) {}
 
 private:
   const TargetLowering *TL = nullptr;
   const TargetLibraryInfo *TLI = nullptr;
   SmallVector<NodePtr> CompositeNodes;
+  DenseMap<std::pair<Value *, Value *>, NodePtr> CachedResult;
 
   SmallPtrSet<Instruction *, 16> FinalInstructions;
 
   /// Root instructions are instructions from which complex computation starts
   std::map<Instruction *, NodePtr> RootToNode;
 
   /// Topologically sorted root instructions
   SmallVector<Instruction *, 1> OrderedRoots;
assert(((Operation != ComplexDeinterleavingOperation::ReductionPHI &&
             (R && I)) &&
            "Reduction related nodes must have Real and Imaginary parts");
     return std::make_shared<ComplexDeinterleavingCompositeNode>(Operation, R,
                                                                 I);
   }
 
   NodePtr submitCompositeNode(NodePtr Node) {
     CompositeNodes.push_back(Node);
+    if (Node->Real && Node->Imag)
+      CachedResult[{Node->Real, Node->Imag}] = Node;
     return Node;
   }
 
-  NodePtr getContainingComposite(Value *R, Value *I) {
-    for (const auto &CN : CompositeNodes) {
-      if (CN->Real == R && CN->Imag == I)
-        return CN;
-    }
-    return nullptr;
-  }
-
   /// Identifies a complex partial multiply pattern and its rotation, based on
   /// the following patterns
   ///
   ///  0:  r: cr + ar * br
   ///      i: ci + ar * bi
   /// 90:  r: cr - ai * bi
   ///      i: ci + ai * br
   /// 180: r: cr - ar * br
ComplexDeinterleavingGraph::identifySymmetricOperation(Instruction *Real,
   return submitCompositeNode(Node);
 }
 
 ComplexDeinterleavingGraph::NodePtr
 ComplexDeinterleavingGraph::identifyNode(Value *R, Value *I) {
   LLVM_DEBUG(dbgs() << "identifyNode on " << *R << " / " << *I << "\n");
   assert(R->getType() == I->getType() &&
          "Real and imaginary parts should not have different types");
-  if (NodePtr CN = getContainingComposite(R, I)) {
+
+  auto It = CachedResult.find({R, I});
+  if (It != CachedResult.end()) {
     LLVM_DEBUG(dbgs() << " - Folding to existing node\n");
-    return CN;
+    return It->second;
   }
 
   if (NodePtr CN = identifySplat(R, I))
     return CN;
 
   auto *Real = dyn_cast<Instruction>(R);
   auto *Imag = dyn_cast<Instruction>(I);
   if (!Real || !Imag)
if (HasCMulSupport && HasCAddSupport) {
     if (NodePtr CN = identifyReassocNodes(Real, Imag))
       return CN;
   }
 
   if (NodePtr CN = identifySymmetricOperation(Real, Imag))
     return CN;
 
   LLVM_DEBUG(dbgs() << "  - Not recognised as a valid pattern.\n");
+  CachedResult[{R, I}] = nullptr;
   return nullptr;
 }
 
 ComplexDeinterleavingGraph::NodePtr
 ComplexDeinterleavingGraph::identifyReassocNodes(Instruction *Real,
                                                  Instruction *Imag) {
   auto IsOperationSupported = [](unsigned Opcode) -> bool {
     return Opcode == Instruction::FAdd || Opcode == Instruction::FSub ||