[ir] Add basic block arguments.

This CL adds the ability to add parameters to blocks and emits the
params as needed in the disassembly.

The ShortCircuit is updated to use block arguments instead of creating
a return value.

Bug: tint:1909
Change-Id: I92afa6cf8ff4e01bfa3de46e76c26c465f0d6062
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/133200
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/tint/BUILD.gn

@@ -1185,6 +1185,8 @@ libtint_source_set("libtint_ir_src") {
     "ir/bitcast.h",
     "ir/block.cc",
     "ir/block.h",
+    "ir/block_param.cc",
+    "ir/block_param.h",
     "ir/builder.cc",
     "ir/builder.h",
     "ir/builtin.cc",

src/tint/CMakeLists.txt

@@ -716,6 +716,8 @@ if(${TINT_BUILD_IR})
     ir/bitcast.h
     ir/block.cc
     ir/block.h
+    ir/block_param.cc
+    ir/block_param.h
     ir/builder.cc
     ir/builder.h
     ir/builtin.cc

src/tint/ir/block.h

@@ -15,6 +15,7 @@
 #ifndef SRC_TINT_IR_BLOCK_H_
 #define SRC_TINT_IR_BLOCK_H_
 
+#include "src/tint/ir/block_param.h"
 #include "src/tint/ir/branch.h"
 #include "src/tint/ir/flow_node.h"
 #include "src/tint/ir/instruction.h"
@@ -40,6 +41,9 @@ class Block : public utils::Castable<Block, FlowNode> {
 
     /// The instructions in the block
     utils::Vector<const Instruction*, 16> instructions;
+
+    /// The parameters passed into the block
+    utils::Vector<const BlockParam*, 0> params;
 };
 
 }  // namespace tint::ir

src/tint/ir/block_param.cc

@@ -0,0 +1,25 @@
+// 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/block_param.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::ir::BlockParam);
+
+namespace tint::ir {
+
+BlockParam::BlockParam(const type::Type* ty) : type(ty) {}
+
+BlockParam::~BlockParam() = default;
+
+}  // namespace tint::ir

src/tint/ir/block_param.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_BLOCK_PARAM_H_
+#define SRC_TINT_IR_BLOCK_PARAM_H_
+
+#include "src/tint/ir/value.h"
+#include "src/tint/utils/castable.h"
+
+namespace tint::ir {
+
+/// An instruction in the IR.
+class BlockParam : public utils::Castable<BlockParam, Value> {
+  public:
+    /// Constructor
+    /// @param type the type of the var
+    explicit BlockParam(const type::Type* type);
+    BlockParam(const BlockParam& inst) = delete;
+    BlockParam(BlockParam&& inst) = delete;
+    ~BlockParam() override;
+
+    BlockParam& operator=(const BlockParam& inst) = delete;
+    BlockParam& operator=(BlockParam&& inst) = delete;
+
+    /// @returns the type of the var
+    const type::Type* Type() const override { return type; }
+
+    /// the result type of the instruction
+    const type::Type* type = nullptr;
+};
+
+}  // namespace tint::ir
+
+#endif  // SRC_TINT_IR_BLOCK_PARAM_H_

src/tint/ir/builder.cc

@@ -243,4 +243,8 @@ ir::Var* Builder::Declare(const type::Type* type) {
     return ir.instructions.Create<ir::Var>(type);
 }
 
+ir::BlockParam* Builder::BlockParam(const type::Type* type) {
+    return ir.values.Create<ir::BlockParam>(type);
+}
+
 }  // namespace tint::ir

src/tint/ir/builder.h

@@ -361,6 +361,11 @@ class Builder {
     /// @returns the instruction
     ir::Var* Declare(const type::Type* type);
 
+    /// Creates a new `BlockParam`
+    /// @param type the parameter type
+    /// @returns the value
+    ir::BlockParam* BlockParam(const type::Type* type);
+
     /// Retrieves the root block for the module, creating if necessary
     /// @returns the root block
     ir::Block* CreateRootBlockIfNeeded();

src/tint/ir/disassembler.cc

@@ -154,7 +154,19 @@ void Disassembler::Walk(const FlowNode* node) {
                 return;
             }
 
-            Indent() << "%fn" << IdOf(b) << " = block {" << std::endl;
+            Indent() << "%fn" << IdOf(b) << " = block";
+            if (!b->params.IsEmpty()) {
+                out_ << " (";
+                for (auto* p : b->params) {
+                    if (p != b->params.Front()) {
+                        out_ << ", ";
+                    }
+                    EmitValue(p);
+                }
+                out_ << ")";
+            }
+
+            out_ << " {" << std::endl;
             {
                 ScopedIndent si(indent_size_);
                 EmitBlockInstructions(b);
@@ -248,7 +260,20 @@ void Disassembler::Walk(const FlowNode* node) {
         [&](const ir::If* i) {
             Indent() << "%fn" << IdOf(i) << " = if ";
             EmitValue(i->condition);
-            out_ << " [t: %fn" << IdOf(i->true_.target) << ", f: %fn" << IdOf(i->false_.target);
+
+            bool has_true = !i->true_.target->IsDead();
+            bool has_false = !i->false_.target->IsDead();
+
+            out_ << " [";
+            if (has_true) {
+                out_ << "t: %fn" << IdOf(i->true_.target);
+            }
+            if (has_false) {
+                if (has_true) {
+                    out_ << ", ";
+                }
+                out_ << "f: %fn" << IdOf(i->false_.target);
+            }
             if (i->merge.target->IsConnected()) {
                 out_ << ", m: %fn" << IdOf(i->merge.target);
             }
@@ -258,10 +283,12 @@ void Disassembler::Walk(const FlowNode* node) {
                 ScopedIndent if_indent(indent_size_);
                 ScopedStopNode scope(stop_nodes_, i->merge.target);
 
-                Indent() << "# true branch" << std::endl;
-                Walk(i->true_.target);
+                if (has_true) {
+                    Indent() << "# true branch" << std::endl;
+                    Walk(i->true_.target);
+                }
 
-                if (!i->false_.target->IsDead()) {
+                if (has_false) {
                     Indent() << "# false branch" << std::endl;
                     Walk(i->false_.target);
                 }
@@ -373,7 +400,11 @@ void Disassembler::EmitValue(const Value* val) {
             if (i->Type() != nullptr) {
                 out_ << ":" << i->Type()->FriendlyName();
             }
-        });
+        },
+        [&](const ir::BlockParam* p) {
+            out_ << "%" << IdOf(p) << ":" << p->Type()->FriendlyName();
+        },
+        [&](Default) { out_ << "Unknown value: " << val->TypeInfo().name; });
 }
 
 void Disassembler::EmitInstruction(const Instruction* inst) {

src/tint/ir/from_program.cc

@@ -59,6 +59,7 @@
 #include "src/tint/ast/var.h"
 #include "src/tint/ast/variable_decl_statement.h"
 #include "src/tint/ast/while_statement.h"
+#include "src/tint/ir/block_param.h"
 #include "src/tint/ir/builder.h"
 #include "src/tint/ir/function.h"
 #include "src/tint/ir/if.h"
@@ -944,28 +945,34 @@ class Impl {
             return utils::Failure;
         }
 
-        // Generate a variable to store the short-circut into
-        auto* ty = builder_.ir.types.Get<type::Reference>(builder_.ir.types.Get<type::Bool>(),
-                                                          builtin::AddressSpace::kFunction,
-                                                          builtin::Access::kReadWrite);
-        auto* result_var = builder_.Declare(ty);
-        current_flow_block_->instructions.Push(result_var);
-
-        auto* lhs_store = builder_.Store(result_var, lhs.Get());
-        current_flow_block_->instructions.Push(lhs_store);
-
         auto* if_node = builder_.CreateIf(lhs.Get());
         BranchTo(if_node);
 
+        auto* result = builder_.BlockParam(builder_.ir.types.Get<type::Bool>());
+        if_node->merge.target->As<Block>()->params.Push(result);
+
         utils::Result<Value*> rhs;
         {
             FlowStackScope scope(this, if_node);
 
+            utils::Vector<Value*, 1> alt_args;
+            alt_args.Push(lhs.Get());
+
             // If this is an `&&` then we only evaluate the RHS expression in the true block.
             // If this is an `||` then we only evaluate the RHS expression in the false block.
             if (expr->op == ast::BinaryOp::kLogicalAnd) {
+                // If the lhs is false, then that is the result we want to pass to the merge block
+                // as our argument
+                current_flow_block_ = if_node->false_.target->As<Block>();
+                BranchTo(if_node->merge.target, std::move(alt_args));
+
                 current_flow_block_ = if_node->true_.target->As<Block>();
             } else {
+                // If the lhs is true, then that is the result we want to pass to the merge block
+                // as our argument
+                current_flow_block_ = if_node->true_.target->As<Block>();
+                BranchTo(if_node->merge.target, std::move(alt_args));
+
                 current_flow_block_ = if_node->false_.target->As<Block>();
             }
 
@@ -973,14 +980,14 @@ class Impl {
             if (!rhs) {
                 return utils::Failure;
             }
-            auto* rhs_store = builder_.Store(result_var, rhs.Get());
-            current_flow_block_->instructions.Push(rhs_store);
+            utils::Vector<Value*, 1> args;
+            args.Push(rhs.Get());
 
-            BranchTo(if_node->merge.target);
+            BranchTo(if_node->merge.target, std::move(args));
         }
         current_flow_block_ = if_node->merge.target->As<Block>();
 
-        return result_var;
+        return result;
     }
 
     utils::Result<Value*> EmitBinary(const ast::BinaryExpression* expr) {

src/tint/ir/from_program_binary_test.cc

@@ -442,7 +442,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_CompoundXor) {
 
 TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalAnd) {
     Func("my_func", utils::Empty, ty.bool_(), utils::Vector{Return(true)});
-    auto* expr = LogicalAnd(Call("my_func"), false);
+    auto* expr = If(LogicalAnd(Call("my_func"), false), Block());
     WrapInFunction(expr);
 
     auto m = Build();
@@ -456,18 +456,32 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalAnd) {
 %fn3 = func test_function():void [@compute @workgroup_size(1, 1, 1)] {
   %fn4 = block {
     %1:bool = call my_func
-    %tint_symbol:ref<function, bool, read_write> = var
-    store %tint_symbol:ref<function, bool, read_write>, %1:bool
   } -> %fn5 # branch
 
   %fn5 = if %1:bool [t: %fn6, f: %fn7, m: %fn8]
     # true branch
     %fn6 = block {
-      store %tint_symbol:ref<function, bool, read_write>, false
-    } -> %fn8 # branch
+    } -> %fn8 false # branch
+
+    # false branch
+    %fn7 = block {
+    } -> %fn8 %1:bool # branch
 
   # if merge
-  %fn8 = block {
+  %fn8 = block (%2:bool) {
+  } -> %fn9 # branch
+
+  %fn9 = if %2:bool [t: %fn10, f: %fn11, m: %fn12]
+    # true branch
+    %fn10 = block {
+    } -> %fn12 # branch
+
+    # false branch
+    %fn11 = block {
+    } -> %fn12 # branch
+
+  # if merge
+  %fn12 = block {
   } -> %func_end # return
 } %func_end
 
@@ -476,7 +490,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalAnd) {
 
 TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalOr) {
     Func("my_func", utils::Empty, ty.bool_(), utils::Vector{Return(true)});
-    auto* expr = LogicalOr(Call("my_func"), true);
+    auto* expr = If(LogicalOr(Call("my_func"), true), Block());
     WrapInFunction(expr);
 
     auto m = Build();
@@ -490,19 +504,32 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalOr) {
 %fn3 = func test_function():void [@compute @workgroup_size(1, 1, 1)] {
   %fn4 = block {
     %1:bool = call my_func
-    %tint_symbol:ref<function, bool, read_write> = var
-    store %tint_symbol:ref<function, bool, read_write>, %1:bool
   } -> %fn5 # branch
 
   %fn5 = if %1:bool [t: %fn6, f: %fn7, m: %fn8]
     # true branch
+    %fn6 = block {
+    } -> %fn8 %1:bool # branch
+
     # false branch
     %fn7 = block {
-      store %tint_symbol:ref<function, bool, read_write>, true
-    } -> %fn8 # branch
+    } -> %fn8 true # branch
 
   # if merge
-  %fn8 = block {
+  %fn8 = block (%2:bool) {
+  } -> %fn9 # branch
+
+  %fn9 = if %2:bool [t: %fn10, f: %fn11, m: %fn12]
+    # true branch
+    %fn10 = block {
+    } -> %fn12 # branch
+
+    # false branch
+    %fn11 = block {
+    } -> %fn12 # branch
+
+  # if merge
+  %fn12 = block {
   } -> %func_end # return
 } %func_end
 
@@ -758,23 +785,24 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_Compound) {
   %fn4 = block {
     %1:f32 = call my_func
     %2:bool = lt %1:f32, 2.0f
-    %tint_symbol:ref<function, bool, read_write> = var
-    store %tint_symbol:ref<function, bool, read_write>, %2:bool
   } -> %fn5 # branch
 
   %fn5 = if %2:bool [t: %fn6, f: %fn7, m: %fn8]
     # true branch
     %fn6 = block {
+      %3:f32 = call my_func
       %4:f32 = call my_func
-      %5:f32 = call my_func
-      %6:f32 = mul 2.29999995231628417969f, %5:f32
-      %7:f32 = div %4:f32, %6:f32
-      %8:bool = gt 2.5f, %7:f32
-      store %tint_symbol:ref<function, bool, read_write>, %8:bool
-    } -> %fn8 # branch
+      %5:f32 = mul 2.29999995231628417969f, %4:f32
+      %6:f32 = div %3:f32, %5:f32
+      %7:bool = gt 2.5f, %6:f32
+    } -> %fn8 %7:bool # branch
+
+    # false branch
+    %fn7 = block {
+    } -> %fn8 %2:bool # branch
 
   # if merge
-  %fn8 = block {
+  %fn8 = block (%tint_symbol:bool) {
   } -> %func_end # return
 } %func_end