[CodeGen] Add support for reductions in ComplexDeinterleaving pass

[igor.kirillov]

Summary

This commit enhances the ComplexDeinterleaving pass to handle unordered
reductions in simple one-block vectorized loops, supporting both
SVE and Neon architectures.

[igor.kirillov]

Add tests

[NickGuy]

Looks good, just a couple of small points from me today.

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

386–391

Might be worth adding a comment here saying what it returns, it's not immediately clear to me

1734–1787

This function is getting big. Solely in terms of readability, I'd suggest moving the code for each condition to separate functions, leaving the ifs here (at least the ReductionOperation branch, ReductionPHI is small enough on its own)

[igor.kirillov]

Refactor replaceNode. Add comments

[igor.kirillov]

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

1734–1787

What do you think about this version?

[igor.kirillov]

Made changes to the checkNodes method in how it verifies reduction nodes.
Additionally, the final reduction operation is now generated at the beginning of the basic block to prevent broken IR.

[NickGuy]

LGTM, nice work

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

1808

Nit: May be worth adding a default case for this switch. Something simple like the following

default:
  llvm_unreachable("Unhandled case in ComplexDeinterleavingGraph::replaceNode");
  break;

[igor.kirillov]

Rebase before pushing

[igor.kirillov]

Add default unreachable case

[fhahn]

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

1435

Iterating over std::map<Instruction *, ... here causes codegen non-determinism as the order depends on the memory addresses, so the order may different across different runs.

I pushed a fix to use MapVector instead: https://github.com/llvm/llvm-project/commit/4c9223c77062bc93f63039c9ebdd885d0f50de59

Spotting such issues can be very difficult and we need to be careful when the iteration order can be non-deterministic

llvm/include/llvm/CodeGen/ComplexDeinterleavingPass.h

 };
 
 enum class ComplexDeinterleavingOperation {
   CAdd,
   CMulPartial,
   // The following 'operations' are used to represent internal states. Backends
   // are not expected to try and support these in any capacity.
   Deinterleave,
-  Symmetric
+  Symmetric,
+  ReductionPHI,
+  ReductionOperation,
 };
 
 enum class ComplexDeinterleavingRotation {
   Rotation_0 = 0,
   Rotation_90 = 1,
   Rotation_180 = 2,
   Rotation_270 = 3,
 };
 
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_COMPLEXDEINTERLEAVING_H

llvm/lib/CodeGen/ComplexDeinterleavingPass.cpp

private:
   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;
 
+  /// When examining a basic block for complex deinterleaving, if it is a simple
+  /// one-block loop, then the only incoming block is 'Incoming' and the
+  /// 'BackEdge' block is the block itself."
+  BasicBlock *BackEdge = nullptr;
+  BasicBlock *Incoming = nullptr;
+
+  /// ReductionInfo maps from %ReductionOp to %PHInode and Instruction
+  /// %OutsideUser as it is shown in the IR:
+  ///
+  /// vector.body:
+  ///   %PHInode = phi <vector type> [ zeroinitializer, %entry ],
+  ///                                [ %ReductionOp, %vector.body ]
+  ///   ...
+  ///   %ReductionOp = fadd i64 ...
+  ///   ...
+  ///   br i1 %condition, label %vector.body, %middle.block
+  ///
+  /// middle.block:
+  ///   %OutsideUser = llvm.vector.reduce.fadd(..., %ReductionOp)
+  ///
+  /// %OutsideUser can be `llvm.vector.reduce.fadd` or `fadd` preceding
+  /// `llvm.vector.reduce.fadd` when unroll factor isn't one.
+  std::map<Instruction *, std::pair<PHINode *, Instruction *>> ReductionInfo;
+
+  /// In the process of detecting a reduction, we consider a pair of
+  /// %ReductionOP, which we refer to as real and imag (or vice versa), and
+  /// traverse the use-tree to detect complex operations. As this is a reduction
+  /// operation, it will eventually reach RealPHI and ImagPHI, which corresponds
+  /// to the %ReductionOPs that we suspect to be complex.
+  /// RealPHI and ImagPHI are used by the identifyPHINode method.
+  PHINode *RealPHI = nullptr;
+  PHINode *ImagPHI = nullptr;
+
+  /// OldToNewPHI maps the original real PHINode to a new, double-sized PHINode.
+  /// The new PHINode corresponds to a vector of deinterleaved complex numbers.
+  /// This mapping is populated during
+  /// ComplexDeinterleavingOperation::ReductionPHI node replacement. It is then
+  /// used in the ComplexDeinterleavingOperation::ReductionOperation node
+  /// replacement process.
+  std::map<PHINode *, PHINode *> OldToNewPHI;
+
   NodePtr prepareCompositeNode(ComplexDeinterleavingOperation Operation,
                                Instruction *R, Instruction *I) {
+    assert(((Operation != ComplexDeinterleavingOperation::ReductionPHI &&
+             Operation != ComplexDeinterleavingOperation::ReductionOperation) ||
+            (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);
     return Node;
   }
private:
   /// complex numbers. There are two ways to represent the Deinterleave
   /// operation:
   /// * Using two shufflevectors with even indices for /pReal instruction and
   /// odd indices for /pImag instructions (only for fixed-width vectors)
   /// * Using two extractvalue instructions applied to `vector.deinterleave2`
   /// intrinsic (for both fixed and scalable vectors)
   NodePtr identifyDeinterleave(Instruction *Real, Instruction *Imag);
 
+  NodePtr identifyPHINode(Instruction *Real, Instruction *Imag);
+
   Value *replaceNode(IRBuilderBase &Builder, RawNodePtr Node);
 
 public:
   void dump() { dump(dbgs()); }
   void dump(raw_ostream &OS) {
     for (const auto &Node : CompositeNodes)
       Node->dump(OS);
   }
 
   /// Returns false if the deinterleaving operation should be cancelled for the
   /// current graph.
   bool identifyNodes(Instruction *RootI);
 
+  bool collectPotentialReductions(BasicBlock *B);
+
+  void identifyReductionNodes();
+
   /// Check that every instruction, from the roots to the leaves, has internal
   /// uses.
   bool checkNodes();
 
   /// Perform the actual replacement of the underlying instruction graph.
   void replaceNodes();
 };
 
 class ComplexDeinterleaving {
if (Mask[Idx] != (Idx * 2) + Offset)
       return false;
   }
 
   return true;
 }
 
 bool ComplexDeinterleaving::evaluateBasicBlock(BasicBlock *B) {
   ComplexDeinterleavingGraph Graph(TL, TLI);
+  if (Graph.collectPotentialReductions(B))
+    Graph.identifyReductionNodes();
+
   for (auto &I : *B)
     Graph.identifyNodes(&I);
 
   if (Graph.checkNodes()) {
     Graph.replaceNodes();
     return true;
   }
 
ComplexDeinterleavingGraph::identifyNode(Instruction *Real, Instruction *Imag) {
   if (NodePtr CN = getContainingComposite(Real, Imag)) {
     LLVM_DEBUG(dbgs() << " - Folding to existing node\n");
     return CN;
   }
 
   if (NodePtr CN = identifyDeinterleave(Real, Imag))
     return CN;
 
+  if (NodePtr CN = identifyPHINode(Real, Imag))
+    return CN;
+
   auto *VTy = cast<VectorType>(Real->getType());
   auto *NewVTy = VectorType::getDoubleElementsVectorType(VTy);
 
   bool HasCMulSupport = TL->isComplexDeinterleavingOperationSupported(
       ComplexDeinterleavingOperation::CMulPartial, NewVTy);
   bool HasCAddSupport = TL->isComplexDeinterleavingOperationSupported(
       ComplexDeinterleavingOperation::CAdd, NewVTy);
 
for (auto ItI = ImagAddends.begin(); ItI != ImagAddends.end(); ++ItI) {
         }
       }
     }
   }
   return nullptr;
 }
 
 bool ComplexDeinterleavingGraph::identifyNodes(Instruction *RootI) {
+  // This potential root instruction might already have been recognized as
+  // reduction. Because RootToNode maps both Real and Imaginary parts to
+  // CompositeNode we should choose only one either Real or Imag instruction to
+  // use as an anchor for generating complex instruction.
+  auto It = RootToNode.find(RootI);
+  if (It != RootToNode.end() && It->second->Real == RootI) {
+    OrderedRoots.push_back(RootI);
+    return true;
+  }
+
   auto RootNode = identifyRoot(RootI);
   if (!RootNode)
     return false;
 
   LLVM_DEBUG({
     Function *F = RootI->getFunction();
     BasicBlock *B = RootI->getParent();
     dbgs() << "Complex deinterleaving graph for " << F->getName()
            << "::" << B->getName() << ".\n";
     dump(dbgs());
     dbgs() << "\n";
   });
   RootToNode[RootI] = RootNode;
   OrderedRoots.push_back(RootI);
   return true;
 }
 
+bool ComplexDeinterleavingGraph::collectPotentialReductions(BasicBlock *B) {
+  bool FoundPotentialReduction = false;
+
+  auto *Br = dyn_cast<BranchInst>(B->getTerminator());
+  if (!Br || Br->getNumSuccessors() != 2)
+    return false;
+
+  // Identify simple one-block loop
+  if (Br->getSuccessor(0) != B && Br->getSuccessor(1) != B)
+    return false;
+
+  SmallVector<PHINode *> PHIs;
+  for (auto &PHI : B->phis()) {
+    if (PHI.getNumIncomingValues() != 2)
+      continue;
+
+    if (!PHI.hasOneUser())
+      continue;
+
+    if (!PHI.getType()->isVectorTy())
+      continue;
+
+    auto *ReductionOp = dyn_cast<Instruction>(PHI.getIncomingValueForBlock(B));
+    if (!ReductionOp)
+      continue;
+
+    // Check if final instruction is reduced outside of current block
+    Instruction *FinalReduction = nullptr;
+    auto NumUsers = 0u;
+    for (auto *U : ReductionOp->users()) {
+      ++NumUsers;
+      if (U == &PHI)
+        continue;
+      FinalReduction = dyn_cast<Instruction>(U);
+    }
+
+    if (NumUsers != 2 || !FinalReduction || FinalReduction->getParent() == B)
+      continue;
+
+    ReductionInfo[ReductionOp] = {&PHI, FinalReduction};
+    FinalInstructions.insert(&PHI);
+
+    BackEdge = B;
+    auto BackEdgeIdx = PHI.getBasicBlockIndex(B);
+    auto IncomingIdx = BackEdgeIdx == 0 ? 1 : 0;
+    Incoming = PHI.getIncomingBlock(IncomingIdx);
+    FoundPotentialReduction = true;
+  }
+  return FoundPotentialReduction;
+}
+
+void ComplexDeinterleavingGraph::identifyReductionNodes() {
+  SmallVector<bool> Processed(ReductionInfo.size(), false);
+  SmallVector<Instruction *> OperationInstruction;
+  for (auto &P : ReductionInfo)
+    OperationInstruction.push_back(P.first);
+
+  // Identify a complex computation by evaluating two reduction operations that
+  // potentially could be involved
+  for (size_t i = 0; i < OperationInstruction.size(); ++i) {
+    if (Processed[i])
+      continue;
+    for (size_t j = i + 1; j < OperationInstruction.size(); ++j) {
+      if (Processed[j])
+        continue;
+
+      auto *Real = OperationInstruction[i];
+      auto *Imag = OperationInstruction[j];
+
+      RealPHI = ReductionInfo[Real].first;
+      ImagPHI = ReductionInfo[Imag].first;
+      auto Node = identifyNode(Real, Imag);
+      if (!Node) {
+        std::swap(Real, Imag);
+        std::swap(RealPHI, ImagPHI);
+        Node = identifyNode(Real, Imag);
+      }
+
+      // If a node is identified, mark its operation instructions as used to
+      // prevent re-identification and attach the node to the real part
+      if (Node) {
+        LLVM_DEBUG(dbgs() << "Identified reduction starting from instructions: "
+                          << *Real << " / " << *Imag << "\n");
+        Processed[i] = true;
+        Processed[j] = true;
+        auto RootNode = prepareCompositeNode(
+            ComplexDeinterleavingOperation::ReductionOperation, Real, Imag);
+        RootNode->addOperand(Node);
+        RootToNode[Real] = RootNode;
+        RootToNode[Imag] = RootNode;
+        submitCompositeNode(RootNode);
+        break;
+      }
+    }
+  }
+
+  RealPHI = nullptr;
+  ImagPHI = nullptr;
+}
+
 bool ComplexDeinterleavingGraph::checkNodes() {
   // Collect all instructions from roots to leaves
   SmallPtrSet<Instruction *, 16> AllInstructions;
   SmallVector<Instruction *, 8> Worklist;
   for (auto *I : OrderedRoots)
     Worklist.push_back(I);
 
   // Extract all instructions that are used by all XCMLA/XCADD/ADD/SUB/NEG
NodePtr PlaceholderNode =
       prepareCompositeNode(llvm::ComplexDeinterleavingOperation::Deinterleave,
                            RealShuffle, ImagShuffle);
   PlaceholderNode->ReplacementNode = RealShuffle->getOperand(0);
   FinalInstructions.insert(RealShuffle);
   FinalInstructions.insert(ImagShuffle);
   return submitCompositeNode(PlaceholderNode);
 }
 
+ComplexDeinterleavingGraph::NodePtr
+ComplexDeinterleavingGraph::identifyPHINode(Instruction *Real,
+                                            Instruction *Imag) {
+  if (Real != RealPHI || Imag != ImagPHI)
+    return nullptr;
+
+  NodePtr PlaceholderNode = prepareCompositeNode(
+      ComplexDeinterleavingOperation::ReductionPHI, Real, Imag);
+  return submitCompositeNode(PlaceholderNode);
+}
+
 static Value *replaceSymmetricNode(IRBuilderBase &B, unsigned Opcode,
                                    FastMathFlags Flags, Value *InputA,
                                    Value *InputB) {
   Value *I;
   switch (Opcode) {
   case Instruction::FNeg:
     I = B.CreateFNeg(InputA);
     break;
static Value *replaceSymmetricNode(IRBuilderBase &B, unsigned Opcode,
   return I;
 }
 
 Value *ComplexDeinterleavingGraph::replaceNode(IRBuilderBase &Builder,
                                                RawNodePtr Node) {
   if (Node->ReplacementNode)
     return Node->ReplacementNode;
 
+  // If Operation is ReductionPHI, a new empty PHINode is created.
+  // It is filled later when the ReductionOperation is processed.
+  if (Node->Operation == ComplexDeinterleavingOperation::ReductionPHI) {
+    auto *VTy = cast<VectorType>(Node->Real->getType());
+    auto *NewVTy = VectorType::getDoubleElementsVectorType(VTy);
+    auto *NewPHI = PHINode::Create(NewVTy, 0, "", BackEdge->getFirstNonPHI());
+    OldToNewPHI[dyn_cast<PHINode>(Node->Real)] = NewPHI;
+    Node->ReplacementNode = NewPHI;
+    return NewPHI;
+  }
+
+  // If Operation is ReductionOperation, find the previously generated empty
+  // PHINode and interleave the initial values.
+  if (Node->Operation == ComplexDeinterleavingOperation::ReductionOperation) {
+    auto Result = replaceNode(Builder, Node->Operands[0]);
+    auto *OldPHIReal = ReductionInfo[Node->Real].first;
+    auto *OldPHIImag = ReductionInfo[Node->Imag].first;
+    auto *NewPHI = OldToNewPHI[OldPHIReal];
+
+    auto *VTy = cast<VectorType>(Node->Real->getType());
+    auto *NewVTy = VectorType::getDoubleElementsVectorType(VTy);
+
+    // We have to interleave initial origin values coming from IncomingBlock
+    Value *InitReal = OldPHIReal->getIncomingValueForBlock(Incoming);
+    Value *InitImag = OldPHIImag->getIncomingValueForBlock(Incoming);
+
+    IRBuilder<> Builder(Incoming->getTerminator());
+    auto *NewInit =
+        Builder.CreateIntrinsic(Intrinsic::experimental_vector_interleave2,
+                                NewVTy, {InitReal, InitImag});
+
+    NewPHI->addIncoming(NewInit, Incoming);
+    NewPHI->addIncoming(Result, BackEdge);
+
+    // Deinterleave complex vector outside of loop so that it can be finally
+    // reduced
+    auto *FinalReductionReal = ReductionInfo[Node->Real].second;
+    auto *FinalReductionImag = ReductionInfo[Node->Imag].second;
+
+    Builder.SetInsertPoint(FinalReductionReal);
+    auto *Deinterleave =
+        Builder.CreateIntrinsic(Intrinsic::experimental_vector_deinterleave2,
+                                Result->getType(), Result);
+
+    auto *NewReal = Builder.CreateExtractValue(Deinterleave, (uint64_t)0);
+    FinalReductionReal->replaceUsesOfWith(Node->Real, NewReal);
+
+    Builder.SetInsertPoint(FinalReductionImag);
+    auto *NewImag = Builder.CreateExtractValue(Deinterleave, 1);
+    FinalReductionImag->replaceUsesOfWith(Node->Imag, NewImag);
+
+    Node->ReplacementNode = Result;
+    return Result;
+  }
+
   Value *Input0 = replaceNode(Builder, Node->Operands[0]);
   Value *Input1 = Node->Operands.size() > 1
                       ? replaceNode(Builder, Node->Operands[1])
                       : nullptr;
   Value *Accumulator = Node->Operands.size() > 2
                            ? replaceNode(Builder, Node->Operands[2])
                            : nullptr;
   if (Input1)
     assert(Input0->getType() == Input1->getType() &&
            "Node inputs need to be of the same type");
 
   if (Node->Operation == ComplexDeinterleavingOperation::Symmetric)
     Node->ReplacementNode = replaceSymmetricNode(Builder, Node->Opcode,
                                                  Node->Flags, Input0, Input1);
   else
     Node->ReplacementNode = TL->createComplexDeinterleavingIR(
         Builder, Node->Operation, Node->Rotation, Input0, Input1, Accumulator);
 
   assert(Node->ReplacementNode && "Target failed to create Intrinsic call.");
   NumComplexTransformations += 1;
   return Node->ReplacementNode;
 }
 
 void ComplexDeinterleavingGraph::replaceNodes() {
   SmallVector<Instruction *, 16> DeadInstrRoots;
   for (auto *RootInstruction : OrderedRoots) {
     // Check if this potential root went through check process and we can
     // deinterleave it
     if (!RootToNode.count(RootInstruction))
       continue;
 
     IRBuilder<> Builder(RootInstruction);
     auto RootNode = RootToNode[RootInstruction];
     Value *R = replaceNode(Builder, RootNode.get());
+
+    if (RootNode->Operation ==
+        ComplexDeinterleavingOperation::ReductionOperation) {
+      ReductionInfo[RootNode->Real].first->removeIncomingValue(BackEdge);
+      ReductionInfo[RootNode->Imag].first->removeIncomingValue(BackEdge);
+      DeadInstrRoots.push_back(RootNode->Real);
+      DeadInstrRoots.push_back(RootNode->Imag);
+    } else {
       assert(R && "Unable to find replacement for RootInstruction");
       DeadInstrRoots.push_back(RootInstruction);
       RootInstruction->replaceAllUsesWith(R);
     }
+  }
 
   for (auto *I : DeadInstrRoots)
     RecursivelyDeleteTriviallyDeadInstructions(I, TLI);
 }