spirv: Rewrite CFG construction

This commit completely rewrites the way we extract a structured CFG from
SPIR-V.  The new approach is different in a few ways:

 1. It does a breadth-first search instead of depth-first.  This means
    that we've visited the merge node for a construct before we visit
    any of the nodes inside the construct.  This makes it easier to
    validate things like loop and switch nesting.

 2. We record more information in the CFG.  Earlier commits added a
    parent pointer to vtn_cf_node but we now record all of the merge and
    other special blocks for each CFG node.  This lets us validate
    things more precisely.

 3. It makes heavy use of merge blocks for walking the CFG.  Previously,
    we sort of used them as hints for trying to guess the CFG structure
    but things got dicey whenever a merge was missing.  We had some
    heuristics for how to handle short-circuiting if statements but it
    was a bunch of special cases.

    Now, we make them a fundamental part of walking the CFG.  When we
    encounter a control-flow construct, we add the body components of
    the construct to the BFS work list and then jump to the merge block
    if one exists to continue scanning the current CFG nesting level.
    If no merge block exists, we assume that means that control-flow
    never re-converges in a normal way and that the only way to get back
    to normality is with a direct jump such as a loop break or continue.
    This should make things far more robust when trying to deal with the
    more creative placement (or lack thereof) of merge instructions.

Reviewed-by: Alan Baker <[email protected]>
Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>
Tested-by: Marge Bot <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/3820>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/3820>

2 files changed, 503 insertions, 294 deletions

diff --git a/src/compiler/spirv/vtn_cfg.c b/src/compiler/spirv/vtn_cfg.c
index dca8a70dda5..935de9c3807 100644
--- a/src/compiler/spirv/vtn_cfg.c
+++ b/src/compiler/spirv/vtn_cfg.c
@@ -397,41 +397,6 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
    return true;
 }
 
-static void
-vtn_add_case(struct vtn_builder *b, struct vtn_switch *swtch,
-             struct vtn_block *break_block,
-             uint32_t block_id, uint64_t val, bool is_default)
-{
-   struct vtn_block *case_block = vtn_block(b, block_id);
-
-   /* Don't create dummy cases that just break */
-   if (case_block == break_block)
-      return;
-
-   if (case_block->switch_case == NULL) {
-      struct vtn_case *c = ralloc(b, struct vtn_case);
-
-      c->node.type = vtn_cf_node_type_case;
-      c->node.parent = &swtch->node;
-      list_inithead(&c->body);
-      c->start_block = case_block;
-      c->fallthrough = NULL;
-      util_dynarray_init(&c->values, b);
-      c->is_default = false;
-      c->visited = false;
-
-      list_addtail(&c->node.link, &swtch->cases);
-
-      case_block->switch_case = c;
-   }
-
-   if (is_default) {
-      case_block->switch_case->is_default = true;
-   } else {
-      util_dynarray_append(&case_block->switch_case->values, uint64_t, val);
-   }
-}
-
 /* This function performs a depth-first search of the cases and puts them
  * in fall-through order.
  */
@@ -461,303 +426,502 @@ vtn_order_case(struct vtn_switch *swtch, struct vtn_case *cse)
    }
 }
 
-static enum vtn_branch_type
-vtn_get_branch_type(struct vtn_builder *b,
-                    struct vtn_block *block,
-                    struct vtn_case *swcase, struct vtn_block *switch_break,
-                    struct vtn_block *loop_break, struct vtn_block *loop_cont)
+static void
+vtn_switch_order_cases(struct vtn_switch *swtch)
 {
-   if (block->switch_case) {
-      /* This branch is actually a fallthrough */
-      vtn_assert(swcase->fallthrough == NULL ||
-                 swcase->fallthrough == block->switch_case);
-      swcase->fallthrough = block->switch_case;
-      return vtn_branch_type_switch_fallthrough;
-   } else if (block == loop_break) {
-      return vtn_branch_type_loop_break;
-   } else if (block == loop_cont) {
-      return vtn_branch_type_loop_continue;
-   } else if (block == switch_break) {
-      return vtn_branch_type_switch_break;
-   } else {
-      return vtn_branch_type_none;
+   struct list_head cases;
+   list_replace(&swtch->cases, &cases);
+   list_inithead(&swtch->cases);
+   while (!list_is_empty(&cases)) {
+      struct vtn_case *cse =
+         list_first_entry(&cases, struct vtn_case, node.link);
+      vtn_order_case(swtch, cse);
    }
 }
 
 static void
-vtn_cfg_walk_blocks(struct vtn_builder *b,
-                    struct vtn_cf_node *cf_parent,
-                    struct list_head *cf_list,
-                    struct vtn_block *start, struct vtn_case *switch_case,
-                    struct vtn_block *switch_break,
-                    struct vtn_block *loop_break, struct vtn_block *loop_cont,
-                    struct vtn_block *end)
+vtn_block_set_merge_cf_node(struct vtn_builder *b, struct vtn_block *block,
+                            struct vtn_cf_node *cf_node)
 {
-   struct vtn_block *block = start;
-   while (block != end) {
-      if (block->merge && (*block->merge & SpvOpCodeMask) == SpvOpLoopMerge &&
-          !block->loop) {
-         struct vtn_loop *loop = ralloc(b, struct vtn_loop);
-
-         loop->node.type = vtn_cf_node_type_loop;
-         loop->node.parent = cf_parent;
-         list_inithead(&loop->body);
-         list_inithead(&loop->cont_body);
-         loop->control = block->merge[3];
-
-         list_addtail(&loop->node.link, cf_list);
-         block->loop = loop;
-
-         struct vtn_block *new_loop_break = vtn_block(b, block->merge[1]);
-         struct vtn_block *new_loop_cont = vtn_block(b, block->merge[2]);
-
-         /* Note: This recursive call will start with the current block as
-          * its start block.  If we weren't careful, we would get here
-          * again and end up in infinite recursion.  This is why we set
-          * block->loop above and check for it before creating one.  This
-          * way, we only create the loop once and the second call that
-          * tries to handle this loop goes to the cases below and gets
-          * handled as a regular block.
-          *
-          * Note: When we make the recursive walk calls, we pass NULL for
-          * the switch break since you have to break out of the loop first.
-          * We do, however, still pass the current switch case because it's
-          * possible that the merge block for the loop is the start of
-          * another case.
-          */
-         vtn_cfg_walk_blocks(b, &loop->node, &loop->body,
-                             block, switch_case, NULL,
-                             new_loop_break, new_loop_cont, NULL );
-         vtn_cfg_walk_blocks(b, &loop->node, &loop->cont_body,
-                             new_loop_cont, NULL, NULL,
-                             new_loop_break, NULL, block);
-
-         enum vtn_branch_type branch_type =
-            vtn_get_branch_type(b, new_loop_break, switch_case, switch_break,
-                                loop_break, loop_cont);
-
-         if (branch_type != vtn_branch_type_none) {
-            /* Stop walking through the CFG when this inner loop's break block
-             * ends up as the same block as the outer loop's continue block
-             * because we are already going to visit it.
-             */
-            vtn_assert(branch_type == vtn_branch_type_loop_continue);
-            return;
+   vtn_fail_if(block->merge_cf_node != NULL,
+               "The merge block declared by a header block cannot be a "
+               "merge block declared by any other header block.");
+
+   block->merge_cf_node = cf_node;
+}
+
+#define VTN_DECL_CF_NODE_FIND(_type)                        \
+static inline struct vtn_##_type *                          \
+vtn_cf_node_find_##_type(struct vtn_cf_node *node)          \
+{                                                           \
+   while (node && node->type != vtn_cf_node_type_##_type)   \
+      node = node->parent;                                  \
+   return (struct vtn_##_type *)node;                       \
+}
+
+VTN_DECL_CF_NODE_FIND(if)
+VTN_DECL_CF_NODE_FIND(loop)
+VTN_DECL_CF_NODE_FIND(case)
+VTN_DECL_CF_NODE_FIND(switch)
+VTN_DECL_CF_NODE_FIND(function)
+
+static enum vtn_branch_type
+vtn_handle_branch(struct vtn_builder *b,
+                  struct vtn_cf_node *cf_parent,
+                  struct vtn_block *target_block)
+{
+   struct vtn_loop *loop = vtn_cf_node_find_loop(cf_parent);
+
+   /* Detect a loop back-edge first.  That way none of the code below
+    * accidentally operates on a loop back-edge.
+    */
+   if (loop && target_block == loop->header_block)
+      return vtn_branch_type_loop_back_edge;
+
+   /* Try to detect fall-through */
+   if (target_block->switch_case) {
+      /* When it comes to handling switch cases, we can break calls to
+       * vtn_handle_branch into two cases: calls from within a case construct
+       * and calls for the jump to each case construct.  In the second case,
+       * cf_parent is the vtn_switch itself and vtn_cf_node_find_case() will
+       * return the outer switch case in which this switch is contained.  It's
+       * fine if the target block is a switch case from an outer switch as
+       * long as it is also the switch break for this switch.
+       */
+      struct vtn_case *switch_case = vtn_cf_node_find_case(cf_parent);
+
+      /* This doesn't get called for the OpSwitch */
+      vtn_fail_if(switch_case == NULL,
+                  "A switch case can only be entered through an OpSwitch or "
+                  "falling through from another switch case.");
+
+      /* Because block->switch_case is only set on the entry block for a given
+       * switch case, we only ever get here if we're jumping to the start of a
+       * switch case.  It's possible, however, that a switch case could jump
+       * to itself via a back-edge.  That *should* get caught by the loop
+       * handling case above but if we have a back edge without a loop merge,
+       * we could en up here.
+       */
+      vtn_fail_if(target_block->switch_case == switch_case,
+                  "A switch cannot fall-through to itself.  Likely, there is "
+                  "a back-edge which is not to a loop header.");
+
+      vtn_fail_if(target_block->switch_case->node.parent !=
+                     switch_case->node.parent,
+                  "A switch case fall-through must come from the same "
+                  "OpSwitch construct");
+
+      vtn_fail_if(switch_case->fallthrough != NULL &&
+                  switch_case->fallthrough != target_block->switch_case,
+                  "Each case construct can have at most one branch to "
+                  "another case construct");
+
+      switch_case->fallthrough = target_block->switch_case;
+
+      /* We don't immediately return vtn_branch_type_switch_fallthrough
+       * because it may also be a loop or switch break for an inner loop or
+       * switch and that takes precedence.
+       */
+   }
+
+   if (loop && target_block == loop->cont_block)
+      return vtn_branch_type_loop_continue;
+
+   /* We walk blocks as a breadth-first search on the control-flow construct
+    * tree where, when we find a construct, we add the vtn_cf_node for that
+    * construct and continue iterating at the merge target block (if any).
+    * Therefore, we want merges whose with parent == cf_parent to be treated
+    * as regular branches.  We only want to consider merges if they break out
+    * of the current CF construct.
+    */
+   if (target_block->merge_cf_node != NULL &&
+       target_block->merge_cf_node->parent != cf_parent) {
+      switch (target_block->merge_cf_node->type) {
+      case vtn_cf_node_type_if:
+         for (struct vtn_cf_node *node = cf_parent;
+              node != target_block->merge_cf_node; node = node->parent) {
+            vtn_fail_if(node == NULL || node->type != vtn_cf_node_type_if,
+                        "Branching to the merge block of a selection "
+                        "construct can only be used to break out of a "
+                        "selection construct");
+
+            struct vtn_if *if_stmt = vtn_cf_node_as_if(node);
+
+            /* This should be guaranteed by our iteration */
+            assert(if_stmt->merge_block != target_block);
+
+            vtn_fail_if(if_stmt->merge_block != NULL,
+                        "Branching to the merge block of a selection "
+                        "construct can only be used to break out of the "
+                        "inner most nested selection level");
          }
+         return vtn_branch_type_if_merge;
 
-         block = new_loop_break;
-         continue;
+      case vtn_cf_node_type_loop:
+         vtn_fail_if(target_block->merge_cf_node != &loop->node,
+                     "Loop breaks can only break out of the inner most "
+                     "nested loop level");
+         return vtn_branch_type_loop_break;
+
+      case vtn_cf_node_type_switch: {
+         struct vtn_switch *swtch = vtn_cf_node_find_switch(cf_parent);
+         vtn_fail_if(target_block->merge_cf_node != &swtch->node,
+                     "Switch breaks can only break out of the inner most "
+                     "nested switch level");
+         return vtn_branch_type_switch_break;
       }
 
-      vtn_assert(block->node.link.next == NULL);
-      block->node.parent = cf_parent;
-      list_addtail(&block->node.link, cf_list);
+      default:
+         unreachable("Invalid CF node type for a merge");
+      }
+   }
 
-      switch (*block->branch & SpvOpCodeMask) {
-      case SpvOpBranch: {
-         struct vtn_block *branch_block = vtn_block(b, block->branch[1]);
+   if (target_block->switch_case)
+      return vtn_branch_type_switch_fallthrough;
 
-         block->branch_type = vtn_get_branch_type(b, branch_block,
-                                                  switch_case, switch_break,
-                                                  loop_break, loop_cont);
+   return vtn_branch_type_none;
+}
 
-         if (block->branch_type != vtn_branch_type_none)
-            return;
+struct vtn_cfg_work_item {
+   struct list_head link;
 
-         block = branch_block;
-         continue;
+   struct vtn_cf_node *cf_parent;
+   struct list_head *cf_list;
+   struct vtn_block *start_block;
+};
+
+static void
+vtn_add_cfg_work_item(struct vtn_builder *b,
+                      struct list_head *work_list,
+                      struct vtn_cf_node *cf_parent,
+                      struct list_head *cf_list,
+                      struct vtn_block *start_block)
+{
+   struct vtn_cfg_work_item *work = ralloc(b, struct vtn_cfg_work_item);
+   work->cf_parent = cf_parent;
+   work->cf_list = cf_list;
+   work->start_block = start_block;
+   list_addtail(&work->link, work_list);
+}
+
+/* Processes a block and returns the next block to process or NULL if we've
+ * reached the end of the construct.
+ */
+static struct vtn_block *
+vtn_process_block(struct vtn_builder *b,
+                  struct list_head *work_list,
+                  struct vtn_cf_node *cf_parent,
+                  struct list_head *cf_list,
+                  struct vtn_block *block)
+{
+   if (!list_is_empty(cf_list)) {
+      /* vtn_process_block() acts like an iterator: it processes the given
+       * block and then returns the next block to process.  For a given
+       * control-flow construct, vtn_build_cfg() calls vtn_process_block()
+       * repeatedly until it finally returns NULL.  Therefore, we know that
+       * the only blocks on which vtn_process_block() can be called are either
+       * the first block in a construct or a block that vtn_process_block()
+       * returned for the current construct.  If cf_list is empty then we know
+       * that we're processing the first block in the construct and we have to
+       * add it to the list.
+       *
+       * If cf_list is not empty, then it must be the block returned by the
+       * previous call to vtn_process_block().  We know a priori that
+       * vtn_process_block only returns either normal branches
+       * (vtn_branch_type_none) or merge target blocks.
+       */
+      switch (vtn_handle_branch(b, cf_parent, block)) {
+      case vtn_branch_type_none:
+         /* For normal branches, we want to process them and add them to the
+          * current construct.  Merge target blocks also look like normal
+          * branches from the perspective of this construct.  See also
+          * vtn_handle_branch().
+          */
+         break;
+
+      case vtn_branch_type_loop_continue:
+      case vtn_branch_type_switch_fallthrough:
+         /* The two cases where we can get early exits from a construct that
+          * are not to that construct's merge target are loop continues and
+          * switch fall-throughs.  In these cases, we need to break out of the
+          * current construct by returning NULL.
+          */
+         return NULL;
+
+      default:
+         /* The only way we can get here is if something was used as two kinds
+          * of merges at the same time and that's illegal.
+          */
+         vtn_fail("A block was used as a merge target from two or more "
+                  "structured control-flow constructs");
       }
+   }
 
-      case SpvOpReturn:
-      case SpvOpReturnValue:
-         block->branch_type = vtn_branch_type_return;
-         return;
+   /* Once a block has been processed, it is placed into and the list link
+    * will point to something non-null.  If we see a node we've already
+    * processed here, it either exists in multiple functions or it's an
+    * invalid back-edge.
+    */
+   if (block->node.parent != NULL) {
+      vtn_fail_if(vtn_cf_node_find_function(&block->node) !=
+                  vtn_cf_node_find_function(cf_parent),
+                  "A block cannot exist in two functions at the "
+                  "same time");
 
-      case SpvOpKill:
-         block->branch_type = vtn_branch_type_discard;
-         return;
+      vtn_fail("Invalid back or cross-edge in the CFG");
+   }
 
-      case SpvOpBranchConditional: {
-         struct vtn_block *then_block = vtn_block(b, block->branch[2]);
-         struct vtn_block *else_block = vtn_block(b, block->branch[3]);
+   if (block->merge && (*block->merge & SpvOpCodeMask) == SpvOpLoopMerge &&
+       block->loop == NULL) {
+      vtn_fail_if((*block->branch & SpvOpCodeMask) != SpvOpBranch &&
+                  (*block->branch & SpvOpCodeMask) != SpvOpBranchConditional,
+                  "An OpLoopMerge instruction must immediately precede "
+                  "either an OpBranch or OpBranchConditional instruction.");
+
+      struct vtn_loop *loop = rzalloc(b, struct vtn_loop);
+
+      loop->node.type = vtn_cf_node_type_loop;
+      loop->node.parent = cf_parent;
+      list_inithead(&loop->body);
+      list_inithead(&loop->cont_body);
+      loop->header_block = block;
+      loop->break_block = vtn_block(b, block->merge[1]);
+      loop->cont_block = vtn_block(b, block->merge[2]);
+      loop->control = block->merge[3];
+
+      list_addtail(&loop->node.link, cf_list);
+      block->loop = loop;
+
+      /* Note: The work item for the main loop body will start with the
+       * current block as its start block.  If we weren't careful, we would
+       * get here again and end up in an infinite loop.  This is why we set
+       * block->loop above and check for it before creating one.  This way,
+       * we only create the loop once and the second iteration that tries to
+       * handle this loop goes to the cases below and gets handled as a
+       * regular block.
+       */
+      vtn_add_cfg_work_item(b, work_list, &loop->node,
+                            &loop->body, loop->header_block);
+
+      /* For continue targets, SPIR-V guarantees the following:
+       *
+       *  - the Continue Target must dominate the back-edge block
+       *  - the back-edge block must post dominate the Continue Target
+       *
+       * If the header block is the same as the continue target, this
+       * condition is trivially satisfied and there is no real continue
+       * section.
+       */
+      if (loop->cont_block != loop->header_block) {
+         vtn_add_cfg_work_item(b, work_list, &loop->node,
+                               &loop->cont_body, loop->cont_block);
+      }
 
-         struct vtn_if *if_stmt = ralloc(b, struct vtn_if);
+      vtn_block_set_merge_cf_node(b, loop->break_block, &loop->node);
 
-         if_stmt->node.type = vtn_cf_node_type_if;
-         if_stmt->node.parent = cf_parent;
-         if_stmt->condition = block->branch[1];
-         list_inithead(&if_stmt->then_body);
-         list_inithead(&if_stmt->else_body);
+      return loop->break_block;
+   }
 
-         list_addtail(&if_stmt->node.link, cf_list);
+   /* Add the block to the CF list */
+   block->node.parent = cf_parent;
+   list_addtail(&block->node.link, cf_list);
 
-         if (block->merge &&
-             (*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge) {
-            if_stmt->control = block->merge[2];
-         } else {
-            if_stmt->control = SpvSelectionControlMaskNone;
-         }
+   switch (*block->branch & SpvOpCodeMask) {
+   case SpvOpBranch: {
+      struct vtn_block *branch_block = vtn_block(b, block->branch[1]);
 
-         if_stmt->then_type = vtn_get_branch_type(b, then_block,
-                                                  switch_case, switch_break,
-                                                  loop_break, loop_cont);
-         if_stmt->else_type = vtn_get_branch_type(b, else_block,
-                                                  switch_case, switch_break,
-                                                  loop_break, loop_cont);
-
-         if (then_block == else_block) {
-            block->branch_type = if_stmt->then_type;
-            if (block->branch_type == vtn_branch_type_none) {
-               block = then_block;
-               continue;
-            } else {
-               return;
-            }
-         } else if (if_stmt->then_type == vtn_branch_type_none &&
-                    if_stmt->else_type == vtn_branch_type_none) {
-            /* Neither side of the if is something we can short-circuit. */
-            vtn_assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
-            struct vtn_block *merge_block = vtn_block(b, block->merge[1]);
-
-            vtn_cfg_walk_blocks(b, &if_stmt->node, &if_stmt->then_body,
-                                then_block, switch_case, switch_break,
-                                loop_break, loop_cont, merge_block);
-            vtn_cfg_walk_blocks(b, &if_stmt->node, &if_stmt->else_body,
-                                else_block, switch_case, switch_break,
-                                loop_break, loop_cont, merge_block);
-
-            enum vtn_branch_type merge_type =
-               vtn_get_branch_type(b, merge_block, switch_case, switch_break,
-                                   loop_break, loop_cont);
-            if (merge_type == vtn_branch_type_none) {
-               block = merge_block;
-               continue;
-            } else {
-               return;
-            }
-         } else if (if_stmt->then_type != vtn_branch_type_none &&
-                    if_stmt->else_type != vtn_branch_type_none) {
-            /* Both sides were short-circuited.  We're done here. */
-            return;
-         } else {
-            /* Exeactly one side of the branch could be short-circuited.
-             * We set the branch up as a predicated break/continue and we
-             * continue on with the other side as if it were what comes
-             * after the if.
-             */
-            if (if_stmt->then_type == vtn_branch_type_none) {
-               block = then_block;
-            } else {
-               block = else_block;
-            }
-            continue;
-         }
-         vtn_fail("Should have returned or continued");
+      block->branch_type = vtn_handle_branch(b, cf_parent, branch_block);
+
+      if (block->branch_type == vtn_branch_type_none)
+         return branch_block;
+      else
+         return NULL;
+   }
+
+   case SpvOpReturn:
+   case SpvOpReturnValue:
+      block->branch_type = vtn_branch_type_return;
+      return NULL;
+
+   case SpvOpKill:
+      block->branch_type = vtn_branch_type_discard;
+      return NULL;
+
+   case SpvOpBranchConditional: {
+      struct vtn_value *cond_val = vtn_untyped_value(b, block->branch[1]);
+      vtn_fail_if(!cond_val->type ||
+                  cond_val->type->base_type != vtn_base_type_scalar ||
+                  cond_val->type->type != glsl_bool_type(),
+                  "Condition must be a Boolean type scalar");
+
+      struct vtn_block *then_block = vtn_block(b, block->branch[2]);
+      struct vtn_block *else_block = vtn_block(b, block->branch[3]);
+
+      if (then_block == else_block) {
+         /* This is uncommon but it can happen.  We treat this the same way as
+          * an unconditional branch.
+          */
+         block->branch_type = vtn_handle_branch(b, cf_parent, then_block);
+
+         if (block->branch_type == vtn_branch_type_none)
+            return then_block;
+         else
+            return NULL;
       }
 
-      case SpvOpSwitch: {
-         vtn_assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
-         struct vtn_block *break_block = vtn_block(b, block->merge[1]);
-
-         struct vtn_switch *swtch = ralloc(b, struct vtn_switch);
-
-         swtch->node.type = vtn_cf_node_type_switch;
-         swtch->node.parent = cf_parent;
-         swtch->selector = block->branch[1];
-         list_inithead(&swtch->cases);
-
-         list_addtail(&swtch->node.link, cf_list);
-
-         /* First, we go through and record all of the cases. */
-         const uint32_t *branch_end =
-            block->branch + (block->branch[0] >> SpvWordCountShift);
-
-         struct vtn_value *cond_val = vtn_untyped_value(b, block->branch[1]);
-         vtn_fail_if(!cond_val->type ||
-                     cond_val->type->base_type != vtn_base_type_scalar,
-                     "Selector of OpSelect must have a type of OpTypeInt");
-
-         nir_alu_type cond_type =
-            nir_get_nir_type_for_glsl_type(cond_val->type->type);
-         vtn_fail_if(nir_alu_type_get_base_type(cond_type) != nir_type_int &&
-                     nir_alu_type_get_base_type(cond_type) != nir_type_uint,
-                     "Selector of OpSelect must have a type of OpTypeInt");
-
-         bool is_default = true;
-         const unsigned bitsize = nir_alu_type_get_type_size(cond_type);
-         for (const uint32_t *w = block->branch + 2; w < branch_end;) {
-            uint64_t literal = 0;
-            if (!is_default) {
-               if (bitsize <= 32) {
-                  literal = *(w++);
-               } else {
-                  assert(bitsize == 64);
-                  literal = vtn_u64_literal(w);
-                  w += 2;
-               }
-            }
+      struct vtn_if *if_stmt = rzalloc(b, struct vtn_if);
 
-            uint32_t block_id = *(w++);
+      if_stmt->node.type = vtn_cf_node_type_if;
+      if_stmt->node.parent = cf_parent;
+      if_stmt->condition = block->branch[1];
+      list_inithead(&if_stmt->then_body);
+      list_inithead(&if_stmt->else_body);
 
-            vtn_add_case(b, swtch, break_block, block_id, literal, is_default);
-            is_default = false;
-         }
+      list_addtail(&if_stmt->node.link, cf_list);
 
-         /* Now, we go through and walk the blocks.  While we walk through
-          * the blocks, we also gather the much-needed fall-through
-          * information.
+      if (block->merge &&
+          (*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge) {
+         /* We may not always have a merge block and that merge doesn't
+          * technically have to be an OpSelectionMerge.  We could have a block
+          * with an OpLoopMerge which ends in an OpBranchConditional.
           */
-         vtn_foreach_cf_node(case_node, &swtch->cases) {
-            struct vtn_case *cse = vtn_cf_node_as_case(case_node);
-            vtn_assert(cse->start_block != break_block);
-            vtn_cfg_walk_blocks(b, &cse->node, &cse->body, cse->start_block,
-                                cse, break_block, loop_break, loop_cont, NULL);
-         }
+         if_stmt->merge_block = vtn_block(b, block->merge[1]);
+         vtn_block_set_merge_cf_node(b, if_stmt->merge_block, &if_stmt->node);
 
-         /* Finally, we walk over all of the cases one more time and put
-          * them in fall-through order.
-          */
-         for (const uint32_t *w = block->branch + 2; w < branch_end;) {
-            struct vtn_block *case_block = vtn_block(b, *w);
+         if_stmt->control = block->merge[2];
+      }
+
+      if_stmt->then_type = vtn_handle_branch(b, &if_stmt->node, then_block);
+      if (if_stmt->then_type == vtn_branch_type_none) {
+         vtn_add_cfg_work_item(b, work_list, &if_stmt->node,
+                               &if_stmt->then_body, then_block);
+      }
+
+      if_stmt->else_type = vtn_handle_branch(b, &if_stmt->node, else_block);
+      if (if_stmt->else_type == vtn_branch_type_none) {
+         vtn_add_cfg_work_item(b, work_list, &if_stmt->node,
+                               &if_stmt->else_body, else_block);
+      }
+
+      return if_stmt->merge_block;
+   }
+
+   case SpvOpSwitch: {
+      struct vtn_value *sel_val = vtn_untyped_value(b, block->branch[1]);
+      vtn_fail_if(!sel_val->type ||
+                  sel_val->type->base_type != vtn_base_type_scalar,
+                  "Selector of OpSwitch must have a type of OpTypeInt");
+
+      nir_alu_type sel_type =
+         nir_get_nir_type_for_glsl_type(sel_val->type->type);
+      vtn_fail_if(nir_alu_type_get_base_type(sel_type) != nir_type_int &&
+                  nir_alu_type_get_base_type(sel_type) != nir_type_uint,
+                  "Selector of OpSwitch must have a type of OpTypeInt");
+
+      struct vtn_switch *swtch = rzalloc(b, struct vtn_switch);
+
+      swtch->node.type = vtn_cf_node_type_switch;
+      swtch->node.parent = cf_parent;
+      swtch->selector = block->branch[1];
+      list_inithead(&swtch->cases);
+
+      list_addtail(&swtch->node.link, cf_list);
+
+      /* We may not always have a merge block */
+      if (block->merge) {
+         vtn_fail_if((*block->merge & SpvOpCodeMask) != SpvOpSelectionMerge,
+                     "An OpLoopMerge instruction must immediately precede "
+                     "either an OpBranch or OpBranchConditional "
+                     "instruction.");
+         swtch->break_block = vtn_block(b, block->merge[1]);
+         vtn_block_set_merge_cf_node(b, swtch->break_block, &swtch->node);
+      }
+
+      /* First, we go through and record all of the cases. */
+      const uint32_t *branch_end =
+         block->branch + (block->branch[0] >> SpvWordCountShift);
+
+      struct hash_table *block_to_case = _mesa_pointer_hash_table_create(b);
 
+      bool is_default = true;
+      const unsigned bitsize = nir_alu_type_get_type_size(sel_type);
+      for (const uint32_t *w = block->branch + 2; w < branch_end;) {
+         uint64_t literal = 0;
+         if (!is_default) {
             if (bitsize <= 32) {
-               w += 2;
+               literal = *(w++);
             } else {
                assert(bitsize == 64);
-               w += 3;
+               literal = vtn_u64_literal(w);
+               w += 2;
             }
+         }
+         struct vtn_block *case_block = vtn_block(b, *(w++));
 
-            if (case_block == break_block)
-               continue;
+         struct hash_entry *case_entry =
+            _mesa_hash_table_search(block_to_case, case_block);
 
-            vtn_assert(case_block->switch_case);
+         struct vtn_case *cse;
+         if (case_entry) {
+            cse = case_entry->data;
+         } else {
+            cse = rzalloc(b, struct vtn_case);
+
+            cse->node.type = vtn_cf_node_type_case;
+            cse->node.parent = &swtch->node;
+            list_inithead(&cse->body);
+            util_dynarray_init(&cse->values, b);
+
+            cse->type = vtn_handle_branch(b, &swtch->node, case_block);
+            switch (cse->type) {
+            case vtn_branch_type_none:
+               /* This is a "real" cases which has stuff in it */
+               vtn_fail_if(case_block->switch_case != NULL,
+                           "OpSwitch has a case which is also in another "
+                           "OpSwitch construct");
+               case_block->switch_case = cse;
+               vtn_add_cfg_work_item(b, work_list, &cse->node,
+                                     &cse->body, case_block);
+               break;
+
+            case vtn_branch_type_switch_break:
+            case vtn_branch_type_loop_break:
+            case vtn_branch_type_loop_continue:
+               /* Switch breaks as well as loop breaks and continues can be
+                * used to break out of a switch construct or as direct targets
+                * of the OpSwitch.
+                */
+               break;
+
+            default:
+               vtn_fail("Target of OpSwitch is not a valid structured exit "
+                        "from the switch construct.");
+            }
 
-            vtn_order_case(swtch, case_block->switch_case);
-         }
+            list_addtail(&cse->node.link, &swtch->cases);
 
-         enum vtn_branch_type branch_type =
-            vtn_get_branch_type(b, break_block, switch_case, NULL,
-                                loop_break, loop_cont);
+            _mesa_hash_table_insert(block_to_case, case_block, cse);
+         }
 
-         if (branch_type != vtn_branch_type_none) {
-            /* It is possible that the break is actually the continue block
-             * for the containing loop.  In this case, we need to bail and let
-             * the loop parsing code handle the continue properly.
-             */
-            vtn_assert(branch_type == vtn_branch_type_loop_continue);
-            return;
+         if (is_default) {
+            cse->is_default = true;
+         } else {
+            util_dynarray_append(&cse->values, uint64_t, literal);
          }
 
-         block = break_block;
-         continue;
+         is_default = false;
       }
 
-      case SpvOpUnreachable:
-         return;
+      _mesa_hash_table_destroy(block_to_case, NULL);
 
-      default:
-         vtn_fail("Unhandled opcode");
-      }
+      return swtch->break_block;
+   }
+
+   case SpvOpUnreachable:
+      return NULL;
+
+   default:
+      vtn_fail("Block did not end with a valid branch instruction");
    }
 }
 
@@ -767,10 +931,30 @@ vtn_build_cfg(struct vtn_builder *b, const uint32_t *words, const uint32_t *end)
    vtn_foreach_instruction(b, words, end,
                            vtn_cfg_handle_prepass_instruction);
 
-   vtn_foreach_cf_node(node, &b->functions) {
-      struct vtn_function *func = vtn_cf_node_as_function(node);
-      vtn_cfg_walk_blocks(b, &func->node, &func->body, func->start_block,
-                          NULL, NULL, NULL, NULL, NULL);
+   vtn_foreach_cf_node(func_node, &b->functions) {
+      struct vtn_function *func = vtn_cf_node_as_function(func_node);
+
+      /* We build the CFG for each function by doing a breadth-first search on
+       * the control-flow graph.  We keep track of our state using a worklist.
+       * Doing a BFS ensures that we visit each structured control-flow
+       * construct and its merge node before we visit the stuff inside the
+       * construct.
+       */
+      struct list_head work_list;
+      list_inithead(&work_list);
+      vtn_add_cfg_work_item(b, &work_list, &func->node, &func->body,
+                            func->start_block);
+
+      while (!list_is_empty(&work_list)) {
+         struct vtn_cfg_work_item *work =
+            list_first_entry(&work_list, struct vtn_cfg_work_item, link);
+         list_del(&work->link);
+
+         for (struct vtn_block *block = work->start_block; block; ) {
+            block = vtn_process_block(b, &work_list, work->cf_parent,
+                                      work->cf_list, block);
+         }
+      }
    }
 }
 
@@ -841,6 +1025,8 @@ vtn_emit_branch(struct vtn_builder *b, enum vtn_branch_type branch_type,
                 nir_variable *switch_fall_var, bool *has_switch_break)
 {
    switch (branch_type) {
+   case vtn_branch_type_if_merge:
+      break; /* Nothing to do */
    case vtn_branch_type_switch_break:
       nir_store_var(&b->nb, switch_fall_var, nir_imm_false(&b->nb), 1);
       *has_switch_break = true;
@@ -853,6 +1039,8 @@ vtn_emit_branch(struct vtn_builder *b, enum vtn_branch_type branch_type,
    case vtn_branch_type_loop_continue:
       nir_jump(&b->nb, nir_jump_continue);
       break;
+   case vtn_branch_type_loop_back_edge:
+      break;
    case vtn_branch_type_return:
       nir_jump(&b->nb, nir_jump_return);
       break;
@@ -1052,6 +1240,11 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
       case vtn_cf_node_type_switch: {
          struct vtn_switch *vtn_switch = vtn_cf_node_as_switch(node);
 
+         /* Before we can emit anything, we need to sort the list of cases in
+          * fall-through order.
+          */
+         vtn_switch_order_cases(vtn_switch);
+
          /* First, we create a variable to keep track of whether or not the
           * switch is still going at any given point.  Any switch breaks
           * will set this variable to false.
diff --git a/src/compiler/spirv/vtn_private.h b/src/compiler/spirv/vtn_private.h
index b7eb31089cd..0d73c61e219 100644
--- a/src/compiler/spirv/vtn_private.h
+++ b/src/compiler/spirv/vtn_private.h
@@ -123,10 +123,12 @@ enum vtn_value_type {
 
 enum vtn_branch_type {
    vtn_branch_type_none,
+   vtn_branch_type_if_merge,
    vtn_branch_type_switch_break,
    vtn_branch_type_switch_fallthrough,
    vtn_branch_type_loop_break,
    vtn_branch_type_loop_continue,
+   vtn_branch_type_loop_back_edge,
    vtn_branch_type_discard,
    vtn_branch_type_return,
 };
@@ -157,6 +159,10 @@ struct vtn_loop {
     */
    struct list_head cont_body;
 
+   struct vtn_block *header_block;
+   struct vtn_block *cont_block;
+   struct vtn_block *break_block;
+
    SpvLoopControlMask control;
 };
 
@@ -171,17 +177,17 @@ struct vtn_if {
    enum vtn_branch_type else_type;
    struct list_head else_body;
 
+   struct vtn_block *merge_block;
+
    SpvSelectionControlMask control;
 };
 
 struct vtn_case {
    struct vtn_cf_node node;
 
+   enum vtn_branch_type type;
    struct list_head body;
 
-   /* The block that starts this case */
-   struct vtn_block *start_block;
-
    /* The fallthrough case, if any */
    struct vtn_case *fallthrough;
 
@@ -201,6 +207,8 @@ struct vtn_switch {
    uint32_t selector;
 
    struct list_head cases;
+
+   struct vtn_block *break_block;
 };
 
 struct vtn_block {
@@ -217,6 +225,14 @@ struct vtn_block {
 
    enum vtn_branch_type branch_type;
 
+   /* The CF node for which this is a merge target
+    *
+    * The SPIR-V spec requires that any given block can be the merge target
+    * for at most one merge instruction.  If this block is a merge target,
+    * this points back to the block containing that merge instruction.
+    */
+   struct vtn_cf_node *merge_cf_node;
+
    /** Points to the loop that this block starts (if it starts a loop) */
    struct vtn_loop *loop;