blob: 0964a81bfccc978b693df2710692149e6a7ef5f6 [file] [log] [blame]
// 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/lang/wgsl/ast/transform/clamp_frag_depth.h"
#include <utility>
#include "src/tint/lang/core/builtin/builtin_value.h"
#include "src/tint/lang/wgsl/ast/attribute.h"
#include "src/tint/lang/wgsl/ast/builtin_attribute.h"
#include "src/tint/lang/wgsl/ast/function.h"
#include "src/tint/lang/wgsl/ast/module.h"
#include "src/tint/lang/wgsl/ast/struct.h"
#include "src/tint/lang/wgsl/program/clone_context.h"
#include "src/tint/lang/wgsl/program/program_builder.h"
#include "src/tint/lang/wgsl/sem/function.h"
#include "src/tint/lang/wgsl/sem/statement.h"
#include "src/tint/lang/wgsl/sem/struct.h"
#include "src/tint/utils/containers/vector.h"
#include "src/tint/utils/macros/scoped_assignment.h"
TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::ClampFragDepth);
namespace tint::ast::transform {
/// PIMPL state for the transform
struct ClampFragDepth::State {
/// The source program
const Program* const src;
/// The target program builder
ProgramBuilder b{};
/// The clone context
program::CloneContext ctx = {&b, src, /* auto_clone_symbols */ true};
/// The sem::Info of the program
const sem::Info& sem = src->Sem();
/// The symbols of the program
const SymbolTable& sym = src->Symbols();
/// Runs the transform
/// @returns the new program or SkipTransform if the transform is not required
Transform::ApplyResult Run() {
// Abort on any use of push constants in the module.
for (auto* global : src->AST().GlobalVariables()) {
if (auto* var = global->As<Var>()) {
auto* v = src->Sem().Get(var);
if (TINT_UNLIKELY(v->AddressSpace() == builtin::AddressSpace::kPushConstant)) {
TINT_ICE()
<< "ClampFragDepth doesn't know how to handle module that already use push "
"constants";
return Program(std::move(b));
}
}
}
if (!ShouldRun()) {
return SkipTransform;
}
// At least one entry-point needs clamping. Add the following to the module:
//
// enable chromium_experimental_push_constant;
//
// struct FragDepthClampArgs {
// min : f32,
// max : f32,
// }
// var<push_constant> frag_depth_clamp_args : FragDepthClampArgs;
//
// fn clamp_frag_depth(v : f32) -> f32 {
// return clamp(v, frag_depth_clamp_args.min, frag_depth_clamp_args.max);
// }
b.Enable(builtin::Extension::kChromiumExperimentalPushConstant);
b.Structure(b.Symbols().New("FragDepthClampArgs"),
tint::Vector{b.Member("min", b.ty.f32()), b.Member("max", b.ty.f32())});
auto args_sym = b.Symbols().New("frag_depth_clamp_args");
b.GlobalVar(args_sym, b.ty("FragDepthClampArgs"), builtin::AddressSpace::kPushConstant);
auto base_fn_sym = b.Symbols().New("clamp_frag_depth");
b.Func(base_fn_sym, tint::Vector{b.Param("v", b.ty.f32())}, b.ty.f32(),
tint::Vector{b.Return(b.Call("clamp", "v", b.MemberAccessor(args_sym, "min"),
b.MemberAccessor(args_sym, "max")))});
// If true, the currently cloned function returns frag depth directly as a scalar
bool returns_frag_depth_as_value = false;
// If valid, the currently cloned function returns frag depth in a struct
// The symbol is the name of the helper function to apply the depth clamping.
Symbol returns_frag_depth_as_struct_helper;
// Map of io struct to helper function to return the structure with the depth clamping
// applied.
Hashmap<const Struct*, Symbol, 4u> io_structs_clamp_helpers;
// Register a callback that will be called for each visted AST function.
// This call wraps the cloning of the function's statements, and will assign to
// `returns_frag_depth_as_value` or `returns_frag_depth_as_struct_helper` if the function's
// return value requires depth clamping.
ctx.ReplaceAll([&](const Function* fn) {
if (fn->PipelineStage() != PipelineStage::kFragment) {
return ctx.CloneWithoutTransform(fn);
}
if (ReturnsFragDepthAsValue(fn)) {
TINT_SCOPED_ASSIGNMENT(returns_frag_depth_as_value, true);
return ctx.CloneWithoutTransform(fn);
}
if (ReturnsFragDepthInStruct(fn)) {
// At most once per I/O struct, add the conversion function:
//
// fn clamp_frag_depth_S(s : S) -> S {
// return S(s.first, s.second, clamp_frag_depth(s.frag_depth), s.last);
// }
auto* struct_ty = sem.Get(fn)->ReturnType()->As<sem::Struct>()->Declaration();
auto helper = io_structs_clamp_helpers.GetOrCreate(struct_ty, [&] {
auto return_ty = fn->return_type;
auto fn_sym =
b.Symbols().New("clamp_frag_depth_" + struct_ty->name->symbol.Name());
tint::Vector<const Expression*, 8u> initializer_args;
for (auto* member : struct_ty->members) {
const Expression* arg =
b.MemberAccessor("s", ctx.Clone(member->name->symbol));
if (ContainsFragDepth(member->attributes)) {
arg = b.Call(base_fn_sym, arg);
}
initializer_args.Push(arg);
}
tint::Vector params{b.Param("s", ctx.Clone(return_ty))};
tint::Vector body{
b.Return(b.Call(ctx.Clone(return_ty), std::move(initializer_args))),
};
b.Func(fn_sym, params, ctx.Clone(return_ty), body);
return fn_sym;
});
TINT_SCOPED_ASSIGNMENT(returns_frag_depth_as_struct_helper, helper);
return ctx.CloneWithoutTransform(fn);
}
return ctx.CloneWithoutTransform(fn);
});
// Replace the return statements `return expr` with `return clamp_frag_depth(expr)`.
ctx.ReplaceAll([&](const ReturnStatement* stmt) -> const ReturnStatement* {
if (returns_frag_depth_as_value) {
return b.Return(stmt->source, b.Call(base_fn_sym, ctx.Clone(stmt->value)));
}
if (returns_frag_depth_as_struct_helper.IsValid()) {
return b.Return(stmt->source, b.Call(returns_frag_depth_as_struct_helper,
ctx.Clone(stmt->value)));
}
return nullptr;
});
ctx.Clone();
return Program(std::move(b));
}
private:
/// @returns true if the transform should run
bool ShouldRun() {
for (auto* fn : src->AST().Functions()) {
if (fn->PipelineStage() == PipelineStage::kFragment &&
(ReturnsFragDepthAsValue(fn) || ReturnsFragDepthInStruct(fn))) {
return true;
}
}
return false;
}
/// @param attrs the attributes to examine
/// @returns true if @p attrs contains a `@builtin(frag_depth)` attribute
bool ContainsFragDepth(VectorRef<const Attribute*> attrs) {
for (auto* attribute : attrs) {
if (auto* builtin_attr = attribute->As<BuiltinAttribute>()) {
auto builtin = sem.Get(builtin_attr)->Value();
if (builtin == builtin::BuiltinValue::kFragDepth) {
return true;
}
}
}
return false;
}
/// @param fn the function to examine
/// @returns true if @p fn has a return type with a `@builtin(frag_depth)` attribute
bool ReturnsFragDepthAsValue(const Function* fn) {
return ContainsFragDepth(fn->return_type_attributes);
}
/// @param fn the function to examine
/// @returns true if @p fn has a return structure with a `@builtin(frag_depth)` attribute on one
/// of the members
bool ReturnsFragDepthInStruct(const Function* fn) {
if (auto* struct_ty = sem.Get(fn)->ReturnType()->As<sem::Struct>()) {
for (auto* member : struct_ty->Members()) {
if (ContainsFragDepth(member->Declaration()->attributes)) {
return true;
}
}
}
return false;
}
};
ClampFragDepth::ClampFragDepth() = default;
ClampFragDepth::~ClampFragDepth() = default;
Transform::ApplyResult ClampFragDepth::Apply(const Program* src, const DataMap&, DataMap&) const {
return State{src}.Run();
}
} // namespace tint::ast::transform