[ir] Add explicit continue branches.

This CL adds a `Continue` instruction into the IR and uses it when
branching into the continue block.

Bug: tint:1718
Change-Id: If54afe6f53b587f1d8a99afd920b94ebdcb608e5
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/134462
Reviewed-by: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>

src/tint/BUILD.gn

@@ -1224,6 +1224,8 @@ if (tint_build_ir) {
       "ir/constant.h",
       "ir/construct.cc",
       "ir/construct.h",
+      "ir/continue.cc",
+      "ir/continue.h",
       "ir/convert.cc",
       "ir/convert.h",
       "ir/debug.cc",

src/tint/CMakeLists.txt

@@ -732,6 +732,8 @@ if(${TINT_BUILD_IR})
     ir/constant.h
     ir/construct.cc
     ir/construct.h
+    ir/continue.cc
+    ir/continue.h
     ir/convert.cc
     ir/convert.h
     ir/debug.cc

src/tint/ir/builder.cc

@@ -213,6 +213,10 @@ ir::Return* Builder::Return(Function* func, utils::VectorRef<Value*> args) {
     return ir.values.Create<ir::Return>(func, args);
 }
 
+ir::Continue* Builder::Continue(Loop* loop) {
+    return ir.values.Create<ir::Continue>(loop);
+}
+
 ir::BlockParam* Builder::BlockParam(const type::Type* type) {
     return ir.values.Create<ir::BlockParam>(type);
 }

src/tint/ir/builder.h

@@ -24,6 +24,7 @@
 #include "src/tint/ir/builtin.h"
 #include "src/tint/ir/constant.h"
 #include "src/tint/ir/construct.h"
+#include "src/tint/ir/continue.h"
 #include "src/tint/ir/convert.h"
 #include "src/tint/ir/discard.h"
 #include "src/tint/ir/function.h"
@@ -336,6 +337,11 @@ class Builder {
     /// @returns the instruction
     ir::Return* Return(Function* func, utils::VectorRef<Value*> args = {});
 
+    /// Creates a continue instruction
+    /// @param loop the loop being continued
+    /// @returns the instruction
+    ir::Continue* Continue(Loop* loop);
+
     /// Creates a branch declaration
     /// @param to the node being branched too
     /// @param args the branch arguments

src/tint/ir/continue.cc

@@ -0,0 +1,31 @@
+// Copyright 2023 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/ir/continue.h"
+
+#include "src/tint/ir/loop.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::ir::Continue);
+
+namespace tint::ir {
+
+Continue::Continue(ir::Loop* loop) : Base(utils::Empty), loop_(loop) {
+    TINT_ASSERT(IR, loop_);
+    loop_->AddUsage(this);
+    loop_->Continuing()->AddInboundBranch(this);
+}
+
+Continue::~Continue() = default;
+
+}  // namespace tint::ir

src/tint/ir/continue.h

@@ -0,0 +1,45 @@
+// Copyright 2023 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_TINT_IR_CONTINUE_H_
+#define SRC_TINT_IR_CONTINUE_H_
+
+#include "src/tint/ir/branch.h"
+#include "src/tint/utils/castable.h"
+
+// Forward declarations
+namespace tint::ir {
+class Loop;
+}  // namespace tint::ir
+
+namespace tint::ir {
+
+/// A continue instruction.
+class Continue : public utils::Castable<Continue, Branch> {
+  public:
+    /// Constructor
+    /// @param loop the loop owning the continue block
+    explicit Continue(ir::Loop* loop);
+    ~Continue() override;
+
+    /// @returns the loop owning the continue block
+    const ir::Loop* Loop() const { return loop_; }
+
+  private:
+    ir::Loop* loop_ = nullptr;
+};
+
+}  // namespace tint::ir
+
+#endif  // SRC_TINT_IR_CONTINUE_H_

src/tint/ir/debug.cc

@@ -18,6 +18,7 @@
 #include <unordered_set>
 
 #include "src/tint/ir/block.h"
+#include "src/tint/ir/continue.h"
 #include "src/tint/ir/if.h"
 #include "src/tint/ir/loop.h"
 #include "src/tint/ir/return.h"
@@ -70,6 +71,8 @@ std::string Debug::AsDotGraph(const Module* mod) {
 
                          if (b->Branch()->Is<ir::Return>()) {
                              return;
+                         } else if (auto* cont = b->Branch()->As<ir::Continue>()) {
+                             Graph(cont->Loop()->Continuing());
                          } else {
                              Graph(b->Branch()->To());
                          }

src/tint/ir/disassembler.cc

@@ -23,6 +23,7 @@
 #include "src/tint/ir/block.h"
 #include "src/tint/ir/builtin.h"
 #include "src/tint/ir/construct.h"
+#include "src/tint/ir/continue.h"
 #include "src/tint/ir/convert.h"
 #include "src/tint/ir/discard.h"
 #include "src/tint/ir/if.h"
@@ -411,6 +412,8 @@ void Disassembler::EmitBranch(const Branch* b) {
     std::string suffix = "";
     if (b->Is<ir::Return>()) {
         out_ << "ret";
+    } else if (auto* cont = b->As<ir::Continue>()) {
+        out_ << "continue %b" << IdOf(cont->Loop()->Continuing());
     } else {
         out_ << "br %b" << IdOf(b->To());
         if (b->To()->Is<RootTerminator>()) {

src/tint/ir/from_program.cc

@@ -166,6 +166,8 @@ class Impl {
         diagnostics_.add_error(tint::diag::System::IR, err, s);
     }
 
+    bool NeedBranch() { return current_flow_block_ && !current_flow_block_->HasBranchTarget(); }
+
     void SetBranch(Branch* br) {
         TINT_ASSERT(IR, current_flow_block_);
         TINT_ASSERT(IR, !current_flow_block_->HasBranchTarget());
@@ -174,13 +176,6 @@ class Impl {
         current_flow_block_ = nullptr;
     }
 
-    void SetBranchIfNeeded(Branch* br) {
-        if (!current_flow_block_ || current_flow_block_->HasBranchTarget()) {
-            return;
-        }
-        SetBranch(br);
-    }
-
     void BranchTo(Block* node, utils::VectorRef<Value*> args = {}) {
         TINT_ASSERT(IR, current_flow_block_);
         TINT_ASSERT(IR, !current_flow_block_->HasBranchTarget());
@@ -190,7 +185,7 @@ class Impl {
     }
 
     void BranchToIfNeeded(Block* node) {
-        if (!current_flow_block_ || current_flow_block_->HasBranchTarget()) {
+        if (!NeedBranch()) {
             return;
         }
         BranchTo(node);
@@ -346,13 +341,13 @@ class Impl {
         }
         ir_func->SetParams(params);
 
-        {
+        current_flow_block_ = ir_func->StartTarget();
-            current_flow_block_ = ir_func->StartTarget();
+        EmitBlock(ast_func->body);
-            EmitBlock(ast_func->body);
 
-            // If the branch target has already been set then a `return` was called. Only set in
+        // If the branch target has already been set then a `return` was called. Only set in
-            // the case where `return` wasn't called.
+        // the case where `return` wasn't called.
-            SetBranchIfNeeded(builder_.Return(current_function_));
+        if (NeedBranch()) {
+            SetBranch(builder_.Return(current_function_));
         }
 
         TINT_ASSERT(IR, control_stack_.IsEmpty());
@@ -366,7 +361,7 @@ class Impl {
 
             // If the current flow block has a branch target then the rest of the statements in
             // this block are dead code. Skip them.
-            if (!current_flow_block_ || current_flow_block_->HasBranchTarget()) {
+            if (!NeedBranch()) {
                 break;
             }
         }
@@ -580,22 +575,28 @@ class Impl {
             TINT_DEFER(scopes_.Pop());
             EmitStatements(stmt->body->statements);
 
-            // The current block didn't `break`, `return` or `continue`, go to the continuing
+            // The current block didn't `break`, `return` or `continue`, go to the continuing block.
-            // block.
+            if (NeedBranch()) {
-            BranchToIfNeeded(loop_inst->Continuing());
+                SetBranch(builder_.Continue(loop_inst));
-
-            current_flow_block_ = loop_inst->Continuing();
-            if (stmt->continuing) {
-                EmitBlock(stmt->continuing);
             }
 
-            // Branch back to the start node if the continue target didn't branch out already
+            if (IsConnected(loop_inst->Continuing(), 0)) {
-            BranchToIfNeeded(loop_inst->Start());
+                // Note, even if there is no continuing block, we may have branched into the
+                // continue so we have to set the current block and then emit the branch if needed
+                // below otherwise empty continuing blocks will fail to branch back to the start
+                // block.
+                current_flow_block_ = loop_inst->Continuing();
+                if (stmt->continuing) {
+                    EmitBlock(stmt->continuing);
+                }
+                // Branch back to the start node if the continue target didn't branch out already
+                BranchToIfNeeded(loop_inst->Start());
+            }
         }
 
         // The loop merge can get disconnected if the loop returns directly, or the continuing
         // target branches, eventually, to the merge, but nothing branched to the
-        // Continuing().target.
+        // Continuing() block.
         current_flow_block_ = loop_inst->Merge();
         if (!IsConnected(loop_inst->Merge(), 0)) {
             current_flow_block_ = nullptr;
@@ -629,7 +630,9 @@ class Impl {
             current_flow_block_ = if_inst->Merge();
             EmitBlock(stmt->body);
 
-            BranchToIfNeeded(loop_inst->Continuing());
+            if (NeedBranch()) {
+                SetBranch(builder_.Continue(loop_inst));
+            }
         }
         // The while loop always has a path to the Merge().target as the break statement comes
         // before anything inside the loop.
@@ -640,8 +643,6 @@ class Impl {
         auto* loop_inst = builder_.CreateLoop();
         current_flow_block_->Instructions().Push(loop_inst);
 
-        loop_inst->Continuing()->Instructions().Push(builder_.Branch(loop_inst->Start()));
-
         // Make sure the initializer ends up in a contained scope
         scopes_.Push();
         TINT_DEFER(scopes_.Pop());
@@ -673,11 +674,14 @@ class Impl {
             }
 
             EmitBlock(stmt->body);
-            BranchToIfNeeded(loop_inst->Continuing());
+            if (NeedBranch()) {
+                SetBranch(builder_.Continue(loop_inst));
+            }
 
             if (stmt->continuing) {
                 current_flow_block_ = loop_inst->Continuing();
                 EmitStatement(stmt->continuing);
+                loop_inst->Continuing()->Instructions().Push(builder_.Branch(loop_inst->Start()));
             }
         }
 
@@ -752,7 +756,7 @@ class Impl {
         TINT_ASSERT(IR, current_control);
 
         if (auto* c = current_control->As<Loop>()) {
-            BranchTo(c->Continuing());
+            SetBranch(builder_.Continue(c));
         } else {
             TINT_UNREACHABLE(IR, diagnostics_);
         }

src/tint/ir/from_program_test.cc

@@ -304,18 +304,13 @@ TEST_F(IR_BuilderImplTest, IfStatement_JumpChainToMerge) {
     if true [t: %b2, f: %b3, m: %b4]
       # True block
       %b2 = block {
-        loop [s: %b5, c: %b6, m: %b7]
+        loop [s: %b5, m: %b6]
           %b5 = block {
-            br %b7
+            br %b6
-          }
-
-          # Continuing block
-          %b6 = block {
-            br %b5
           }
 
         # Merge block
-        %b7 = block {
+        %b6 = block {
           br %b4
         }
 
@@ -350,25 +345,20 @@ TEST_F(IR_BuilderImplTest, Loop_WithBreak) {
 
     ASSERT_EQ(1u, m.functions.Length());
 
-    EXPECT_EQ(2u, flow->Start()->InboundBranches().Length());
+    EXPECT_EQ(1u, flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, flow->Continuing()->InboundBranches().Length());
     EXPECT_EQ(1u, flow->Merge()->InboundBranches().Length());
 
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
   %b1 = block {
-    loop [s: %b2, c: %b3, m: %b4]
+    loop [s: %b2, m: %b3]
       %b2 = block {
-        br %b4
+        br %b3
-      }
-
-      # Continuing block
-      %b3 = block {
-        br %b2
       }
 
     # Merge block
-    %b4 = block {
+    %b3 = block {
       ret
     }
 
@@ -418,7 +408,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinue) {
 
         # Merge block
         %b7 = block {
-          br %b3
+          continue %b3
         }
 
 
@@ -466,7 +456,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinuing_BreakIf) {
   %b1 = block {
     loop [s: %b2, c: %b3, m: %b4]
       %b2 = block {
-        br %b3
+        continue %b3
       }
 
       # Continuing block
@@ -516,7 +506,7 @@ TEST_F(IR_BuilderImplTest, Loop_Continuing_Body_Scope) {
   %b1 = block {
     loop [s: %b2, c: %b3, m: %b4]
       %b2 = block {
-        br %b3
+        continue %b3
       }
 
       # Continuing block
@@ -589,7 +579,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithReturn) {
 
         # Merge block
         %b6 = block {
-          br %b3
+          continue %b3
         }
 
 
@@ -619,22 +609,17 @@ TEST_F(IR_BuilderImplTest, Loop_WithOnlyReturn) {
 
     ASSERT_EQ(1u, m.functions.Length());
 
-    EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
+    EXPECT_EQ(1u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, loop_flow->Continuing()->InboundBranches().Length());
     EXPECT_EQ(0u, loop_flow->Merge()->InboundBranches().Length());
 
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
   %b1 = block {
-    loop [s: %b2, c: %b3]
+    loop [s: %b2]
       %b2 = block {
         ret
       }
-      # Continuing block
-      %b3 = block {
-        br %b2
-      }
-
 
   }
 
@@ -663,57 +648,17 @@ TEST_F(IR_BuilderImplTest, Loop_WithOnlyReturn_ContinuingBreakIf) {
 
     ASSERT_EQ(1u, m.functions.Length());
 
-    EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
+    EXPECT_EQ(1u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, loop_flow->Continuing()->InboundBranches().Length());
-    EXPECT_EQ(1u, loop_flow->Merge()->InboundBranches().Length());
+    EXPECT_EQ(0u, loop_flow->Merge()->InboundBranches().Length());
 
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
   %b1 = block {
-    loop [s: %b2, c: %b3, m: %b4]
+    loop [s: %b2]
       %b2 = block {
         ret
       }
-      # Continuing block
-      %b3 = block {
-        if true [t: %b5, f: %b6, m: %b7]
-          # True block
-          %b5 = block {
-            br %b4
-          }
-
-          # False block
-          %b6 = block {
-            br %b7
-          }
-
-        # Merge block
-        %b7 = block {
-          br %b2
-        }
-
-
-      }
-
-    # Merge block
-    %b4 = block {
-      if true [t: %b8, f: %b9, m: %b10]
-        # True block
-        %b8 = block {
-          ret
-        }
-        # False block
-        %b9 = block {
-          br %b10
-        }
-
-      # Merge block
-      %b10 = block {
-        ret
-      }
-
-    }
-
 
   }
 
@@ -735,7 +680,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithIf_BothBranchesBreak) {
 
     ASSERT_EQ(1u, m.functions.Length());
 
-    EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
+    EXPECT_EQ(1u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, loop_flow->Continuing()->InboundBranches().Length());
     EXPECT_EQ(2u, loop_flow->Merge()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->True()->InboundBranches().Length());
@@ -745,29 +690,24 @@ TEST_F(IR_BuilderImplTest, Loop_WithIf_BothBranchesBreak) {
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
   %b1 = block {
-    loop [s: %b2, c: %b3, m: %b4]
+    loop [s: %b2, m: %b3]
       %b2 = block {
-        if true [t: %b5, f: %b6]
+        if true [t: %b4, f: %b5]
           # True block
-          %b5 = block {
+          %b4 = block {
-            br %b4
+            br %b3
           }
 
           # False block
-          %b6 = block {
+          %b5 = block {
-            br %b4
+            br %b3
           }
 
 
       }
 
-      # Continuing block
-      %b3 = block {
-        br %b2
-      }
-
     # Merge block
-    %b4 = block {
+    %b3 = block {
       ret
     }
 
@@ -817,7 +757,7 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
               if true [t: %b11, f: %b12, m: %b13]
                 # True block
                 %b11 = block {
-                  br %b6
+                  continue %b6
                 }
 
                 # False block
@@ -827,7 +767,7 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
               # Merge block
               %b13 = block {
-                br %b6
+                continue %b6
               }
 
 
@@ -838,46 +778,41 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
           # Continuing block
           %b6 = block {
-            loop [s: %b14, c: %b15, m: %b16]
+            loop [s: %b14, m: %b15]
               %b14 = block {
-                br %b16
+                br %b15
-              }
-
-              # Continuing block
-              %b15 = block {
-                br %b14
               }
 
             # Merge block
-            %b16 = block {
+            %b15 = block {
-              loop [s: %b17, c: %b18, m: %b19]
+              loop [s: %b16, c: %b17, m: %b18]
-                %b17 = block {
+                %b16 = block {
-                  br %b18
+                  continue %b17
                 }
 
                 # Continuing block
-                %b18 = block {
+                %b17 = block {
-                  if true [t: %b20, f: %b21, m: %b22]
+                  if true [t: %b19, f: %b20, m: %b21]
                     # True block
-                    %b20 = block {
+                    %b19 = block {
-                      br %b19
+                      br %b18
                     }
 
                     # False block
-                    %b21 = block {
+                    %b20 = block {
-                      br %b22
+                      br %b21
                     }
 
                   # Merge block
-                  %b22 = block {
+                  %b21 = block {
-                    br %b17
+                    br %b16
                   }
 
 
                 }
 
               # Merge block
-              %b19 = block {
+              %b18 = block {
                 br %b5
               }
 
@@ -889,20 +824,20 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
         # Merge block
         %b7 = block {
-          if true [t: %b23, f: %b24, m: %b25]
+          if true [t: %b22, f: %b23, m: %b24]
             # True block
-            %b23 = block {
+            %b22 = block {
               br %b4
             }
 
             # False block
-            %b24 = block {
+            %b23 = block {
-              br %b25
+              br %b24
             }
 
           # Merge block
-          %b25 = block {
+          %b24 = block {
-            br %b3
+            continue %b3
           }
 
 
@@ -968,7 +903,7 @@ TEST_F(IR_BuilderImplTest, While) {
 
         # Merge block
         %b7 = block {
-          br %b3
+          continue %b3
         }
 
 
@@ -1103,25 +1038,20 @@ TEST_F(IR_BuilderImplTest, For_NoInitCondOrContinuing) {
 
     ASSERT_EQ(1u, m.functions.Length());
 
-    EXPECT_EQ(2u, flow->Start()->InboundBranches().Length());
+    EXPECT_EQ(1u, flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, flow->Continuing()->InboundBranches().Length());
     EXPECT_EQ(1u, flow->Merge()->InboundBranches().Length());
 
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
   %b1 = block {
-    loop [s: %b2, c: %b3, m: %b4]
+    loop [s: %b2, m: %b3]
       %b2 = block {
-        br %b4
+        br %b3
-      }
-
-      # Continuing block
-      %b3 = block {
-        br %b2
       }
 
     # Merge block
-    %b4 = block {
+    %b3 = block {
       ret
     }