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>

src/tint/BUILD.gn

@@ -500,6 +500,8 @@ libtint_source_set("libtint_core_all_src") {
     "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 @@ if (tint_build_unittests) {
       "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",

src/tint/CMakeLists.txt

@@ -425,6 +425,8 @@ list(APPEND TINT_LIB_SRCS
   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 @@ if(TINT_BUILD_TESTS)
       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

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

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_