[spirv] Support continue and break statements (#543)

tools/clang/include/clang/SPIRV/ModuleBuilder.h

@@ -71,7 +71,9 @@ public:
 
   /// \brief Creates a SPIR-V basic block. On success, returns the <label-id>
   /// for the basic block. On failure, returns zero.
-  uint32_t createBasicBlock(llvm::StringRef name = "");
+  /// All basic blocks are reachable by default. If isReachable is set to false,
+  /// a debug name will not be created for the basic block.
+  uint32_t createBasicBlock(llvm::StringRef name = "", bool isReachable = true);
 
   /// \brief Adds the basic block with the given label as a successor to the
   /// current basic block.

tools/clang/include/clang/SPIRV/Structure.h

@@ -80,7 +80,9 @@ private:
 class BasicBlock {
 public:
   /// \brief Constructs a basic block with the given <label-id>.
-  inline explicit BasicBlock(uint32_t labelId);
+  /// By default all basic blocks are considered reachable, unless the caller
+  /// specifies otherwise.
+  inline explicit BasicBlock(uint32_t labelId, bool isReachable = true);
 
   // Disable copy constructor/assignment
   BasicBlock(const BasicBlock &) = delete;
@@ -135,6 +137,9 @@ public:
   /// \brief Returns true if this basic block is terminated.
   bool isTerminated() const;
 
+  /// \brief Returns true if this basic block is reachable in the control flow.
+  inline bool isReachable() const;
+
 private:
   uint32_t labelId; ///< The label id for this basic block. Zero means invalid.
   std::deque<Instruction> instructions;
@@ -142,6 +147,7 @@ private:
   llvm::SmallVector<BasicBlock *, 2> successors;
   BasicBlock *mergeTarget;
   BasicBlock *continueTarget;
+  bool reachable;
 };
 
 // === Function definition ===
@@ -351,8 +357,9 @@ std::vector<uint32_t> Instruction::take() { return std::move(words); }
 
 // === Basic block inline implementations ===
 
-BasicBlock::BasicBlock(uint32_t id)
-    : labelId(id), mergeTarget(nullptr), continueTarget(nullptr) {}
+BasicBlock::BasicBlock(uint32_t id, bool isReachable)
+    : labelId(id), mergeTarget(nullptr), continueTarget(nullptr),
+      reachable(isReachable) {}
 
 bool BasicBlock::isEmpty() const {
   return labelId == 0 && instructions.empty();
@@ -384,6 +391,8 @@ void BasicBlock::setContinueTarget(BasicBlock *target) {
 
 BasicBlock *BasicBlock::getContinueTarget() const { return continueTarget; }
 
+bool BasicBlock::isReachable() const { return reachable; }
+
 // === Function inline implementations ===
 
 Function::Function(uint32_t rType, uint32_t rId,

tools/clang/lib/SPIRV/ModuleBuilder.cpp

@@ -97,15 +97,23 @@ bool ModuleBuilder::endFunction() {
   return true;
 }
 
-uint32_t ModuleBuilder::createBasicBlock(llvm::StringRef name) {
+uint32_t ModuleBuilder::createBasicBlock(llvm::StringRef name,
+                                         bool isReachable) {
   if (theFunction == nullptr) {
     assert(false && "found detached basic block");
     return 0;
   }
 
   const uint32_t labelId = theContext.takeNextId();
-  basicBlocks[labelId] = llvm::make_unique<BasicBlock>(labelId);
-  theModule.addDebugName(labelId, name);
+  basicBlocks[labelId] = llvm::make_unique<BasicBlock>(labelId, isReachable);
+
+  // OpName instructions should not be added for unreachable basic blocks
+  // because such blocks are *not* discovered by BlockReadableOrderVisitor and
+  // therefore they are not emitted.
+  // The newly created basic block is unreachable if specified by the caller,
+  // or, if this block is being created by a block that is already unreachable.
+  if (isReachable && (!insertPoint || insertPoint->isReachable()))
+    theModule.addDebugName(labelId, name);
 
   return labelId;
 }

tools/clang/lib/SPIRV/SPIRVEmitter.cpp

@@ -122,6 +122,56 @@ bool isSpirvMatrixOp(spv::Op opcode) {
   return false;
 }
 
+/// \brief Returns the statement that is the immediate parent AST node of the
+/// given statement. Returns nullptr if there are no parents nodes.
+const Stmt *getImmediateParent(ASTContext &astContext, const Stmt *stmt) {
+  const auto &parents = astContext.getParents(*stmt);
+  return parents.empty() ? nullptr : parents[0].get<Stmt>();
+}
+
+/// \brief Returns true if there are no unreachable statements after the given
+/// statement. A "continue" statement and "break" statement cause a branch to a
+/// loop header and a loop merge block, respectively. A "return" statement
+/// causes a branch out of the function. In such situations, any statement that
+/// follows the break/continue/return will not be executed. When this method
+/// returns false, it indicates that there exists statements that are not going
+/// to be executed.
+bool isLastStmtBeforeControlFlowBranching(ASTContext &astContext,
+                                          const Stmt *stmt) {
+  const Stmt *parent = getImmediateParent(astContext, stmt);
+  if (const auto *parentCS = dyn_cast_or_null<CompoundStmt>(parent)) {
+    if (stmt == *(parentCS->body_rbegin())) {
+      // The current statement is the last child node of the parent.
+
+      // Handle nested compound statements. e.g.
+      // while (cond) {
+      //   StmtA;
+      //   {
+      //      StmtB;
+      //      {{continue;}}
+      //   }
+      //   {StmtC;}
+      //   StmtD;
+      // }
+      //
+      // The continue statement is the last statement in its CompoundStmt scope,
+      // but, if nested compound statements are flattened, the continue
+      // statement is not the last statement in the current loop scope.
+      const Stmt *grandparent = getImmediateParent(astContext, parent);
+      if (grandparent && isa<CompoundStmt>(grandparent))
+        return isLastStmtBeforeControlFlowBranching(astContext, parent);
+
+      return true;
+    }
+  }
+
+  return false;
+}
+
+bool isLoopStmt(const Stmt *stmt) {
+  return isa<ForStmt>(stmt) || isa<WhileStmt>(stmt) || isa<DoStmt>(stmt);
+}
+
 } // namespace
 
 SPIRVEmitter::SPIRVEmitter(CompilerInstance &ci)
@@ -212,6 +262,8 @@ void SPIRVEmitter::doStmt(const Stmt *stmt,
     doBreakStmt(breakStmt);
   } else if (const auto *theDoStmt = dyn_cast<DoStmt>(stmt)) {
     doDoStmt(theDoStmt, attrs);
+  } else if (const auto *continueStmt = dyn_cast<ContinueStmt>(stmt)) {
+    doContinueStmt(continueStmt);
   } else if (const auto *whileStmt = dyn_cast<WhileStmt>(stmt)) {
     doWhileStmt(whileStmt, attrs);
   } else if (const auto *forStmt = dyn_cast<ForStmt>(stmt)) {
@@ -531,6 +583,11 @@ void SPIRVEmitter::doDoStmt(const DoStmt *theDoStmt,
   const uint32_t continueBB = theBuilder.createBasicBlock("do_while.continue");
   const uint32_t mergeBB = theBuilder.createBasicBlock("do_while.merge");
 
+  // Make sure any continue statements branch to the continue block, and any
+  // break statements branch to the merge block.
+  continueStack.push(continueBB);
+  breakStack.push(mergeBB);
+
   // Branch from the current insert point to the header block.
   theBuilder.createBranch(headerBB);
   theBuilder.addSuccessor(headerBB);
@@ -550,7 +607,8 @@ void SPIRVEmitter::doDoStmt(const DoStmt *theDoStmt,
   if (const Stmt *body = theDoStmt->getBody()) {
     doStmt(body);
   }
-  theBuilder.createBranch(continueBB);
+  if (!theBuilder.isCurrentBasicBlockTerminated())
+    theBuilder.createBranch(continueBB);
   theBuilder.addSuccessor(continueBB);
 
   // Process the <continue> block. The check for whether the loop should
@@ -571,6 +629,37 @@ void SPIRVEmitter::doDoStmt(const DoStmt *theDoStmt,
 
   // Set insertion point to the <merge> block for subsequent statements
   theBuilder.setInsertPoint(mergeBB);
+
+  // Done with the current scope's continue block and merge block.
+  continueStack.pop();
+  breakStack.pop();
+}
+
+void SPIRVEmitter::doContinueStmt(const ContinueStmt *continueStmt) {
+  assert(!theBuilder.isCurrentBasicBlockTerminated());
+  const uint32_t continueTargetBB = continueStack.top();
+  theBuilder.createBranch(continueTargetBB);
+  theBuilder.addSuccessor(continueTargetBB);
+
+  // If any statements follow a continue statement in a loop, they will not be
+  // executed. For example, StmtB and StmtC below are never executed:
+  //
+  // while (true) {
+  //   StmtA;
+  //   continue;
+  //   StmtB;
+  //   StmtC;
+  // }
+  //
+  // To handle such cases, we do not stop tranlsation. We create a new basic
+  // block in which StmtB and StmtC will be translated.
+  // Note that since this basic block is unreachable, BlockReadableOrderVisitor
+  // will not emit it in the final module binary.
+  if (!isLastStmtBeforeControlFlowBranching(astContext, continueStmt)) {
+    const uint32_t unreachableBB =
+        theBuilder.createBasicBlock("unreachable", /*isReachable*/ false);
+    theBuilder.setInsertPoint(unreachableBB);
+  }
 }
 
 void SPIRVEmitter::doWhileStmt(const WhileStmt *whileStmt,
@@ -616,6 +705,11 @@ void SPIRVEmitter::doWhileStmt(const WhileStmt *whileStmt,
   const uint32_t continueBB = theBuilder.createBasicBlock("while.continue");
   const uint32_t mergeBB = theBuilder.createBasicBlock("while.merge");
 
+  // Make sure any continue statements branch to the continue block, and any
+  // break statements branch to the merge block.
+  continueStack.push(continueBB);
+  breakStack.push(mergeBB);
+
   // Process the <check> block
   theBuilder.createBranch(checkBB);
   theBuilder.addSuccessor(checkBB);
@@ -651,7 +745,8 @@ void SPIRVEmitter::doWhileStmt(const WhileStmt *whileStmt,
   if (const Stmt *body = whileStmt->getBody()) {
     doStmt(body);
   }
-  theBuilder.createBranch(continueBB);
+  if (!theBuilder.isCurrentBasicBlockTerminated())
+    theBuilder.createBranch(continueBB);
   theBuilder.addSuccessor(continueBB);
 
   // Process the <continue> block. While loops do not have an explicit
@@ -662,6 +757,10 @@ void SPIRVEmitter::doWhileStmt(const WhileStmt *whileStmt,
 
   // Set insertion point to the <merge> block for subsequent statements
   theBuilder.setInsertPoint(mergeBB);
+
+  // Done with the current scope's continue and merge blocks.
+  continueStack.pop();
+  breakStack.pop();
 }
 
 void SPIRVEmitter::doForStmt(const ForStmt *forStmt,
@@ -709,6 +808,11 @@ void SPIRVEmitter::doForStmt(const ForStmt *forStmt,
   const uint32_t continueBB = theBuilder.createBasicBlock("for.continue");
   const uint32_t mergeBB = theBuilder.createBasicBlock("for.merge");
 
+  // Make sure any continue statements branch to the continue block, and any
+  // break statements branch to the merge block.
+  continueStack.push(continueBB);
+  breakStack.push(mergeBB);
+
   // Process the <init> block
   if (const Stmt *initStmt = forStmt->getInit()) {
     doStmt(initStmt);
@@ -741,7 +845,8 @@ void SPIRVEmitter::doForStmt(const ForStmt *forStmt,
   if (const Stmt *body = forStmt->getBody()) {
     doStmt(body);
   }
-  theBuilder.createBranch(continueBB);
+  if (!theBuilder.isCurrentBasicBlockTerminated())
+    theBuilder.createBranch(continueBB);
   theBuilder.addSuccessor(continueBB);
 
   // Process the <continue> block
@@ -754,6 +859,10 @@ void SPIRVEmitter::doForStmt(const ForStmt *forStmt,
 
   // Set insertion point to the <merge> block for subsequent statements
   theBuilder.setInsertPoint(mergeBB);
+
+  // Done with the current scope's continue block and merge block.
+  continueStack.pop();
+  breakStack.pop();
 }
 
 void SPIRVEmitter::doIfStmt(const IfStmt *ifStmt) {
@@ -882,10 +991,68 @@ void SPIRVEmitter::doReturnStmt(const ReturnStmt *stmt) {
   }
 }
 
+bool SPIRVEmitter::breakStmtIsLastStmtInCaseStmt(const BreakStmt *breakStmt,
+                                                 const SwitchStmt *switchStmt) {
+  std::vector<const Stmt *> flatSwitch;
+  flattenSwitchStmtAST(switchStmt->getBody(), &flatSwitch);
+  auto iter =
+      std::find(flatSwitch.begin(), flatSwitch.end(), cast<Stmt>(breakStmt));
+  assert(iter != std::end(flatSwitch));
+
+  // We are interested to know what comes after the BreakStmt.
+  ++iter;
+
+  // The break statement is the last statement in its case statement iff:
+  // it is the last statement in the switch, or,
+  // its next statement is either 'default' or 'case'
+  return iter == flatSwitch.end() || isa<CaseStmt>(*iter) ||
+         isa<DefaultStmt>(*iter);
+}
+
+const Stmt *SPIRVEmitter::breakStmtScope(const BreakStmt *breakStmt) {
+  const Stmt *curNode = breakStmt;
+  do {
+    curNode = getImmediateParent(astContext, curNode);
+  } while (curNode && !isLoopStmt(curNode) && !isa<SwitchStmt>(curNode));
+
+  return curNode;
+}
+
 void SPIRVEmitter::doBreakStmt(const BreakStmt *breakStmt) {
+  assert(!theBuilder.isCurrentBasicBlockTerminated());
   uint32_t breakTargetBB = breakStack.top();
   theBuilder.addSuccessor(breakTargetBB);
   theBuilder.createBranch(breakTargetBB);
+
+  // If any statements follow a break statement in a loop or switch, they will
+  // not be executed. For example, StmtB and StmtC below are never executed:
+  //
+  // while (true) {
+  //   StmtA;
+  //   break;
+  //   StmtB;
+  //   StmtC;
+  // }
+  //
+  // To handle such cases, we do not stop tranlsation. We create a new basic
+  // block in which StmtB and StmtC will be translated.
+  // Note that since this basic block is unreachable, BlockReadableOrderVisitor
+  // will not emit it in the final module binary.
+  const Stmt *scope = breakStmtScope(breakStmt);
+  if (
+      // Have unreachable instructions after a break statement in a case of a
+      // switch statement
+      (isa<SwitchStmt>(scope) &&
+       !breakStmtIsLastStmtInCaseStmt(breakStmt,
+                                      dyn_cast<SwitchStmt>(scope))) ||
+      // Have unreachable instructions after a break statement in a loop
+      // statement
+      (isLoopStmt(scope) &&
+       !isLastStmtBeforeControlFlowBranching(astContext, breakStmt))) {
+    const uint32_t unreachableBB =
+        theBuilder.createBasicBlock("unreachable", false);
+    theBuilder.setInsertPoint(unreachableBB);
+  }
 }
 
 void SPIRVEmitter::doSwitchStmt(const SwitchStmt *switchStmt,

tools/clang/lib/SPIRV/SPIRVEmitter.h

@@ -82,6 +82,7 @@ private:
                     llvm::ArrayRef<const Attr *> attrs = {});
   void doWhileStmt(const WhileStmt *, llvm::ArrayRef<const Attr *> attrs = {});
   void doDoStmt(const DoStmt *, llvm::ArrayRef<const Attr *> attrs = {});
+  void doContinueStmt(const ContinueStmt *);
 
   uint32_t doBinaryOperator(const BinaryOperator *expr);
   uint32_t doCallExpr(const CallExpr *callExpr);
@@ -375,6 +376,18 @@ private:
   /// statement.
   void processSwitchStmtUsingIfStmts(const SwitchStmt *switchStmt);
 
+private:
+  /// \brief Returns the statement that the given break statement applies to.
+  /// According to the spec, break statements can only apply to loops (do, for,
+  /// while) or case statements inside a switch statement. The frontend ensures
+  /// this is true (errors out otherwise).
+  const Stmt *breakStmtScope(const BreakStmt *);
+
+  /// \brief Returns true if the given BreakStmt is the last statement inside
+  /// its case statement of the given switch statement. Panics if the given
+  /// break statement is not inside the tree of the given switch statement.
+  bool breakStmtIsLastStmtInCaseStmt(const BreakStmt *, const SwitchStmt *);
+
 private:
   /// \brief Wrapper method to create an error message and report it
   /// in the diagnostic engine associated with this consumer.
@@ -438,6 +451,13 @@ private:
   /// This stack keeps track of the basic blocks to which branching could occur.
   std::stack<uint32_t> breakStack;
 
+  /// Loops (do, while, for) may encounter "continue" statements that alter
+  /// their control flow. At any point the continue statement is observed, the
+  /// control flow jumps to the inner-most scope's continue block.
+  /// This stack keeps track of the basic blocks to which such branching could
+  /// occur.
+  std::stack<uint32_t> continueStack;
+
   /// Maps a given statement to the basic block that is associated with it.
   llvm::DenseMap<const Stmt *, uint32_t> stmtBasicBlock;
 };

tools/clang/test/CodeGenSPIRV/break-stmt.mixed.hlsl

@@ -0,0 +1,87 @@
+// Run: %dxc -T ps_6_0 -E main
+
+void main() {
+  int a, b;
+  bool cond = true;
+
+// CHECK:      OpBranch %while_check
+// CHECK-NEXT: %while_check = OpLabel
+// CHECK-NEXT: [[cond:%\d+]] = OpLoad %bool %cond
+// CHECK-NEXT: OpLoopMerge %while_merge %while_continue None
+// CHECK-NEXT: OpBranchConditional [[cond]] %while_body %while_merge
+  while(cond) {
+// CHECK-NEXT: %while_body = OpLabel
+// CHECK-NEXT: [[b:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: OpSelectionMerge %switch_merge None
+// CHECK-NEXT: OpSwitch [[b]] %switch_default 1 %switch_1 2 %switch_2 5 %switch_5
+    switch(b) {
+// CHECK-NEXT: %switch_1 = OpLabel
+// CHECK-NEXT: OpStore %a %int_1
+// CHECK-NEXT: OpBranch %switch_merge
+      case 1:
+        a = 1;
+        break;    // Break from the case statement.
+// CHECK-NEXT: %switch_2 = OpLabel
+// CHECK-NEXT: OpStore %a %int_3
+// CHECK-NEXT: OpBranch %while_continue
+      case 2: {
+        a = 3;
+        {continue;} // continue for the while loop.
+        a = 4;      // No SPIR-V should be emitted for this statement.
+        break;      // No SPIR-V should be emitted for this statement.
+      }
+// CHECK-NEXT: %switch_5 = OpLabel
+// CHECK-NEXT: OpStore %a %int_5
+// CHECK-NEXT: OpBranch %switch_merge
+      case 5 : {
+        a = 5;
+        {{break;}} // Break from the case statement.
+        a = 6;     // No SPIR-V should be emitted for this statement.
+      }
+// CHECK-NEXT: %switch_default = OpLabel
+// CHECK-NEXT: OpBranch %for_check
+      default:
+      // CHECK-NEXT: %for_check = OpLabel
+      // CHECK-NEXT: {{%\d+}} = OpLoad %int %i
+      // CHECK-NEXT: [[i_lt_10:%\d+]] = OpSLessThan %bool %37 %int_10
+      // CHECK-NEXT: OpLoopMerge %for_merge %for_continue None
+      // CHECK-NEXT: OpBranchConditional [[i_lt_10]] %for_body %for_merge
+        for (int i=0; i<10; ++i) {
+        // CHECK-NEXT: %for_body = OpLabel
+        // CHECK-NEXT: [[cond1:%\d+]] = OpLoad %bool %cond
+        // CHECK-NEXT: OpSelectionMerge %if_merge None
+        // CHECK-NEXT: OpBranchConditional [[cond1]] %if_true %if_false
+          if (cond) {
+          // CHECK-NEXT: %if_true = OpLabel
+          // CHECK-NEXT: OpBranch %for_merge
+            break;    // Break the for loop.
+            break;    // No SPIR-V should be emitted for this statement.
+            continue; // No SPIR-V should be emitted for this statement.
+            ++a;      // No SPIR-V should be emitted for this statement.
+          } else {
+          // CHECK-NEXT: %if_false = OpLabel
+          // CHECK-NEXT: OpBranch %for_continue
+            continue; // continue for the for loop.
+            continue; // No SPIR-V should be emitted for this statement.
+            break;    // No SPIR-V should be emitted for this statement.
+            ++a;      // No SPIR-V should be emitted for this statement.
+          }
+        // CHECK-NEXT: %if_merge = OpLabel
+        // CHECK-NEXT: OpBranch %for_continue
+        // CHECK-NEXT: %for_continue = OpLabel
+        // CHECK:      OpBranch %for_check
+        }
+        // CHECK-NEXT: %for_merge = OpLabel
+        // CHECK-NEXT: OpBranch %switch_merge
+        break;        // Break from the default statement.
+    }
+    // CHECK-NEXT: %switch_merge = OpLabel
+
+    // CHECK-NEXT: OpBranch %while_merge
+    break; // Break the while loop.
+    
+  // CHECK-NEXT: %while_continue = OpLabel
+  // CHECK-NEXT: OpBranch %while_check
+  }
+  // CHECK-NEXT: %while_merge = OpLabel
+}

tools/clang/test/CodeGenSPIRV/do-stmt.break.hlsl

@@ -0,0 +1,71 @@
+// Run: %dxc -T ps_6_0 -E main
+
+int foo() { return true; }
+
+void main() {
+  int val = 0;
+  int i = 0;
+
+// CHECK:      OpBranch %do_while_header
+// CHECK-NEXT: %do_while_header = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge %do_while_continue None
+// CHECK-NEXT: OpBranch %do_while_body
+  do {
+// CHECK-NEXT: %do_while_body = OpLabel
+    ++i;
+// CHECK:      OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %if_true %if_merge
+    if (i > 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %do_while_merge
+      break;
+      break;       // No SPIR-V should be emitted for this statement.
+      val = i;     // No SPIR-V should be emitted for this statement.
+      while(true); // No SPIR-V should be emitted for this statement.
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+    val = i;
+// CHECK:      OpBranch %do_while_merge
+    {{break;}}
+    val = val * 2; // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %do_while_continue = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header %do_while_merge
+  } while (i < 10);
+// CHECK-NEXT: %do_while_merge = OpLabel
+  
+  
+  ////////////////////////////////////////////////////////////////////////////////
+  // Nested do-while loops with break statements                                //
+  // Each break statement should branch to the corresponding loop's break block //
+  ////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: OpBranch %do_while_header_0
+
+// CHECK-NEXT: %do_while_header_0 = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge_0 %do_while_continue_0 None
+// CHECK-NEXT: OpBranch %do_while_body_0
+  do {
+// CHECK-NEXT: %do_while_body_0 = OpLabel
+    ++i;
+// CHECK:      OpBranch %do_while_header_1
+// CHECK-NEXT: %do_while_header_1 = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge_1 %do_while_continue_1 None
+// CHECK-NEXT: OpBranch %do_while_body_1
+    do {
+// CHECK-NEXT: %do_while_body_1 = OpLabel
+      ++val;
+// CHECK:      OpBranch %do_while_merge_1
+      break;
+// CHECK-NEXT: %do_while_continue_1 = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header_1 %do_while_merge_1
+    } while (i < 10);
+// CHECK-NEXT: %do_while_merge_1 = OpLabel
+    --i;
+// CHECK:      OpBranch %do_while_merge_0
+    {break;}
+// CHECK-NEXT: %do_while_continue_0 = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header_0 %do_while_merge_0
+  } while(val < 10);
+// CHECK-NEXT: %do_while_merge_0 = OpLabel
+}

tools/clang/test/CodeGenSPIRV/do-stmt.continue.hlsl

@@ -0,0 +1,74 @@
+// Run: %dxc -T ps_6_0 -E main
+
+int foo() { return true; }
+
+void main() {
+  int val = 0;
+  int i = 0;
+
+// CHECK:      OpBranch %do_while_header
+// CHECK-NEXT: %do_while_header = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge %do_while_continue None
+// CHECK-NEXT: OpBranch %do_while_body
+  do {
+// CHECK-NEXT: %do_while_body = OpLabel
+    ++i;
+// CHECK:      OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %if_true %if_merge
+    if (i > 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %do_while_continue
+      {{continue;}}
+      val = i;     // No SPIR-V should be emitted for this statement.
+      while(true); // No SPIR-V should be emitted for this statement.
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+    val = i;
+// CHECK:      OpBranch %do_while_continue
+    continue;
+    val = val * 2; // No SPIR-V should be emitted for this statement.
+    continue;      // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %do_while_continue = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header %do_while_merge
+  } while (i < 10);
+// CHECK-NEXT: %do_while_merge = OpLabel
+
+
+
+  //////////////////////////////////////////////////////////////////////////////////////
+  // Nested do-while loops with continue statements                                   //
+  // Each continue statement should branch to the corresponding loop's continue block //
+  //////////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: OpBranch %do_while_header_0
+
+// CHECK-NEXT: %do_while_header_0 = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge_0 %do_while_continue_0 None
+// CHECK-NEXT: OpBranch %do_while_body_0
+  do {
+// CHECK-NEXT: %do_while_body_0 = OpLabel
+    ++i;
+// CHECK:      OpBranch %do_while_header_1
+// CHECK-NEXT: %do_while_header_1 = OpLabel
+// CHECK-NEXT: OpLoopMerge %do_while_merge_1 %do_while_continue_1 None
+// CHECK-NEXT: OpBranch %do_while_body_1
+    do {
+// CHECK-NEXT: %do_while_body_1 = OpLabel
+      ++val;
+// CHECK:      OpBranch %do_while_continue_1
+      continue;
+// CHECK-NEXT: %do_while_continue_1 = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header_1 %do_while_merge_1
+    } while (i < 10);
+// CHECK-NEXT: %do_while_merge_1 = OpLabel
+    --i;
+// CHECK:      OpBranch %do_while_continue_0
+    continue;
+    continue;  // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %do_while_continue_0 = OpLabel
+// CHECK:      OpBranchConditional {{%\d+}} %do_while_header_0 %do_while_merge_0
+  } while(val < 10);
+// CHECK-NEXT: %do_while_merge_0 = OpLabel
+}

tools/clang/test/CodeGenSPIRV/for-stmt.break.hlsl

@@ -0,0 +1,71 @@
+// Run: %dxc -T ps_6_0 -E main
+
+void main() {
+  int val = 0;
+  
+// CHECK:      OpBranch %for_check
+// CHECK-NEXT: %for_check = OpLabel
+// CHECK:      OpLoopMerge %for_merge %for_continue None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body %for_merge
+  for (int i = 0; i < 10; ++i) {
+// CHECK-NEXT: %for_body = OpLabel
+// CHECK:      OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %if_true %if_merge
+    if (i < 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %for_merge
+      break;
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+    val = i;
+// CHECK:      OpBranch %for_merge
+    {break;}
+    break;       // No SPIR-V should be emitted for this statement.
+    val++;       // No SPIR-V should be emitted for this statement.
+    while(true); // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue = OpLabel
+// CHECK:      OpBranch %for_check
+  }
+// CHECK-NEXT: %for_merge = OpLabel
+
+// CHECK-NEXT: OpBranch %for_check_0
+
+  ////////////////////////////////////////////////////////////////////////////////
+  // Nested for loops with break statements                                     //
+  // Each break statement should branch to the corresponding loop's break block //
+  ////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: %for_check_0 = OpLabel
+// CHECK:      OpLoopMerge %for_merge_0 %for_continue_0 None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body_0 %for_merge_0
+  for (int j = 0; j < 10; ++j) {
+// CHECK-NEXT: %for_body_0 = OpLabel
+    val = j+5;
+// CHECK:      OpBranch %for_check_1
+
+// CHECK-NEXT: %for_check_1 = OpLabel
+// CHECK:      OpLoopMerge %for_merge_1 %for_continue_1 None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body_1 %for_merge_1
+    for ( ; val < 20; ++val) {
+// CHECK-NEXT: %for_body_1 = OpLabel
+      int k = val + j;
+// CHECK:      OpBranch %for_merge_1
+      {{break;}}
+      k++; // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue_1 = OpLabel
+// CHECK:      OpBranch %for_check_1
+    }
+// CHECK-NEXT: %for_merge_1 = OpLabel
+    val--;
+// CHECK:      OpBranch %for_merge_0
+    break;
+    break;        // No SPIR-V should be emitted for this statement.
+    val = val*10; // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue_0 = OpLabel
+// CHECK:      OpBranch %for_check_0
+  }
+// CHECK-NEXT: %for_merge_0 = OpLabel
+}

tools/clang/test/CodeGenSPIRV/for-stmt.continue.hlsl

@@ -0,0 +1,73 @@
+// Run: %dxc -T ps_6_0 -E main
+
+void main() {
+  int val = 0;
+  
+// CHECK:      OpBranch %for_check
+// CHECK-NEXT: %for_check = OpLabel
+// CHECK:      OpLoopMerge %for_merge %for_continue None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body %for_merge
+  for (int i = 0; i < 10; ++i) {
+// CHECK-NEXT: %for_body = OpLabel
+// CHECK:      OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %if_true %if_merge
+    if (i < 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %for_continue
+      continue;
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+    val = i;
+// CHECK:      OpBranch %for_continue
+    {continue;}
+    val++;       // No SPIR-V should be emitted for this statement.
+    continue;    // No SPIR-V should be emitted for this statement.
+    while(true); // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue = OpLabel
+// CHECK:      OpBranch %for_check
+  }
+// CHECK-NEXT: %for_merge = OpLabel
+
+// CHECK-NEXT: OpBranch %for_check_0
+
+
+
+  //////////////////////////////////////////////////////////////////////////////////////
+  // Nested for loops with continue statements                                        //
+  // Each continue statement should branch to the corresponding loop's continue block //
+  //////////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: %for_check_0 = OpLabel
+// CHECK:      OpLoopMerge %for_merge_0 %for_continue_0 None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body_0 %for_merge_0
+  for (int j = 0; j < 10; ++j) {
+// CHECK-NEXT: %for_body_0 = OpLabel
+    val = j+5;
+// CHECK:      OpBranch %for_check_1
+
+// CHECK-NEXT: %for_check_1 = OpLabel
+// CHECK:      OpLoopMerge %for_merge_1 %for_continue_1 None
+// CHECK-NEXT: OpBranchConditional {{%\d+}} %for_body_1 %for_merge_1
+    for ( ; val < 20; ++val) {
+// CHECK-NEXT: %for_body_1 = OpLabel
+      int k = val + j;
+// CHECK:      OpBranch %for_continue_1
+      continue;
+      k++;      // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue_1 = OpLabel
+// CHECK:      OpBranch %for_check_1
+    }
+// CHECK-NEXT: %for_merge_1 = OpLabel
+    val--;
+// CHECK:      OpBranch %for_continue_0
+    continue;
+    continue;     // No SPIR-V should be emitted for this statement.
+    val = val*10; // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %for_continue_0 = OpLabel
+// CHECK:      OpBranch %for_check_0
+  }
+// CHECK-NEXT: %for_merge_0 = OpLabel
+}

tools/clang/test/CodeGenSPIRV/while-stmt.break.hlsl

@@ -0,0 +1,86 @@
+// Run: %dxc -T ps_6_0 -E main
+
+int foo() { return true; }
+
+void main() {
+  int val = 0;
+  int i = 0;
+
+  
+// CHECK:      OpBranch %while_check
+// CHECK:      %while_check = OpLabel
+// CHECK:      [[i_lt_10:%\d+]] = OpSLessThan %bool {{%\d+}} %int_10
+// CHECK-NEXT: OpLoopMerge %while_merge %while_continue None
+// CHECK-NEXT: OpBranchConditional [[i_lt_10]] %while_body %while_merge
+  while (i < 10) {
+// CHECK-NEXT: %while_body = OpLabel
+    val = i;
+// CHECK:      [[val_gt_5:%\d+]] = OpSGreaterThan %bool {{%\d+}} %int_5
+// CHECK-NEXT: OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional [[val_gt_5]] %if_true %if_merge
+    if (val > 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %while_merge
+      break;
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+
+// CHECK:      [[val_gt_6:%\d+]] = OpSGreaterThan %bool {{%\d+}} %int_6
+// CHECK-NEXT: OpSelectionMerge %if_merge_0 None
+// CHECK-NEXT: OpBranchConditional [[val_gt_6]] %if_true_0 %if_merge_0
+    if (val > 6) {
+// CHECK-NEXT: %if_true_0 = OpLabel
+// CHECK-NEXT: OpBranch %while_merge
+      break;
+      break;       // No SPIR-V should be emitted for this statement.
+      val++;       // No SPIR-V should be emitted for this statement.
+      while(true); // No SPIR-V should be emitted for this statement.
+      --i;         // No SPIR-V should be emitted for this statement.
+    }
+// CHECK-NEXT: %if_merge_0 = OpLabel
+
+// CHECK-NEXT: OpBranch %while_continue
+// CHECK-NEXT: %while_continue = OpLabel
+// CHECK-NEXT: OpBranch %while_check
+  }
+
+// CHECK-NEXT: %while_merge = OpLabel
+
+
+
+  ////////////////////////////////////////////////////////////////////////////////
+  // Nested while loops with break statements                                   //
+  // Each break statement should branch to the corresponding loop's break block //
+  ////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: OpBranch %while_check_0
+// CHECK-NEXT: %while_check_0 = OpLabel
+// CHECK-NEXT: OpLoopMerge %while_merge_0 %while_continue_0 None
+// CHECK-NEXT: OpBranchConditional %true %while_body_0 %while_merge_0
+  while (true) {
+// CHECK-NEXT: %while_body_0 = OpLabel
+    i++;
+
+// CHECK:      OpBranch %while_check_1
+// CHECK-NEXT: %while_check_1 = OpLabel
+// CHECK:      [[i_lt_20:%\d+]] = OpSLessThan %bool {{%\d+}} %int_20
+// CHECK-NEXT: OpLoopMerge %while_merge_1 %while_continue_1 None
+// CHECK-NEXT: OpBranchConditional [[i_lt_20]] %while_body_1 %while_merge_1
+    while(i<20) {
+// CHECK-NEXT: %while_body_1 = OpLabel
+      val = i;
+// CHECK:      OpBranch %while_merge_1
+      {{break;}}
+// CHECK-NEXT: %while_continue_1 = OpLabel
+// CHECK-NEXT: OpBranch %while_check_1
+    }
+// CHECK-NEXT: %while_merge_1 = OpLabel
+    --i;
+// CHECK:      OpBranch %while_merge_0
+    break;
+// CHECK-NEXT: %while_continue_0 = OpLabel
+// CHECK-NEXT: OpBranch %while_check_0
+  }
+// CHECK-NEXT: %while_merge_0 = OpLabel
+
+}

tools/clang/test/CodeGenSPIRV/while-stmt.continue.hlsl

@@ -0,0 +1,88 @@
+// Run: %dxc -T ps_6_0 -E main
+
+int foo() { return true; }
+
+void main() {
+  int val = 0;
+  int i = 0;
+
+  
+// CHECK:      OpBranch %while_check
+// CHECK:      %while_check = OpLabel
+// CHECK:      [[i_lt_10:%\d+]] = OpSLessThan %bool {{%\d+}} %int_10
+// CHECK-NEXT: OpLoopMerge %while_merge %while_continue None
+// CHECK-NEXT: OpBranchConditional [[i_lt_10]] %while_body %while_merge
+  while (i < 10) {
+// CHECK-NEXT: %while_body = OpLabel
+    val = i;
+// CHECK:      [[val_gt_5:%\d+]] = OpSGreaterThan %bool {{%\d+}} %int_5
+// CHECK-NEXT: OpSelectionMerge %if_merge None
+// CHECK-NEXT: OpBranchConditional [[val_gt_5]] %if_true %if_merge
+    if (val > 5) {
+// CHECK-NEXT: %if_true = OpLabel
+// CHECK-NEXT: OpBranch %while_continue
+      continue;
+    }
+// CHECK-NEXT: %if_merge = OpLabel
+
+// CHECK:      [[val_gt_6:%\d+]] = OpSGreaterThan %bool {{%\d+}} %int_6
+// CHECK-NEXT: OpSelectionMerge %if_merge_0 None
+// CHECK-NEXT: OpBranchConditional [[val_gt_6]] %if_true_0 %if_merge_0
+    if (val > 6) {
+// CHECK-NEXT: %if_true_0 = OpLabel
+// CHECK-NEXT: OpBranch %while_continue
+      {{continue;}}
+      val++;       // No SPIR-V should be emitted for this statement.
+      continue;    // No SPIR-V should be emitted for this statement.
+      while(true); // No SPIR-V should be emitted for this statement.
+      --i;         // No SPIR-V should be emitted for this statement.
+    }
+// CHECK-NEXT: %if_merge_0 = OpLabel
+
+// CHECK-NEXT: OpBranch %while_continue
+// CHECK-NEXT: %while_continue = OpLabel
+// CHECK-NEXT: OpBranch %while_check
+  }
+
+// CHECK-NEXT: %while_merge = OpLabel
+
+
+
+  //////////////////////////////////////////////////////////////////////////////////////
+  // Nested while loops with continue statements                                      //
+  // Each continue statement should branch to the corresponding loop's continue block //
+  //////////////////////////////////////////////////////////////////////////////////////
+
+// CHECK-NEXT: OpBranch %while_check_0
+// CHECK-NEXT: %while_check_0 = OpLabel
+// CHECK-NEXT: OpLoopMerge %while_merge_0 %while_continue_0 None
+// CHECK-NEXT: OpBranchConditional %true %while_body_0 %while_merge_0
+  while (true) {
+// CHECK-NEXT: %while_body_0 = OpLabel
+    i++;
+
+// CHECK:      OpBranch %while_check_1
+// CHECK-NEXT: %while_check_1 = OpLabel
+// CHECK:      [[i_lt_20:%\d+]] = OpSLessThan %bool {{%\d+}} %int_20
+// CHECK-NEXT: OpLoopMerge %while_merge_1 %while_continue_1 None
+// CHECK-NEXT: OpBranchConditional [[i_lt_20]] %while_body_1 %while_merge_1
+    while(i<20) {
+// CHECK-NEXT: %while_body_1 = OpLabel
+      val = i;
+// CHECK:      OpBranch %while_continue_1
+      continue;
+// CHECK-NEXT: %while_continue_1 = OpLabel
+// CHECK-NEXT: OpBranch %while_check_1
+    }
+// CHECK-NEXT: %while_merge_1 = OpLabel
+    --i;
+// CHECK:      OpBranch %while_continue_0
+    continue;
+    continue;  // No SPIR-V should be emitted for this statement.
+
+// CHECK-NEXT: %while_continue_0 = OpLabel
+// CHECK-NEXT: OpBranch %while_check_0
+  }
+// CHECK-NEXT: %while_merge_0 = OpLabel
+
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -206,14 +206,23 @@ TEST_F(FileTest, SwitchStmtUsingIfStmt) {
 // For for statements
 TEST_F(FileTest, ForStmtPlainAssign) { runFileTest("for-stmt.plain.hlsl"); }
 TEST_F(FileTest, ForStmtNestedForStmt) { runFileTest("for-stmt.nested.hlsl"); }
+TEST_F(FileTest, ForStmtContinue) { runFileTest("for-stmt.continue.hlsl"); }
+TEST_F(FileTest, ForStmtBreak) { runFileTest("for-stmt.break.hlsl"); }
 
 // For while statements
 TEST_F(FileTest, WhileStmtPlain) { runFileTest("while-stmt.plain.hlsl"); }
 TEST_F(FileTest, WhileStmtNested) { runFileTest("while-stmt.nested.hlsl"); }
+TEST_F(FileTest, WhileStmtContinue) { runFileTest("while-stmt.continue.hlsl"); }
+TEST_F(FileTest, WhileStmtBreak) { runFileTest("while-stmt.break.hlsl"); }
 
 // For do statements
 TEST_F(FileTest, DoStmtPlain) { runFileTest("do-stmt.plain.hlsl"); }
 TEST_F(FileTest, DoStmtNested) { runFileTest("do-stmt.nested.hlsl"); }
+TEST_F(FileTest, DoStmtContinue) { runFileTest("do-stmt.continue.hlsl"); }
+TEST_F(FileTest, DoStmtBreak) { runFileTest("do-stmt.break.hlsl"); }
+
+// For break statements (mix of breaks in loops and switch)
+TEST_F(FileTest, BreakStmtMixed) { runFileTest("break-stmt.mixed.hlsl"); }
 
 // For control flows
 TEST_F(FileTest, ControlFlowNestedIfForStmt) { runFileTest("cf.if.for.hlsl"); }