SPIR-V: Aggressively prune unreachable merge, continue target

More aggressively prune unreachable code as follows.
When no control flow edges reach a merge block or continue target:
- delete their contents so that:
  - a merge block becomes OpLabel, then OpUnreachable
  - a continue target becomes OpLabel, then an OpBranch back to the
    loop header
- any basic block which is dominated by such a merge block or continue
  target is removed as well.
- decorations targeting the removed instructions are removed.

Enables the SPIR-V builder post-processing step the GLSLANG_WEB case.

SPIRV/GlslangToSpv.cpp

@@ -1630,11 +1630,11 @@ void TGlslangToSpvTraverser::finishSpv()
     for (auto it = iOSet.cbegin(); it != iOSet.cend(); ++it)
         entryPoint->addIdOperand(*it);
 
-#ifndef GLSLANG_WEB
-    // Add capabilities, extensions, remove unneeded decorations, etc., 
+    // Add capabilities, extensions, remove unneeded decorations, etc.,
     // based on the resulting SPIR-V.
+    // Note: WebGPU code generation must have the opportunity to aggressively
+    // prune unreachable merge blocks and continue targets.
     builder.postProcess();
-#endif
 }
 
 // Write the SPV into 'out'.

SPIRV/InReadableOrder.cpp

@@ -61,17 +61,22 @@ namespace {
 // Use by calling visit() on the root block.
 class ReadableOrderTraverser {
 public:
-    explicit ReadableOrderTraverser(std::function<void(Block*)> callback) : callback_(callback) {}
+    ReadableOrderTraverser(std::function<void(Block*, spv::ReachReason, Block*)> callback)
+      : callback_(callback) {}
     // Visits the block if it hasn't been visited already and isn't currently
-    // being delayed.  Invokes callback(block), then descends into its
+    // being delayed.  Invokes callback(block, why, header), then descends into its
     // successors.  Delays merge-block and continue-block processing until all
-    // the branches have been completed.
-    void visit(Block* block)
+    // the branches have been completed.  If |block| is an unreachable merge block or
+    // an unreachable continue target, then |header| is the corresponding header block.
+    void visit(Block* block, spv::ReachReason why, Block* header)
     {
         assert(block);
+        if (why == spv::ReachViaControlFlow) {
+            reachableViaControlFlow_.insert(block);
+        }
         if (visited_.count(block) || delayed_.count(block))
             return;
-        callback_(block);
+        callback_(block, why, header);
         visited_.insert(block);
         Block* mergeBlock = nullptr;
         Block* continueBlock = nullptr;
@@ -87,27 +92,40 @@ public:
                 delayed_.insert(continueBlock);
             }
         }
-        const auto successors = block->getSuccessors();
-        for (auto it = successors.cbegin(); it != successors.cend(); ++it)
-            visit(*it);
+        if (why == spv::ReachViaControlFlow) {
+            const auto& successors = block->getSuccessors();
+            for (auto it = successors.cbegin(); it != successors.cend(); ++it)
+                visit(*it, why, nullptr);
+        }
         if (continueBlock) {
+            const spv::ReachReason continueWhy =
+                (reachableViaControlFlow_.count(continueBlock) > 0)
+                    ? spv::ReachViaControlFlow
+                    : spv::ReachDeadContinue;
             delayed_.erase(continueBlock);
-            visit(continueBlock);
+            visit(continueBlock, continueWhy, block);
         }
         if (mergeBlock) {
+            const spv::ReachReason mergeWhy =
+                (reachableViaControlFlow_.count(mergeBlock) > 0)
+                    ? spv::ReachViaControlFlow
+                    : spv::ReachDeadMerge;
             delayed_.erase(mergeBlock);
-            visit(mergeBlock);
+            visit(mergeBlock, mergeWhy, block);
         }
     }
 
 private:
-    std::function<void(Block*)> callback_;
+    std::function<void(Block*, spv::ReachReason, Block*)> callback_;
     // Whether a block has already been visited or is being delayed.
     std::unordered_set<Block *> visited_, delayed_;
+
+    // The set of blocks that actually are reached via control flow.
+    std::unordered_set<Block *> reachableViaControlFlow_;
 };
 }
 
-void spv::inReadableOrder(Block* root, std::function<void(Block*)> callback)
+void spv::inReadableOrder(Block* root, std::function<void(Block*, spv::ReachReason, Block*)> callback)
 {
-    ReadableOrderTraverser(callback).visit(root);
+    ReadableOrderTraverser(callback).visit(root, spv::ReachViaControlFlow, nullptr);
 }

SPIRV/SpvBuilder.h

@@ -683,14 +683,12 @@ public:
     // based on the type of the base and the chain of dereferences.
     Id accessChainGetInferredType();
 
-    // Add capabilities, extensions, remove unneeded decorations, etc., 
+    // Add capabilities, extensions, remove unneeded decorations, etc.,
     // based on the resulting SPIR-V.
     void postProcess();
 
     // Hook to visit each instruction in a block in a function
     void postProcess(Instruction&);
-    // Hook to visit each instruction in a reachable block in a function.
-    void postProcessReachable(const Instruction&);
     // Hook to visit each non-32-bit sized float/int operation in a block.
     void postProcessType(const Instruction&, spv::Id typeId);
 

SPIRV/SpvPostProcess.cpp

@@ -39,6 +39,7 @@
 #include <cassert>
 #include <cstdlib>
 
+#include <unordered_map>
 #include <unordered_set>
 #include <algorithm>
 
@@ -319,16 +320,14 @@ void Builder::postProcess(Instruction& inst)
     }
 }
 
-// Called for each instruction in a reachable block.
-void Builder::postProcessReachable(const Instruction&)
-{
-    // did have code here, but questionable to do so without deleting the instructions
-}
-
 // comment in header
 void Builder::postProcess()
 {
+    // reachableBlocks is the set of blockss reached via control flow, or which are
+    // unreachable continue targert or unreachable merge.
     std::unordered_set<const Block*> reachableBlocks;
+    std::unordered_map<Block*, Block*> headerForUnreachableContinue;
+    std::unordered_set<Block*> unreachableMerges;
     std::unordered_set<Id> unreachableDefinitions;
     // Collect IDs defined in unreachable blocks. For each function, label the
     // reachable blocks first. Then for each unreachable block, collect the
@@ -336,16 +335,41 @@ void Builder::postProcess()
     for (auto fi = module.getFunctions().cbegin(); fi != module.getFunctions().cend(); fi++) {
         Function* f = *fi;
         Block* entry = f->getEntryBlock();
-        inReadableOrder(entry, [&reachableBlocks](const Block* b) { reachableBlocks.insert(b); });
+        inReadableOrder(entry,
+            [&reachableBlocks, &unreachableMerges, &headerForUnreachableContinue]
+            (Block* b, ReachReason why, Block* header) {
+               reachableBlocks.insert(b);
+               if (why == ReachDeadContinue) headerForUnreachableContinue[b] = header;
+               if (why == ReachDeadMerge) unreachableMerges.insert(b);
+            });
         for (auto bi = f->getBlocks().cbegin(); bi != f->getBlocks().cend(); bi++) {
             Block* b = *bi;
-            if (reachableBlocks.count(b) == 0) {
-                for (auto ii = b->getInstructions().cbegin(); ii != b->getInstructions().cend(); ii++)
+            if (unreachableMerges.count(b) != 0 || headerForUnreachableContinue.count(b) != 0) {
+                auto ii = b->getInstructions().cbegin();
+                ++ii; // Keep potential decorations on the label.
+                for (; ii != b->getInstructions().cend(); ++ii)
+                    unreachableDefinitions.insert(ii->get()->getResultId());
+            } else if (reachableBlocks.count(b) == 0) {
+                // The normal case for unreachable code.  All definitions are considered dead.
+                for (auto ii = b->getInstructions().cbegin(); ii != b->getInstructions().cend(); ++ii)
                     unreachableDefinitions.insert(ii->get()->getResultId());
             }
         }
     }
 
+    // Modify unreachable merge blocks and unreachable continue targets.
+    // Delete their contents.
+    for (auto mergeIter = unreachableMerges.begin(); mergeIter != unreachableMerges.end(); ++mergeIter) {
+        (*mergeIter)->rewriteAsCanonicalUnreachableMerge();
+    }
+    for (auto continueIter = headerForUnreachableContinue.begin();
+         continueIter != headerForUnreachableContinue.end();
+         ++continueIter) {
+        Block* continue_target = continueIter->first;
+        Block* header = continueIter->second;
+        continue_target->rewriteAsCanonicalUnreachableContinue(header);
+    }
+
     // Remove unneeded decorations, for unreachable instructions
     decorations.erase(std::remove_if(decorations.begin(), decorations.end(),
         [&unreachableDefinitions](std::unique_ptr<Instruction>& I) -> bool {
@@ -374,13 +398,6 @@ void Builder::postProcess()
         }
     }
 
-    // process all reachable instructions...
-    for (auto bi = reachableBlocks.cbegin(); bi != reachableBlocks.cend(); ++bi) {
-        const Block* block = *bi;
-        const auto function = [this](const std::unique_ptr<Instruction>& inst) { postProcessReachable(*inst.get()); };
-        std::for_each(block->getInstructions().begin(), block->getInstructions().end(), function);
-    }
-
     // process all block-contained instructions
     for (auto fi = module.getFunctions().cbegin(); fi != module.getFunctions().cend(); fi++) {
         Function* f = *fi;

SPIRV/spvIR.h

@@ -226,6 +226,36 @@ public:
         return nullptr;
     }
 
+    // Change this block into a canonical dead merge block.  Delete instructions
+    // as necessary.  A canonical dead merge block has only an OpLabel and an
+    // OpUnreachable.
+    void rewriteAsCanonicalUnreachableMerge() {
+        assert(localVariables.empty());
+        // Delete all instructions except for the label.
+        assert(instructions.size() > 0);
+        instructions.resize(1);
+        successors.clear();
+        Instruction* unreachable = new Instruction(OpUnreachable);
+        addInstruction(std::unique_ptr<Instruction>(unreachable));
+    }
+    // Change this block into a canonical dead continue target branching to the
+    // given header ID.  Delete instructions as necessary.  A canonical dead continue
+    // target has only an OpLabel and an unconditional branch back to the corresponding
+    // header.
+    void rewriteAsCanonicalUnreachableContinue(Block* header) {
+        assert(localVariables.empty());
+        // Delete all instructions except for the label.
+        assert(instructions.size() > 0);
+        instructions.resize(1);
+        successors.clear();
+        // Add OpBranch back to the header.
+        assert(header != nullptr);
+        Instruction* branch = new Instruction(OpBranch);
+        branch->addIdOperand(header->getId());
+        addInstruction(std::move(std::unique_ptr<Instruction>(branch)));
+        successors.push_back(header);
+    }
+
     bool isTerminated() const
     {
         switch (instructions.back()->getOpCode()) {
@@ -235,6 +265,7 @@ public:
         case OpKill:
         case OpReturn:
         case OpReturnValue:
+        case OpUnreachable:
             return true;
         default:
             return false;
@@ -268,10 +299,24 @@ protected:
     bool unreachable;
 };
 
+// The different reasons for reaching a block in the inReadableOrder traversal.
+typedef enum ReachReason {
+    // Reachable from the entry block via transfers of control, i.e. branches.
+    ReachViaControlFlow = 0,
+    // A continue target that is not reachable via control flow.
+    ReachDeadContinue,
+    // A merge block that is not reachable via control flow.
+    ReachDeadMerge
+};
+
 // Traverses the control-flow graph rooted at root in an order suited for
 // readable code generation.  Invokes callback at every node in the traversal
-// order.
-void inReadableOrder(Block* root, std::function<void(Block*)> callback);
+// order.  The callback arguments are:
+// - the block,
+// - the reason we reached the block,
+// - if the reason was that block is an unreachable continue or unreachable merge block
+//   then the last parameter is the corresponding header block.
+void inReadableOrder(Block* root, std::function<void(Block*, ReachReason, Block* header)> callback);
 
 //
 // SPIR-V IR Function.
@@ -321,7 +366,7 @@ public:
             parameterInstructions[p]->dump(out);
 
         // Blocks
-        inReadableOrder(blocks[0], [&out](const Block* b) { b->dump(out); });
+        inReadableOrder(blocks[0], [&out](const Block* b, ReachReason, Block*) { b->dump(out); });
         Instruction end(0, 0, OpFunctionEnd);
         end.dump(out);
     }