[ir] Add BreakIf instruction.

This CL adds a BreakIf instruction to the IR to instruct a loop to
go break based on a condition or to iterate the loop.

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

src/tint/BUILD.gn

@@ -1214,6 +1214,8 @@ if (tint_build_ir) {
       "ir/block_param.h",
       "ir/branch.cc",
       "ir/branch.h",
+      "ir/break_if.cc",
+      "ir/break_if.h",
       "ir/builder.cc",
       "ir/builder.h",
       "ir/builtin.cc",

src/tint/CMakeLists.txt

@@ -722,6 +722,8 @@ if(${TINT_BUILD_IR})
     ir/block_param.h
     ir/branch.cc
     ir/branch.h
+    ir/break_if.cc
+    ir/break_if.h
     ir/builder.cc
     ir/builder.h
     ir/builtin.cc

src/tint/ir/break_if.cc

@@ -0,0 +1,35 @@
+// 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/break_if.h"
+
+#include "src/tint/ir/loop.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::ir::BreakIf);
+
+namespace tint::ir {
+
+BreakIf::BreakIf(Value* condition, ir::Loop* loop)
+    : Base(utils::Empty), condition_(condition), loop_(loop) {
+    TINT_ASSERT(IR, condition_);
+    TINT_ASSERT(IR, loop_);
+    condition_->AddUsage(this);
+    loop_->AddUsage(this);
+    loop_->Start()->AddInboundBranch(this);
+    loop_->Merge()->AddInboundBranch(this);
+}
+
+BreakIf::~BreakIf() = default;
+
+}  // namespace tint::ir

src/tint/ir/break_if.h

@@ -0,0 +1,51 @@
+// 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_BREAK_IF_H_
+#define SRC_TINT_IR_BREAK_IF_H_
+
+#include "src/tint/ir/branch.h"
+#include "src/tint/ir/value.h"
+#include "src/tint/utils/castable.h"
+
+// Forward declarations
+namespace tint::ir {
+class Loop;
+}  // namespace tint::ir
+
+namespace tint::ir {
+
+/// A break-if iteration instruction.
+class BreakIf : public utils::Castable<BreakIf, Branch> {
+  public:
+    /// Constructor
+    /// @param condition the break condition
+    /// @param loop the loop containing the break-if
+    BreakIf(Value* condition, ir::Loop* loop);
+    ~BreakIf() override;
+
+    /// @returns the break condition
+    const Value* Condition() const { return condition_; }
+
+    /// @returns the loop containing the break-if
+    const ir::Loop* Loop() const { return loop_; }
+
+  private:
+    Value* condition_ = nullptr;
+    ir::Loop* loop_ = nullptr;
+};
+
+}  // namespace tint::ir
+
+#endif  // SRC_TINT_IR_BREAK_IF_H_

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::BreakIf* Builder::BreakIf(Value* condition, Loop* loop) {
+    return ir.values.Create<ir::BreakIf>(condition, loop);
+}
+
 ir::Continue* Builder::Continue(Loop* loop) {
     return ir.values.Create<ir::Continue>(loop);
 }

src/tint/ir/builder.h

@@ -21,6 +21,7 @@
 #include "src/tint/ir/binary.h"
 #include "src/tint/ir/bitcast.h"
 #include "src/tint/ir/block_param.h"
+#include "src/tint/ir/break_if.h"
 #include "src/tint/ir/builtin.h"
 #include "src/tint/ir/constant.h"
 #include "src/tint/ir/construct.h"
@@ -337,6 +338,12 @@ class Builder {
     /// @returns the instruction
     ir::Return* Return(Function* func, utils::VectorRef<Value*> args = {});
 
+    /// Creates a loop break-if instruction
+    /// @param condition the break condition
+    /// @param loop the loop being iterated
+    /// @returns the instruction
+    ir::BreakIf* BreakIf(Value* condition, Loop* loop);
+
     /// Creates a continue instruction
     /// @param loop the loop being continued
     /// @returns the instruction

src/tint/ir/disassembler.cc

@@ -21,6 +21,7 @@
 #include "src/tint/ir/binary.h"
 #include "src/tint/ir/bitcast.h"
 #include "src/tint/ir/block.h"
+#include "src/tint/ir/break_if.h"
 #include "src/tint/ir/builtin.h"
 #include "src/tint/ir/construct.h"
 #include "src/tint/ir/continue.h"
@@ -414,6 +415,10 @@ void Disassembler::EmitBranch(const Branch* b) {
         out_ << "ret";
     } else if (auto* cont = b->As<ir::Continue>()) {
         out_ << "continue %b" << IdOf(cont->Loop()->Continuing());
+    } else if (auto* bi = b->As<ir::BreakIf>()) {
+        out_ << "break_if ";
+        EmitValue(bi->Condition());
+        out_ << " %b" << IdOf(bi->Loop()->Start());
     } else {
         out_ << "br %b" << IdOf(b->To());
         if (b->To()->Is<RootTerminator>()) {

src/tint/ir/from_program.cc

@@ -590,7 +590,9 @@ class Impl {
                     EmitBlock(stmt->continuing);
                 }
                 // Branch back to the start node if the continue target didn't branch out already
-                BranchToIfNeeded(loop_inst->Start());
+                if (NeedBranch()) {
+                    SetBranch(builder_.BreakIf(builder_.Constant(false), loop_inst));
+                }
             }
         }
 
@@ -608,7 +610,8 @@ class Impl {
         current_flow_block_->Instructions().Push(loop_inst);
 
         // Continue is always empty, just go back to the start
-        loop_inst->Continuing()->Instructions().Push(builder_.Branch(loop_inst->Start()));
+        current_flow_block_ = loop_inst->Continuing();
+        SetBranch(builder_.BreakIf(builder_.Constant(false), loop_inst));
 
         {
             ControlStackScope scope(this, loop_inst);
@@ -681,7 +684,7 @@ class Impl {
             if (stmt->continuing) {
                 current_flow_block_ = loop_inst->Continuing();
                 EmitStatement(stmt->continuing);
-                loop_inst->Continuing()->Instructions().Push(builder_.Branch(loop_inst->Start()));
+                SetBranch(builder_.BreakIf(builder_.Constant(false), loop_inst));
             }
         }
 
@@ -772,31 +775,14 @@ class Impl {
     }
 
     void EmitBreakIf(const ast::BreakIfStatement* stmt) {
+        auto* current_control = FindEnclosingControl(ControlFlags::kExcludeSwitch);
+
         // Emit the break-if condition into the end of the preceding block
-        auto reg = EmitExpression(stmt->condition);
+        auto cond = EmitExpression(stmt->condition);
-        if (!reg) {
+        if (!cond) {
             return;
         }
-        auto* if_inst = builder_.CreateIf(reg.Get());
+        SetBranch(builder_.BreakIf(cond.Get(), current_control->As<ir::Loop>()));
-        current_flow_block_->Instructions().Push(if_inst);
-
-        auto* current_control = FindEnclosingControl(ControlFlags::kExcludeSwitch);
-        TINT_ASSERT(IR, current_control);
-        TINT_ASSERT(IR, current_control->Is<Loop>());
-
-        auto* loop = current_control->As<Loop>();
-
-        current_flow_block_ = if_inst->True();
-        BranchTo(loop->Merge());
-
-        current_flow_block_ = if_inst->False();
-        BranchTo(if_inst->Merge());
-
-        current_flow_block_ = if_inst->Merge();
-
-        // The `break-if` has to be the last item in the continuing block. The false branch of
-        // the `break-if` will always take us back to the start of the loop.
-        BranchTo(loop->Start());
     }
 
     utils::Result<Value*> EmitExpression(const ast::Expression* expr) {

src/tint/ir/from_program_test.cc

@@ -385,7 +385,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinue) {
 
     EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(1u, loop_flow->Continuing()->InboundBranches().Length());
-    EXPECT_EQ(1u, loop_flow->Merge()->InboundBranches().Length());
+    EXPECT_EQ(2u, loop_flow->Merge()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->True()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->False()->InboundBranches().Length());
     EXPECT_EQ(2u, if_flow->Merge()->InboundBranches().Length());
@@ -416,7 +416,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinue) {
 
       # Continuing block
       %b3 = block {
-        br %b2
+        break_if false %b2
       }
 
     # Merge block
@@ -440,16 +440,12 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinuing_BreakIf) {
 
     auto m = res.Move();
     auto* loop_flow = FindSingleValue<ir::Loop>(m);
-    auto* break_if_flow = FindSingleValue<ir::If>(m);
 
     ASSERT_EQ(1u, m.functions.Length());
 
     EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(1u, loop_flow->Continuing()->InboundBranches().Length());
     EXPECT_EQ(1u, loop_flow->Merge()->InboundBranches().Length());
-    EXPECT_EQ(1u, break_if_flow->True()->InboundBranches().Length());
-    EXPECT_EQ(1u, break_if_flow->False()->InboundBranches().Length());
-    EXPECT_EQ(2u, break_if_flow->Merge()->InboundBranches().Length());
 
     EXPECT_EQ(Disassemble(m),
               R"(%test_function = func():void [@compute @workgroup_size(1, 1, 1)] -> %b1 {
@@ -461,23 +457,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinuing_BreakIf) {
 
       # Continuing block
       %b3 = block {
-        if true [t: %b5, f: %b6, m: %b7]
+        break_if true %b2
-          # True block
-          %b5 = block {
-            br %b4
-          }
-
-          # False block
-          %b6 = block {
-            br %b7
-          }
-
-        # Merge block
-        %b7 = block {
-          br %b2
-        }
-
-
       }
 
     # Merge block
@@ -511,23 +491,7 @@ TEST_F(IR_BuilderImplTest, Loop_Continuing_Body_Scope) {
 
       # Continuing block
       %b3 = block {
-        if true [t: %b5, f: %b6, m: %b7]
+        break_if true %b2
-          # True block
-          %b5 = block {
-            br %b4
-          }
-
-          # False block
-          %b6 = block {
-            br %b7
-          }
-
-        # Merge block
-        %b7 = block {
-          br %b2
-        }
-
-
       }
 
     # Merge block
@@ -557,7 +521,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithReturn) {
 
     EXPECT_EQ(2u, loop_flow->Start()->InboundBranches().Length());
     EXPECT_EQ(1u, loop_flow->Continuing()->InboundBranches().Length());
-    EXPECT_EQ(0u, loop_flow->Merge()->InboundBranches().Length());
+    EXPECT_EQ(1u, loop_flow->Merge()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->True()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->False()->InboundBranches().Length());
     EXPECT_EQ(2u, if_flow->Merge()->InboundBranches().Length());
@@ -565,20 +529,20 @@ TEST_F(IR_BuilderImplTest, Loop_WithReturn) {
     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, c: %b3, m: %b4]
       %b2 = block {
-        if true [t: %b4, f: %b5, m: %b6]
+        if true [t: %b5, f: %b6, m: %b7]
           # True block
-          %b4 = block {
+          %b5 = block {
             ret
           }
           # False block
-          %b5 = block {
+          %b6 = block {
-            br %b6
+            br %b7
           }
 
         # Merge block
-        %b6 = block {
+        %b7 = block {
           continue %b3
         }
 
@@ -587,9 +551,13 @@ TEST_F(IR_BuilderImplTest, Loop_WithReturn) {
 
       # Continuing block
       %b3 = block {
-        br %b2
+        break_if false %b2
       }
 
+    # Merge block
+    %b4 = block {
+      ret
+    }
 
   }
 
@@ -792,28 +760,12 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
                 # Continuing block
                 %b17 = block {
-                  if true [t: %b19, f: %b20, m: %b21]
+                  break_if true %b16
-                    # True block
-                    %b19 = block {
-                      br %b18
-                    }
-
-                    # False block
-                    %b20 = block {
-                      br %b21
-                    }
-
-                  # Merge block
-                  %b21 = block {
-                    br %b16
-                  }
-
-
                 }
 
               # Merge block
               %b18 = block {
-                br %b5
+                break_if false %b5
               }
 
 
@@ -824,19 +776,19 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
         # Merge block
         %b7 = block {
-          if true [t: %b22, f: %b23, m: %b24]
+          if true [t: %b19, f: %b20, m: %b21]
             # True block
-            %b22 = block {
+            %b19 = block {
               br %b4
             }
 
             # False block
-            %b23 = block {
+            %b20 = block {
-              br %b24
+              br %b21
             }
 
           # Merge block
-          %b24 = block {
+          %b21 = block {
             continue %b3
           }
 
@@ -848,7 +800,7 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
 
       # Continuing block
       %b3 = block {
-        br %b2
+        break_if false %b2
       }
 
     # Merge block
@@ -880,7 +832,7 @@ TEST_F(IR_BuilderImplTest, While) {
 
     EXPECT_EQ(2u, flow->Start()->InboundBranches().Length());
     EXPECT_EQ(1u, flow->Continuing()->InboundBranches().Length());
-    EXPECT_EQ(1u, flow->Merge()->InboundBranches().Length());
+    EXPECT_EQ(2u, flow->Merge()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->True()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->False()->InboundBranches().Length());
     EXPECT_EQ(2u, if_flow->Merge()->InboundBranches().Length());
@@ -911,7 +863,7 @@ TEST_F(IR_BuilderImplTest, While) {
 
       # Continuing block
       %b3 = block {
-        br %b2
+        break_if false %b2
       }
 
     # Merge block
@@ -943,7 +895,7 @@ TEST_F(IR_BuilderImplTest, While_Return) {
 
     EXPECT_EQ(2u, flow->Start()->InboundBranches().Length());
     EXPECT_EQ(0u, flow->Continuing()->InboundBranches().Length());
-    EXPECT_EQ(1u, flow->Merge()->InboundBranches().Length());
+    EXPECT_EQ(2u, flow->Merge()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->True()->InboundBranches().Length());
     EXPECT_EQ(1u, if_flow->False()->InboundBranches().Length());
     EXPECT_EQ(2u, if_flow->Merge()->InboundBranches().Length());
@@ -973,7 +925,7 @@ TEST_F(IR_BuilderImplTest, While_Return) {
 
       # Continuing block
       %b3 = block {
-        br %b2
+        break_if false %b2
       }
 
     # Merge block