tint: Add DemoteToHelper transform
This transform implements demote-to-helper semantics for discard
statements, by setting a global flag and masking any writes to
host-visible memory.
Bug: tint:1723
Change-Id: I619c5661c6ffcc6978aa4f7635b0d680d83de3dd
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/109002
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/BUILD.gn b/src/tint/BUILD.gn
index eb863de..b51e531 100644
--- a/src/tint/BUILD.gn
+++ b/src/tint/BUILD.gn
@@ -500,6 +500,8 @@
"transform/decompose_strided_array.h",
"transform/decompose_strided_matrix.cc",
"transform/decompose_strided_matrix.h",
+ "transform/demote_to_helper.cc",
+ "transform/demote_to_helper.h",
"transform/disable_uniformity_analysis.cc",
"transform/disable_uniformity_analysis.h",
"transform/expand_compound_assignment.cc",
@@ -1223,6 +1225,7 @@
"transform/decompose_memory_access_test.cc",
"transform/decompose_strided_array_test.cc",
"transform/decompose_strided_matrix_test.cc",
+ "transform/demote_to_helper_test.cc",
"transform/disable_uniformity_analysis_test.cc",
"transform/expand_compound_assignment_test.cc",
"transform/first_index_offset_test.cc",
diff --git a/src/tint/CMakeLists.txt b/src/tint/CMakeLists.txt
index 69712dd..a42d58e 100644
--- a/src/tint/CMakeLists.txt
+++ b/src/tint/CMakeLists.txt
@@ -425,6 +425,8 @@
transform/decompose_strided_array.h
transform/decompose_strided_matrix.cc
transform/decompose_strided_matrix.h
+ transform/demote_to_helper.cc
+ transform/demote_to_helper.h
transform/disable_uniformity_analysis.cc
transform/disable_uniformity_analysis.h
transform/expand_compound_assignment.cc
@@ -1183,6 +1185,7 @@
transform/decompose_memory_access_test.cc
transform/decompose_strided_array_test.cc
transform/decompose_strided_matrix_test.cc
+ transform/demote_to_helper_test.cc
transform/disable_uniformity_analysis_test.cc
transform/expand_compound_assignment_test.cc
transform/first_index_offset_test.cc
diff --git a/src/tint/transform/demote_to_helper.cc b/src/tint/transform/demote_to_helper.cc
new file mode 100644
index 0000000..2516a76
--- /dev/null
+++ b/src/tint/transform/demote_to_helper.cc
@@ -0,0 +1,200 @@
+// Copyright 2022 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/transform/demote_to_helper.h"
+
+#include <unordered_set>
+#include <utility>
+
+#include "src/tint/program_builder.h"
+#include "src/tint/sem/block_statement.h"
+#include "src/tint/sem/call.h"
+#include "src/tint/sem/function.h"
+#include "src/tint/sem/reference.h"
+#include "src/tint/sem/statement.h"
+#include "src/tint/transform/utils/hoist_to_decl_before.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::transform::DemoteToHelper);
+
+using namespace tint::number_suffixes; // NOLINT
+
+namespace tint::transform {
+
+DemoteToHelper::DemoteToHelper() = default;
+
+DemoteToHelper::~DemoteToHelper() = default;
+
+Transform::ApplyResult DemoteToHelper::Apply(const Program* src, const DataMap&, DataMap&) const {
+ auto& sem = src->Sem();
+
+ // Collect the set of functions that need to be processed.
+ // A function needs to be processed if it is reachable by a shader that contains a discard at
+ // any point in its call hierarchy.
+ std::unordered_set<const sem::Function*> functions_to_process;
+ for (auto* func : src->AST().Functions()) {
+ if (!func->IsEntryPoint()) {
+ continue;
+ }
+
+ // Determine whether this entry point and its callees need to be transformed.
+ bool needs_transform = false;
+ if (sem.Get(func)->DiscardStatement()) {
+ needs_transform = true;
+ } else {
+ for (auto* callee : sem.Get(func)->TransitivelyCalledFunctions()) {
+ if (callee->DiscardStatement()) {
+ needs_transform = true;
+ break;
+ }
+ }
+ }
+ if (!needs_transform) {
+ continue;
+ }
+
+ // Process the entry point and its callees.
+ functions_to_process.insert(sem.Get(func));
+ for (auto* callee : sem.Get(func)->TransitivelyCalledFunctions()) {
+ functions_to_process.insert(callee);
+ }
+ }
+
+ if (functions_to_process.empty()) {
+ return SkipTransform;
+ }
+
+ ProgramBuilder b;
+ CloneContext ctx{&b, src, /* auto_clone_symbols */ true};
+
+ // Create a module-scope flag that indicates whether the current invocation has been discarded.
+ auto flag = b.Symbols().New("tint_discarded");
+ b.GlobalVar(flag, ast::AddressSpace::kPrivate, b.Expr(false));
+
+ // Replace all discard statements with a statement that marks the invocation as discarded.
+ ctx.ReplaceAll([&](const ast::DiscardStatement*) -> const ast::Statement* {
+ return b.Assign(flag, b.Expr(true));
+ });
+
+ // Insert a conditional discard at the end of each entry point that does not end with a return.
+ for (auto* func : functions_to_process) {
+ if (func->Declaration()->IsEntryPoint()) {
+ auto* sem_body = sem.Get(func->Declaration()->body);
+ if (sem_body->Behaviors().Contains(sem::Behavior::kNext)) {
+ ctx.InsertBack(func->Declaration()->body->statements,
+ b.If(flag, b.Block(b.Discard())));
+ }
+ }
+ }
+
+ HoistToDeclBefore hoist_to_decl_before(ctx);
+
+ // Mask all writes to host-visible memory using the discarded flag.
+ // We also insert a discard statement before all return statements in entry points for shaders
+ // that discard.
+ for (auto* node : src->ASTNodes().Objects()) {
+ Switch(
+ node,
+
+ // Mask assignments to storage buffer variables.
+ [&](const ast::AssignmentStatement* assign) {
+ // Skip writes in functions that are not called from shaders that discard.
+ auto* func = sem.Get(assign)->Function();
+ if (functions_to_process.count(func) == 0) {
+ return;
+ }
+
+ // Skip writes to invocation-private address spaces.
+ auto* ref = sem.Get(assign->lhs)->Type()->As<sem::Reference>();
+ switch (ref->AddressSpace()) {
+ case ast::AddressSpace::kStorage:
+ // Need to mask these.
+ break;
+ case ast::AddressSpace::kFunction:
+ case ast::AddressSpace::kPrivate:
+ case ast::AddressSpace::kOut:
+ // Skip these.
+ return;
+ default:
+ TINT_UNREACHABLE(Transform, b.Diagnostics())
+ << "write to unhandled address space: " << ref->AddressSpace();
+ }
+
+ // Mask the assignment using the invocation-discarded flag.
+ ctx.Replace(assign, b.If(b.Not(flag), b.Block(ctx.Clone(assign))));
+ },
+
+ // Mask builtins that write to host-visible memory.
+ [&](const ast::CallExpression* call) {
+ auto* sem_call = sem.Get<sem::Call>(call);
+ auto* stmt = sem_call ? sem_call->Stmt() : nullptr;
+ auto* func = stmt ? stmt->Function() : nullptr;
+ auto* builtin = sem_call ? sem_call->Target()->As<sem::Builtin>() : nullptr;
+ if (functions_to_process.count(func) == 0 || !builtin) {
+ return;
+ }
+
+ if (builtin->Type() == sem::BuiltinType::kTextureStore) {
+ // A call to textureStore() will always be a statement.
+ // Wrap it inside a conditional block.
+ auto* masked_call = b.If(b.Not(flag), b.Block(ctx.Clone(stmt->Declaration())));
+ ctx.Replace(stmt->Declaration(), masked_call);
+ } else if (builtin->IsAtomic() &&
+ builtin->Type() != sem::BuiltinType::kAtomicLoad) {
+ // A call to an atomic builtin can be a statement or an expression.
+ if (auto* call_stmt = stmt->Declaration()->As<ast::CallStatement>();
+ call_stmt && call_stmt->expr == call) {
+ // This call is a statement.
+ // Wrap it inside a conditional block.
+ auto* masked_call = b.If(b.Not(flag), b.Block(ctx.Clone(call_stmt)));
+ ctx.Replace(stmt->Declaration(), masked_call);
+ } else {
+ // This call is an expression.
+ // We transform:
+ // let y = x + atomicAdd(&p, 1);
+ // Into:
+ // var tmp : i32;
+ // if (!tint_discarded) {
+ // tmp = atomicAdd(&p, 1);
+ // }
+ // let y = x + tmp;
+ auto result = b.Sym();
+ auto result_decl =
+ b.Decl(b.Var(result, CreateASTTypeFor(ctx, sem_call->Type())));
+ auto* masked_call =
+ b.If(b.Not(flag),
+ b.Block(b.Assign(result, ctx.CloneWithoutTransform(call))));
+ hoist_to_decl_before.Prepare(sem_call);
+ hoist_to_decl_before.InsertBefore(stmt, result_decl);
+ hoist_to_decl_before.InsertBefore(stmt, masked_call);
+ ctx.Replace(call, b.Expr(result));
+ }
+ }
+ },
+
+ // Insert a conditional discard before all return statements in entry points.
+ [&](const ast::ReturnStatement* ret) {
+ auto* func = sem.Get(ret)->Function();
+ if (func->Declaration()->IsEntryPoint() && functions_to_process.count(func)) {
+ auto* discard = b.If(flag, b.Block(b.Discard()));
+ ctx.InsertBefore(sem.Get(ret)->Block()->Declaration()->statements, ret,
+ discard);
+ }
+ });
+ }
+
+ ctx.Clone();
+ return Program(std::move(b));
+}
+
+} // namespace tint::transform
diff --git a/src/tint/transform/demote_to_helper.h b/src/tint/transform/demote_to_helper.h
new file mode 100644
index 0000000..8968e78
--- /dev/null
+++ b/src/tint/transform/demote_to_helper.h
@@ -0,0 +1,43 @@
+// Copyright 2022 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_TRANSFORM_DEMOTE_TO_HELPER_H_
+#define SRC_TINT_TRANSFORM_DEMOTE_TO_HELPER_H_
+
+#include "src/tint/transform/transform.h"
+
+namespace tint::transform {
+
+/// Implement demote-to-helper semantics for discard statements.
+///
+/// For backend targets that implement discard by terminating the invocation, we need to change the
+/// program to ensure that discarding the fragment does not affect uniformity with respect to
+/// derivative operations. We do this by setting a global flag and masking all writes to storage
+/// buffers and textures.
+class DemoteToHelper final : public Castable<DemoteToHelper, Transform> {
+ public:
+ /// Constructor
+ DemoteToHelper();
+ /// Destructor
+ ~DemoteToHelper() override;
+
+ /// @copydoc Transform::Apply
+ ApplyResult Apply(const Program* program,
+ const DataMap& inputs,
+ DataMap& outputs) const override;
+};
+
+} // namespace tint::transform
+
+#endif // SRC_TINT_TRANSFORM_DEMOTE_TO_HELPER_H_
diff --git a/src/tint/transform/demote_to_helper_test.cc b/src/tint/transform/demote_to_helper_test.cc
new file mode 100644
index 0000000..e493410
--- /dev/null
+++ b/src/tint/transform/demote_to_helper_test.cc
@@ -0,0 +1,1100 @@
+// Copyright 2022 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/transform/demote_to_helper.h"
+
+#include <utility>
+
+#include "src/tint/transform/test_helper.h"
+
+namespace tint::transform {
+namespace {
+
+using DemoteToHelperTest = TransformTest;
+
+TEST_F(DemoteToHelperTest, ShouldRunEmptyModule) {
+ auto* src = R"()";
+
+ EXPECT_FALSE(ShouldRun<DemoteToHelper>(src));
+}
+
+TEST_F(DemoteToHelperTest, ShouldRunNoDiscard) {
+ auto* src = R"(
+@group(0) @binding(0)
+var<storage, read_write> v : f32;
+
+@fragment
+fn foo() {
+ v = 42;
+}
+)";
+
+ EXPECT_FALSE(ShouldRun<DemoteToHelper>(src));
+}
+
+TEST_F(DemoteToHelperTest, ShouldRunDiscardInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0)
+var<storage, read_write> v : f32;
+
+@fragment
+fn foo() {
+ discard;
+ v = 42;
+}
+)";
+
+ EXPECT_TRUE(ShouldRun<DemoteToHelper>(src));
+}
+
+TEST_F(DemoteToHelperTest, ShouldRunDiscardInHelper) {
+ auto* src = R"(
+@group(0) @binding(0)
+var<storage, read_write> v : f32;
+
+fn bar() {
+ discard;
+}
+
+@fragment
+fn foo() {
+ bar();
+ v = 42;
+}
+)";
+
+ EXPECT_TRUE(ShouldRun<DemoteToHelper>(src));
+}
+
+TEST_F(DemoteToHelperTest, EmptyModule) {
+ auto* src = R"()";
+
+ auto* expect = src;
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, WriteInEntryPoint_DiscardInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ v = ret.x;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ v = ret.x;
+ }
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, WriteInEntryPoint_DiscardInHelper) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar() {
+ discard;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ bar();
+ }
+ let ret = textureSample(t, s, coord);
+ v = ret.x;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar() {
+ tint_discarded = true;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ bar();
+ }
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ v = ret.x;
+ }
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, WriteInHelper_DiscardInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar(coord : vec2<f32>) {
+ let ret = textureSample(t, s, coord);
+ v = ret.x;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ bar(coord);
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar(coord : vec2<f32>) {
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ v = ret.x;
+ }
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ bar(coord);
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, WriteInHelper_DiscardInHelper) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar(in : f32, coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ v = ret.x;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ bar(in, coord);
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar(in : f32, coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ v = ret.x;
+ }
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ bar(in, coord);
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, WriteInEntryPoint_NoDiscard) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ let ret = textureSample(t, s, coord);
+ v = ret.x;
+}
+)";
+
+ auto* expect = src;
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that no additional discards are inserted when the function unconditionally returns in a
+// nested block.
+TEST_F(DemoteToHelperTest, EntryPointReturn_NestedInBlock) {
+ auto* src = R"(
+@fragment
+fn foo() {
+ {
+ discard;
+ return;
+ }
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@fragment
+fn foo() {
+ {
+ tint_discarded = true;
+ if (tint_discarded) {
+ discard;
+ }
+ return;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that a discard statement is inserted before every return statement in an entry point that
+// contains a discard.
+TEST_F(DemoteToHelperTest, EntryPointReturns_DiscardInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) f32 {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ if (in < 1.0) {
+ return ret.x;
+ }
+ return 2.0;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) f32 {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ if ((in < 1.0)) {
+ if (tint_discarded) {
+ discard;
+ }
+ return ret.x;
+ }
+ if (tint_discarded) {
+ discard;
+ }
+ return 2.0;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that a discard statement is inserted before every return statement in an entry point that
+// calls a function that contains a discard.
+TEST_F(DemoteToHelperTest, EntryPointReturns_DiscardInHelper) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar() {
+ discard;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) f32 {
+ if (in == 0.0) {
+ bar();
+ }
+ let ret = textureSample(t, s, coord);
+ if (in < 1.0) {
+ return ret.x;
+ }
+ return 2.0;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar() {
+ tint_discarded = true;
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) f32 {
+ if ((in == 0.0)) {
+ bar();
+ }
+ let ret = textureSample(t, s, coord);
+ if ((in < 1.0)) {
+ if (tint_discarded) {
+ discard;
+ }
+ return ret.x;
+ }
+ if (tint_discarded) {
+ discard;
+ }
+ return 2.0;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that no return statements are modified in an entry point that does not discard.
+TEST_F(DemoteToHelperTest, EntryPointReturns_NoDiscard) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+fn bar() {
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) f32 {
+ if ((in == 0.0)) {
+ bar();
+ }
+ let ret = textureSample(t, s, coord);
+ if ((in < 1.0)) {
+ return ret.x;
+ }
+ return 2.0;
+}
+)";
+
+ auto* expect = src;
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that only functions that are part of a shader that discards are transformed.
+// Functions in non-discarding stages should not have their writes masked, and non-discarding entry
+// points should not have their return statements replaced.
+TEST_F(DemoteToHelperTest, MultipleShaders) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v1 : f32;
+
+@group(0) @binding(3) var<storage, read_write> v2 : f32;
+
+fn bar_discard(in : f32, coord : vec2<f32>) -> f32 {
+ let ret = textureSample(t, s, coord);
+ v1 = ret.x * 2.0;
+ return ret.y * 2.0;
+}
+
+fn bar_no_discard(in : f32, coord : vec2<f32>) -> f32 {
+ let ret = textureSample(t, s, coord);
+ v1 = ret.x * 2.0;
+ return ret.y * 2.0;
+}
+
+@fragment
+fn foo_discard(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = bar_discard(in, coord);
+ v2 = ret;
+}
+
+@fragment
+fn foo_no_discard(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ return;
+ }
+ let ret = bar_no_discard(in, coord);
+ v2 = ret;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v1 : f32;
+
+@group(0) @binding(3) var<storage, read_write> v2 : f32;
+
+fn bar_discard(in : f32, coord : vec2<f32>) -> f32 {
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ v1 = (ret.x * 2.0);
+ }
+ return (ret.y * 2.0);
+}
+
+fn bar_no_discard(in : f32, coord : vec2<f32>) -> f32 {
+ let ret = textureSample(t, s, coord);
+ v1 = (ret.x * 2.0);
+ return (ret.y * 2.0);
+}
+
+@fragment
+fn foo_discard(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = bar_discard(in, coord);
+ if (!(tint_discarded)) {
+ v2 = ret;
+ }
+ if (tint_discarded) {
+ discard;
+ }
+}
+
+@fragment
+fn foo_no_discard(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ return;
+ }
+ let ret = bar_no_discard(in, coord);
+ v2 = ret;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that we do not mask writes to invocation-private address spaces.
+TEST_F(DemoteToHelperTest, InvocationPrivateWrites) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+var<private> vp : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ var vf : f32;
+ vf = ret.x;
+ vp = ret.y;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+var<private> vp : f32;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ var vf : f32;
+ vf = ret.x;
+ vp = ret.y;
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, TextureStoreInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var t2 : texture_storage_2d<rgba8unorm, write>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ textureStore(t2, vec2<u32>(coord), ret);
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var t2 : texture_storage_2d<rgba8unorm, write>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ textureStore(t2, vec2<u32>(coord), ret);
+ }
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, TextureStoreInHelper) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var t2 : texture_storage_2d<rgba8unorm, write>;
+
+fn bar(coord : vec2<f32>, value : vec4<f32>) {
+ textureStore(t2, vec2<u32>(coord), value);
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ bar(coord, ret);
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var t2 : texture_storage_2d<rgba8unorm, write>;
+
+fn bar(coord : vec2<f32>, value : vec4<f32>) {
+ if (!(tint_discarded)) {
+ textureStore(t2, vec2<u32>(coord), value);
+ }
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ bar(coord, ret);
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, TextureStore_NoDiscard) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var t2 : texture_storage_2d<rgba8unorm, write>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ let ret = textureSample(t, s, coord);
+ textureStore(t2, vec2<u32>(coord), ret);
+}
+)";
+
+ auto* expect = src;
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, AtomicStoreInEntryPoint) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ atomicStore(&a, i32(ret.x));
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ if (!(tint_discarded)) {
+ atomicStore(&(a), i32(ret.x));
+ }
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, AtomicStoreInHelper) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+fn bar(value : vec4<f32>) {
+ atomicStore(&a, i32(value.x));
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if (in == 0.0) {
+ discard;
+ }
+ let ret = textureSample(t, s, coord);
+ bar(ret);
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+fn bar(value : vec4<f32>) {
+ if (!(tint_discarded)) {
+ atomicStore(&(a), i32(value.x));
+ }
+}
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let ret = textureSample(t, s, coord);
+ bar(ret);
+ if (tint_discarded) {
+ discard;
+ }
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, AtomicStore_NoDiscard) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) {
+ let ret = textureSample(t, s, coord);
+ atomicStore(&(a), i32(ret.x));
+}
+)";
+
+ auto* expect = src;
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, AtomicBuiltinExpression) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if (in == 0.0) {
+ discard;
+ }
+ let v = i32(textureSample(t, s, coord).x);
+ let result = v + atomicAdd(&a, v);
+ return result;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let v = i32(textureSample(t, s, coord).x);
+ var tint_symbol : i32;
+ if (!(tint_discarded)) {
+ tint_symbol = atomicAdd(&(a), v);
+ }
+ let result = (v + tint_symbol);
+ if (tint_discarded) {
+ discard;
+ }
+ return result;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DemoteToHelperTest, AtomicBuiltinExpression_InForLoopContinuing) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if (in == 0.0) {
+ discard;
+ }
+ var result = 0;
+ for (var i = 0; i < 10; i = atomicAdd(&a, 1)) {
+ result += i32(textureSample(t, s, coord).x);
+ }
+ return result;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ var result = 0;
+ {
+ var i = 0;
+ loop {
+ if (!((i < 10))) {
+ break;
+ }
+ {
+ result += i32(textureSample(t, s, coord).x);
+ }
+
+ continuing {
+ var tint_symbol : i32;
+ if (!(tint_discarded)) {
+ tint_symbol = atomicAdd(&(a), 1);
+ }
+ i = tint_symbol;
+ }
+ }
+ }
+ if (tint_discarded) {
+ discard;
+ }
+ return result;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+// Test that no masking is generated for calls to `atomicLoad()`.
+TEST_F(DemoteToHelperTest, AtomicLoad) {
+ auto* src = R"(
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if (in == 0.0) {
+ discard;
+ }
+ let v = i32(textureSample(t, s, coord).x);
+ let result = v + atomicLoad(&a);
+ return result;
+}
+)";
+
+ auto* expect = R"(
+var<private> tint_discarded = false;
+
+@group(0) @binding(0) var t : texture_2d<f32>;
+
+@group(0) @binding(1) var s : sampler;
+
+@group(0) @binding(2) var<storage, read_write> v : f32;
+
+@group(0) @binding(3) var<storage, read_write> a : atomic<i32>;
+
+@fragment
+fn foo(@location(0) in : f32, @location(1) coord : vec2<f32>) -> @location(0) i32 {
+ if ((in == 0.0)) {
+ tint_discarded = true;
+ }
+ let v = i32(textureSample(t, s, coord).x);
+ let result = (v + atomicLoad(&(a)));
+ if (tint_discarded) {
+ discard;
+ }
+ return result;
+}
+)";
+
+ auto got = Run<DemoteToHelper>(src);
+
+ EXPECT_EQ(expect, str(got));
+}
+
+} // namespace
+} // namespace tint::transform
