[ir] Add ir::Discard

This CL adds support for discard into the IR. The `discard` statement is
handled as an instruction in the current block. The `discard` is a
`demote_to_helper` in WGSL so control flow has to continue after the
discard, it just predicates writes. So, an instruction seems like the
most logical way to express.

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

src/tint/CMakeLists.txt

@@ -720,6 +720,8 @@ if(${TINT_BUILD_IR})
     ir/debug.h
     ir/disassembler.cc
     ir/disassembler.h
+    ir/discard.cc
+    ir/discard.h
     ir/flow_node.cc
     ir/flow_node.h
     ir/function.cc
@@ -1413,6 +1415,7 @@ if(TINT_BUILD_TESTS)
       ir/bitcast_test.cc
       ir/builder_impl_test.cc
       ir/constant_test.cc
+      ir/discard_test.cc
       ir/runtime_test.cc
       ir/test_helper.h
     )

src/tint/ir/binary.h

@@ -16,8 +16,6 @@
 #define SRC_TINT_IR_BINARY_H_
 
 #include "src/tint/ir/instruction.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/bitcast.h

@@ -16,8 +16,6 @@
 #define SRC_TINT_IR_BITCAST_H_
 
 #include "src/tint/ir/instruction.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/builder.cc

@@ -177,6 +177,10 @@ ir::Bitcast* Builder::Bitcast(const type::Type* type, Value* val) {
     return ir.instructions.Create<ir::Bitcast>(Runtime(type), val);
 }
 
+ir::Discard* Builder::Discard() {
+    return ir.instructions.Create<ir::Discard>(Runtime(ir.types.Get<type::Void>()));
+}
+
 ir::UserCall* Builder::UserCall(const type::Type* type,
                                 Symbol name,
                                 utils::VectorRef<Value*> args) {

src/tint/ir/builder.h

@@ -24,6 +24,7 @@
 #include "src/tint/ir/constant.h"
 #include "src/tint/ir/construct.h"
 #include "src/tint/ir/convert.h"
+#include "src/tint/ir/discard.h"
 #include "src/tint/ir/function.h"
 #include "src/tint/ir/if.h"
 #include "src/tint/ir/loop.h"
@@ -38,6 +39,7 @@
 #include "src/tint/type/f32.h"
 #include "src/tint/type/i32.h"
 #include "src/tint/type/u32.h"
+#include "src/tint/type/void.h"
 
 namespace tint::ir {
 
@@ -284,6 +286,10 @@ class Builder {
     /// @returns the instruction
     ir::Bitcast* Bitcast(const type::Type* type, Value* val);
 
+    /// Creates a discard instruction
+    /// @returns the instruction
+    ir::Discard* Discard();
+
     /// Creates a user function call instruction
     /// @param type the return type of the call
     /// @param name the name of the function being called

src/tint/ir/builder_impl.cc

@@ -27,6 +27,7 @@
 #include "src/tint/ast/call_statement.h"
 #include "src/tint/ast/const_assert.h"
 #include "src/tint/ast/continue_statement.h"
+#include "src/tint/ast/discard_statement.h"
 #include "src/tint/ast/float_literal_expression.h"
 #include "src/tint/ast/for_loop_statement.h"
 #include "src/tint/ast/function.h"
@@ -256,9 +257,7 @@ bool BuilderImpl::EmitStatement(const ast::Statement* stmt) {
         // TODO(dsinclair): Implement
         // },
         [&](const ast::ContinueStatement* c) { return EmitContinue(c); },
-        // [&](const ast::DiscardStatement* d) {
-        // TODO(dsinclair): Implement
-        // },
+        [&](const ast::DiscardStatement* d) { return EmitDiscard(d); },
         [&](const ast::IfStatement* i) { return EmitIf(i); },
         [&](const ast::LoopStatement* l) { return EmitLoop(l); },
         [&](const ast::ForLoopStatement* l) { return EmitForLoop(l); },
@@ -561,6 +560,16 @@ bool BuilderImpl::EmitContinue(const ast::ContinueStatement*) {
     return true;
 }
 
+// Discard is being treated as an instruction. The semantics in WGSL is demote_to_helper, so the
+// code has to continue as before it just predicates writes. If WGSL grows some kind of terminating
+// discard that would probably make sense as a FlowNode but would then require figuring out the
+// multi-level exit that is triggered.
+bool BuilderImpl::EmitDiscard(const ast::DiscardStatement*) {
+    auto* instr = builder.Discard();
+    current_flow_block->instructions.Push(instr);
+    return true;
+}
+
 bool BuilderImpl::EmitBreakIf(const ast::BreakIfStatement* stmt) {
     auto* if_node = builder.CreateIf();
 

src/tint/ir/builder_impl.h

@@ -42,6 +42,7 @@ class BreakStatement;
 class CallExpression;
 class CallStatement;
 class ContinueStatement;
+class DiscardStatement;
 class Expression;
 class ForLoopStatement;
 class Function;
@@ -145,6 +146,10 @@ class BuilderImpl {
     /// @returns true if successful, false otherwise.
     bool EmitContinue(const ast::ContinueStatement* stmt);
 
+    /// Emits a discard statement
+    /// @returns true if successful, false otherwise.
+    bool EmitDiscard(const ast::DiscardStatement*);
+
     /// Emits a break-if statement
     /// @param stmt the break-if statement
     /// @returns true if successful, false otherwise.

src/tint/ir/builder_impl_test.cc

@@ -1858,6 +1858,24 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Bitcast) {
 )");
 }
 
+TEST_F(IR_BuilderImplTest, EmitStatement_Discard) {
+    auto* expr = Discard();
+    Func("test_function", {}, ty.void_(), expr,
+         utils::Vector{
+             create<ast::StageAttribute>(ast::PipelineStage::kFragment),
+         });
+
+    auto& b = CreateBuilder();
+    InjectFlowBlock();
+    auto r = b.EmitStatement(expr);
+    ASSERT_TRUE(r) << b.error();
+
+    Disassembler d(b.builder.ir);
+    d.EmitBlockInstructions(b.current_flow_block->As<ir::Block>());
+    EXPECT_EQ(d.AsString(), R"(%1 (void) = discard
+)");
+}
+
 TEST_F(IR_BuilderImplTest, EmitStatement_UserFunction) {
     Func("my_func", utils::Vector{Param("p", ty.f32())}, ty.void_(), utils::Empty);
 

src/tint/ir/builtin.h

@@ -17,8 +17,6 @@
 
 #include "src/tint/builtin/function.h"
 #include "src/tint/ir/call.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/call.h

@@ -16,8 +16,6 @@
 #define SRC_TINT_IR_CALL_H_
 
 #include "src/tint/ir/instruction.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/constant.h

@@ -17,7 +17,6 @@
 
 #include "src/tint/constant/value.h"
 #include "src/tint/ir/value.h"
-#include "src/tint/symbol_table.h"
 #include "src/tint/utils/string_stream.h"
 
 namespace tint::ir {

src/tint/ir/construct.h

@@ -16,8 +16,6 @@
 #define SRC_TINT_IR_CONSTRUCT_H_
 
 #include "src/tint/ir/call.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/convert.h

@@ -16,7 +16,6 @@
 #define SRC_TINT_IR_CONVERT_H_
 
 #include "src/tint/ir/call.h"
-#include "src/tint/symbol_table.h"
 #include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"

src/tint/ir/discard.cc

@@ -0,0 +1,32 @@
+// 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/discard.h"
+#include "src/tint/debug.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::ir::Discard);
+
+namespace tint::ir {
+
+Discard::Discard(Value* result) : Base(result) {}
+
+Discard::~Discard() = default;
+
+utils::StringStream& Discard::ToString(utils::StringStream& out) const {
+    Result()->ToString(out);
+    out << " = discard";
+    return out;
+}
+
+}  // namespace tint::ir

src/tint/ir/discard.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_DISCARD_H_
+#define SRC_TINT_IR_DISCARD_H_
+
+#include "src/tint/ir/instruction.h"
+#include "src/tint/utils/castable.h"
+#include "src/tint/utils/string_stream.h"
+
+namespace tint::ir {
+
+/// A discard instruction in the IR.
+class Discard : public utils::Castable<Discard, Instruction> {
+  public:
+    /// Constructor
+    /// @param result the result id
+    explicit Discard(Value* result);
+    Discard(const Discard& instr) = delete;
+    Discard(Discard&& instr) = delete;
+    ~Discard() override;
+
+    Discard& operator=(const Discard& instr) = delete;
+    Discard& operator=(Discard&& instr) = delete;
+
+    /// Write the instruction to the given stream
+    /// @param out the stream to write to
+    /// @returns the stream
+    utils::StringStream& ToString(utils::StringStream& out) const override;
+};
+
+}  // namespace tint::ir
+
+#endif  // SRC_TINT_IR_DISCARD_H_

src/tint/ir/discard_test.cc

@@ -0,0 +1,41 @@
+// 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/instruction.h"
+#include "src/tint/ir/test_helper.h"
+#include "src/tint/utils/string_stream.h"
+
+namespace tint::ir {
+namespace {
+
+using IR_InstructionTest = TestHelper;
+
+TEST_F(IR_InstructionTest, Discard) {
+    auto& b = CreateEmptyBuilder();
+
+    b.builder.next_runtime_id = Runtime::Id(42);
+    const auto* instr = b.builder.Discard();
+
+    ASSERT_TRUE(instr->Result()->Is<Runtime>());
+    EXPECT_EQ(Runtime::Id(42), instr->Result()->As<Runtime>()->AsId());
+    ASSERT_NE(instr->Result()->Type(), nullptr);
+    ASSERT_NE(instr->Result()->Type()->As<type::Void>(), nullptr);
+
+    utils::StringStream str;
+    instr->ToString(str);
+    EXPECT_EQ(str.str(), "%42 (void) = discard");
+}
+
+}  // namespace
+}  // namespace tint::ir

src/tint/ir/instruction.h

@@ -16,7 +16,6 @@
 #define SRC_TINT_IR_INSTRUCTION_H_
 
 #include "src/tint/ir/value.h"
-#include "src/tint/symbol_table.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/runtime.h

@@ -16,7 +16,6 @@
 #define SRC_TINT_IR_RUNTIME_H_
 
 #include "src/tint/ir/value.h"
-#include "src/tint/symbol_table.h"
 #include "src/tint/utils/string_stream.h"
 
 namespace tint::ir {

src/tint/ir/user_call.h

@@ -16,8 +16,7 @@
 #define SRC_TINT_IR_USER_CALL_H_
 
 #include "src/tint/ir/call.h"
-#include "src/tint/symbol_table.h"
-#include "src/tint/type/type.h"
+#include "src/tint/symbol.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"
 

src/tint/ir/value.h

@@ -15,7 +15,6 @@
 #ifndef SRC_TINT_IR_VALUE_H_
 #define SRC_TINT_IR_VALUE_H_
 
-#include "src/tint/symbol_table.h"
 #include "src/tint/type/type.h"
 #include "src/tint/utils/castable.h"
 #include "src/tint/utils/string_stream.h"