[SPIR-V] Fix structurizer issues

The "topological" sorting was behaving incorrectly in some cases:
the exit of a loop could have a lower rank than a node in the loop.
This causes issues when structurizing some patterns, and also codegen
issues as we could generate BBs in the incorrect order in regard to the
SPIR-V spec.

Fixing this ordering alone broke other parts of the structurizer, which
by luck worked. Had to fix those.

Added more test cases, especially to test basic patterns.

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

llvm/lib/Target/SPIRV/SPIRVMergeRegionExitTargets.cpp

@@ -130,6 +130,13 @@ class SPIRVMergeRegionExitTargets : public FunctionPass {
     assert(false && "Unhandled terminator type.");
   }
 
+  AllocaInst *CreateVariable(Function &F, Type *Type,
+                             BasicBlock::iterator Position) {
+    const DataLayout &DL = F.getDataLayout();
+    return new AllocaInst(Type, DL.getAllocaAddrSpace(), nullptr, "reg",
+                          Position);
+  }
+
   // Run the pass on the given convergence region, ignoring the sub-regions.
   // Returns true if the CFG changed, false otherwise.
   bool runOnConvergenceRegionNoRecurse(LoopInfo &LI,
@@ -152,6 +159,9 @@ class SPIRVMergeRegionExitTargets : public FunctionPass {
     auto NewExitTarget = BasicBlock::Create(F->getContext(), "new.exit", F);
     IRBuilder<> Builder(NewExitTarget);
 
+    AllocaInst *Variable = CreateVariable(*F, Builder.getInt32Ty(),
+                                          F->begin()->getFirstInsertionPt());
+
     // CodeGen output needs to be stable. Using the set as-is would order
     // the targets differently depending on the allocation pattern.
     // Sorting per basic-block ordering in the function.
@@ -176,18 +186,16 @@ class SPIRVMergeRegionExitTargets : public FunctionPass {
     std::vector<std::pair<BasicBlock *, Value *>> ExitToVariable;
     for (auto Exit : SortedExits) {
       llvm::Value *Value = createExitVariable(Exit, TargetToValue);
+      IRBuilder<> B2(Exit);
+      B2.SetInsertPoint(Exit->getFirstInsertionPt());
+      B2.CreateStore(Value, Variable);
       ExitToVariable.emplace_back(std::make_pair(Exit, Value));
     }
 
-    // Gather the correct value depending on the exit we came from.
-    llvm::PHINode *node =
-        Builder.CreatePHI(Builder.getInt32Ty(), ExitToVariable.size());
-    for (auto [BB, Value] : ExitToVariable) {
-      node->addIncoming(Value, BB);
-    }
+    llvm::Value *Load = Builder.CreateLoad(Builder.getInt32Ty(), Variable);
 
     // Creating the switch to jump to the correct exit target.
-    llvm::SwitchInst *Sw = Builder.CreateSwitch(node, SortedExitTargets[0],
+    llvm::SwitchInst *Sw = Builder.CreateSwitch(Load, SortedExitTargets[0],
                                                 SortedExitTargets.size() - 1);
     for (size_t i = 1; i < SortedExitTargets.size(); i++) {
       BasicBlock *BB = SortedExitTargets[i];

llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp

@@ -87,6 +87,8 @@ BasicBlock *getExitFor(const ConvergenceRegion *CR) {
 // Returns the merge block designated by I if I is a merge instruction, nullptr
 // otherwise.
 BasicBlock *getDesignatedMergeBlock(Instruction *I) {
+  if (I == nullptr)
+    return nullptr;
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(I);
   if (II == nullptr)
     return nullptr;
@@ -102,6 +104,8 @@ BasicBlock *getDesignatedMergeBlock(Instruction *I) {
 // Returns the continue block designated by I if I is an OpLoopMerge, nullptr
 // otherwise.
 BasicBlock *getDesignatedContinueBlock(Instruction *I) {
+  if (I == nullptr)
+    return nullptr;
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(I);
   if (II == nullptr)
     return nullptr;
@@ -447,55 +451,66 @@ class SPIRVStructurizer : public FunctionPass {
     // clang-format on
     std::vector<Edge>
     createAliasBlocksForComplexEdges(std::vector<Edge> Edges) {
-      std::unordered_map<BasicBlock *, BasicBlock *> Seen;
+      std::unordered_set<BasicBlock *> Seen;
       std::vector<Edge> Output;
       Output.reserve(Edges.size());
 
       for (auto &[Src, Dst] : Edges) {
-        auto [iterator, inserted] = Seen.insert({Src, Dst});
-        if (inserted) {
-          Output.emplace_back(Src, Dst);
-          continue;
+        auto [iterator, inserted] = Seen.insert(Src);
+        if (!inserted) {
+          // Src already a source node. Cannot have 2 edges from A to B.
+          // Creating alias source block.
+          BasicBlock *NewSrc =
+              BasicBlock::Create(F.getContext(), "new.src", &F);
+          replaceBranchTargets(Src, Dst, NewSrc);
+          // replacePhiTargets(Dst, Src, NewSrc);
+          IRBuilder<> Builder(NewSrc);
+          Builder.CreateBr(Dst);
+          Src = NewSrc;
         }
 
-        // The exact same edge was already seen. Ignoring.
-        if (iterator->second == Dst)
+        // Dst has a PHI node. We also need to create an alias output block.
+        if (!hasPhiNode(Dst)) {
+          Output.emplace_back(Src, Dst);
           continue;
+        }
 
-        // The same Src block branches to 2 distinct blocks. This will be an
-        // issue for the generated OpPhi. Creating alias block.
+        // Dst already targeted AND contains a PHI node. We'll need alias
+        // blocks.
         BasicBlock *NewSrc =
-            BasicBlock::Create(F.getContext(), "new.exit.src", &F);
+            BasicBlock::Create(F.getContext(), "phi.alias", &F);
         replaceBranchTargets(Src, Dst, NewSrc);
-        replacePhiTargets(Dst, Src, NewSrc);
-
+        // replacePhiTargets(Dst, Src, NewSrc);
         IRBuilder<> Builder(NewSrc);
         Builder.CreateBr(Dst);
-
-        Seen.emplace(NewSrc, Dst);
-        Output.emplace_back(NewSrc, Dst);
+        Output.emplace_back(Src, NewSrc);
       }
 
       return Output;
     }
 
+    AllocaInst *CreateVariable(Function &F, Type *Type,
+                               BasicBlock::iterator Position) {
+      const DataLayout &DL = F.getDataLayout();
+      return new AllocaInst(Type, DL.getAllocaAddrSpace(), nullptr, "reg",
+                            Position);
+    }
+
     // Given a construct defined by |Header|, and a list of exiting edges
     // |Edges|, creates a new single exit node, fixing up those edges.
     BasicBlock *createSingleExitNode(BasicBlock *Header,
                                      std::vector<Edge> &Edges) {
-      auto NewExit = BasicBlock::Create(F.getContext(), "new.exit", &F);
-      IRBuilder<> ExitBuilder(NewExit);
-
-      std::vector<BasicBlock *> Dsts;
-      std::unordered_map<BasicBlock *, ConstantInt *> DstToIndex;
-
       // Given 2 edges: Src1 -> Dst, Src2 -> Dst:
       // If Dst has an PHI node, and Src1 and Src2 are both operands, both Src1
       // and Src2 cannot be hidden by NewExit. Create 2 new nodes: Alias1,
       // Alias2 to which NewExit will branch before going to Dst. Then, patchup
       // Dst PHI node to look for Alias1 and Alias2.
       std::vector<Edge> FixedEdges = createAliasBlocksForComplexEdges(Edges);
 
+      std::vector<BasicBlock *> Dsts;
+      std::unordered_map<BasicBlock *, ConstantInt *> DstToIndex;
+      auto NewExit = BasicBlock::Create(F.getContext(), "new.exit", &F);
+      IRBuilder<> ExitBuilder(NewExit);
       for (auto &[Src, Dst] : FixedEdges) {
         if (DstToIndex.count(Dst) != 0)
           continue;
@@ -506,33 +521,38 @@ class SPIRVStructurizer : public FunctionPass {
       if (Dsts.size() == 1) {
         for (auto &[Src, Dst] : FixedEdges) {
           replaceBranchTargets(Src, Dst, NewExit);
-          replacePhiTargets(Dst, Src, NewExit);
+          // replacePhiTargets(Dst, Src, NewExit);
         }
         ExitBuilder.CreateBr(Dsts[0]);
         return NewExit;
       }
 
-      PHINode *PhiNode =
-          ExitBuilder.CreatePHI(ExitBuilder.getInt32Ty(), FixedEdges.size());
+      AllocaInst *Variable = CreateVariable(F, ExitBuilder.getInt32Ty(),
+                                            F.begin()->getFirstInsertionPt());
+      // PHINode *PhiNode = ExitBuilder.CreatePHI(ExitBuilder.getInt32Ty(),
+      // FixedEdges.size());
 
       for (auto &[Src, Dst] : FixedEdges) {
-        PhiNode->addIncoming(DstToIndex[Dst], Src);
+        IRBuilder<> B2(Src);
+        B2.SetInsertPoint(Src->getFirstInsertionPt());
+        B2.CreateStore(DstToIndex[Dst], Variable);
         replaceBranchTargets(Src, Dst, NewExit);
-        replacePhiTargets(Dst, Src, NewExit);
       }
 
+      llvm::Value *Load =
+          ExitBuilder.CreateLoad(ExitBuilder.getInt32Ty(), Variable);
+
       // If we can avoid an OpSwitch, generate an OpBranch. Reason is some
       // OpBranch are allowed to exist without a new OpSelectionMerge if one of
       // the branch is the parent's merge node, while OpSwitches are not.
       if (Dsts.size() == 2) {
-        Value *Condition = ExitBuilder.CreateCmp(CmpInst::ICMP_EQ,
-                                                 DstToIndex[Dsts[0]], PhiNode);
+        Value *Condition =
+            ExitBuilder.CreateCmp(CmpInst::ICMP_EQ, DstToIndex[Dsts[0]], Load);
         ExitBuilder.CreateCondBr(Condition, Dsts[0], Dsts[1]);
         return NewExit;
       }
 
-      SwitchInst *Sw =
-          ExitBuilder.CreateSwitch(PhiNode, Dsts[0], Dsts.size() - 1);
+      SwitchInst *Sw = ExitBuilder.CreateSwitch(Load, Dsts[0], Dsts.size() - 1);
       for (auto It = Dsts.begin() + 1; It != Dsts.end(); ++It) {
         Sw->addCase(DstToIndex[*It], *It);
       }
@@ -576,7 +596,7 @@ class SPIRVStructurizer : public FunctionPass {
 
   // Creates a new basic block in F with a single OpUnreachable instruction.
   BasicBlock *CreateUnreachable(Function &F) {
-    BasicBlock *BB = BasicBlock::Create(F.getContext(), "new.exit", &F);
+    BasicBlock *BB = BasicBlock::Create(F.getContext(), "unreachable", &F);
     IRBuilder<> Builder(BB);
     Builder.CreateUnreachable();
     return BB;
@@ -1127,6 +1147,18 @@ class SPIRVStructurizer : public FunctionPass {
         continue;
 
       Modified = true;
+
+      if (Merge == nullptr) {
+        Merge = *successors(Header).begin();
+        IRBuilder<> Builder(Header);
+        Builder.SetInsertPoint(Header->getTerminator());
+
+        auto MergeAddress = BlockAddress::get(Merge->getParent(), Merge);
+        SmallVector<Value *, 1> Args = {MergeAddress};
+        Builder.CreateIntrinsic(Intrinsic::spv_selection_merge, {}, {Args});
+        continue;
+      }
+
       Instruction *SplitInstruction = Merge->getTerminator();
       if (isMergeInstruction(SplitInstruction->getPrevNode()))
         SplitInstruction = SplitInstruction->getPrevNode();

llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp

@@ -29,6 +29,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Pass.h"
 #include "llvm/Target/TargetOptions.h"
+#include "llvm/Transforms/Scalar/Reg2Mem.h"
 #include "llvm/Transforms/Utils.h"
 #include <optional>
 
@@ -162,20 +163,30 @@ void SPIRVPassConfig::addIRPasses() {
   TargetPassConfig::addIRPasses();
 
   if (TM.getSubtargetImpl()->isVulkanEnv()) {
+    addPass(createRegToMemWrapperPass());
+
     // 1.  Simplify loop for subsequent transformations. After this steps, loops
     // have the following properties:
     //  - loops have a single entry edge (pre-header to loop header).
     //  - all loop exits are dominated by the loop pre-header.
     //  - loops have a single back-edge.
     addPass(createLoopSimplifyPass());
 
-    // 2. Merge the convergence region exit nodes into one. After this step,
+    // 2. Removes registers whose lifetime spans across basic blocks. Also
+    // removes phi nodes. This will greatly simplify the next steps.
+    addPass(createRegToMemWrapperPass());
+
+    // 3. Merge the convergence region exit nodes into one. After this step,
     // regions are single-entry, single-exit. This will help determine the
     // correct merge block.
     addPass(createSPIRVMergeRegionExitTargetsPass());
 
-    // 3. Structurize.
+    // 4. Structurize.
     addPass(createSPIRVStructurizerPass());
+
+    // 5. Reduce the amount of variables required by pushing some operations
+    // back to virtual registers.
+    addPass(createPromoteMemoryToRegisterPass());
   }
 
   addPass(createSPIRVRegularizerPass());