Move workgroup_size property into sem::Function
The workgroup size should not be a property of the function in the
AST, and this lays the groundwork for allowing both literals and
module-scope constants to be used for this attribute.
Bug: tint:713
Change-Id: I014be879e2adb81cfc5b0ea0e221035fae626223
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51261
Auto-Submit: James Price <jrprice@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
diff --git a/src/ast/function.cc b/src/ast/function.cc
index 24b0227..0896e7a 100644
--- a/src/ast/function.cc
+++ b/src/ast/function.cc
@@ -58,13 +58,6 @@
Function::~Function() = default;
-std::tuple<uint32_t, uint32_t, uint32_t> Function::workgroup_size() const {
- if (auto* workgroup = GetDecoration<WorkgroupDecoration>(decorations_)) {
- return workgroup->values();
- }
- return {1, 1, 1};
-}
-
PipelineStage Function::pipeline_stage() const {
if (auto* stage = GetDecoration<StageDecoration>(decorations_)) {
return stage->value();
diff --git a/src/ast/function.h b/src/ast/function.h
index fd91c7c..397e1c4 100644
--- a/src/ast/function.h
+++ b/src/ast/function.h
@@ -66,10 +66,6 @@
/// @returns the decorations attached to this function
const DecorationList& decorations() const { return decorations_; }
- /// @returns the workgroup size {x, y, z} for the function. {1, 1, 1} will be
- /// return if no workgroup size was set.
- std::tuple<uint32_t, uint32_t, uint32_t> workgroup_size() const;
-
/// @returns the functions pipeline stage or None if not set
PipelineStage pipeline_stage() const;
diff --git a/src/ast/function_test.cc b/src/ast/function_test.cc
index 60c1525..44ba686 100644
--- a/src/ast/function_test.cc
+++ b/src/ast/function_test.cc
@@ -225,31 +225,6 @@
EXPECT_EQ(f->get_last_statement(), nullptr);
}
-TEST_F(FunctionTest, WorkgroupSize_NoneSet) {
- auto* f = Func("func", VariableList{}, ty.void_(), StatementList{},
- DecorationList{});
- uint32_t x = 0;
- uint32_t y = 0;
- uint32_t z = 0;
- std::tie(x, y, z) = f->workgroup_size();
- EXPECT_EQ(x, 1u);
- EXPECT_EQ(y, 1u);
- EXPECT_EQ(z, 1u);
-}
-
-TEST_F(FunctionTest, WorkgroupSize) {
- auto* f = Func("func", VariableList{}, ty.void_(), StatementList{},
- DecorationList{create<WorkgroupDecoration>(2u, 4u, 6u)});
-
- uint32_t x = 0;
- uint32_t y = 0;
- uint32_t z = 0;
- std::tie(x, y, z) = f->workgroup_size();
- EXPECT_EQ(x, 2u);
- EXPECT_EQ(y, 4u);
- EXPECT_EQ(z, 6u);
-}
-
using FunctionListTest = TestHelper;
TEST_F(FunctionListTest, FindSymbol) {
diff --git a/src/inspector/inspector.cc b/src/inspector/inspector.cc
index b8ac9b4..b687b22 100644
--- a/src/inspector/inspector.cc
+++ b/src/inspector/inspector.cc
@@ -198,8 +198,16 @@
entry_point.name = program_->Symbols().NameFor(func->symbol());
entry_point.remapped_name = program_->Symbols().NameFor(func->symbol());
entry_point.stage = func->pipeline_stage();
- std::tie(entry_point.workgroup_size_x, entry_point.workgroup_size_y,
- entry_point.workgroup_size_z) = func->workgroup_size();
+
+ auto wgsize = sem->workgroup_size();
+ entry_point.workgroup_size_x = wgsize[0].value;
+ entry_point.workgroup_size_y = wgsize[1].value;
+ entry_point.workgroup_size_z = wgsize[2].value;
+ if (wgsize[0].overridable_const || wgsize[1].overridable_const ||
+ wgsize[2].overridable_const) {
+ // TODO(crbug.com/tint/713): Handle overridable constants.
+ TINT_ASSERT(false);
+ }
for (auto* param : sem->Parameters()) {
AddEntryPointInOutVariables(
diff --git a/src/resolver/resolver.cc b/src/resolver/resolver.cc
index a80472b..93266aa 100644
--- a/src/resolver/resolver.cc
+++ b/src/resolver/resolver.cc
@@ -1287,6 +1287,20 @@
Mark(deco);
}
+ // Set work-group size defaults.
+ for (int i = 0; i < 3; i++) {
+ info->workgroup_size[i].value = 1;
+ info->workgroup_size[i].overridable_const = nullptr;
+ }
+
+ if (auto* workgroup =
+ ast::GetDecoration<ast::WorkgroupDecoration>(func->decorations())) {
+ // TODO(crbug.com/tint/713): Handle non-literals.
+ info->workgroup_size[0].value = std::get<0>(workgroup->values());
+ info->workgroup_size[1].value = std::get<1>(workgroup->values());
+ info->workgroup_size[2].value = std::get<2>(workgroup->values());
+ }
+
if (!ValidateFunction(func, info)) {
return false;
}
@@ -2517,7 +2531,7 @@
info->declaration, const_cast<sem::Type*>(info->return_type),
remap_vars(info->parameters), remap_vars(info->referenced_module_vars),
remap_vars(info->local_referenced_module_vars), info->return_statements,
- ancestor_entry_points[func->symbol()]);
+ ancestor_entry_points[func->symbol()], info->workgroup_size);
func_info_to_sem_func.emplace(info, sem_func);
sem.Add(func, sem_func);
}
diff --git a/src/resolver/resolver.h b/src/resolver/resolver.h
index 76a722f..ad428a8 100644
--- a/src/resolver/resolver.h
+++ b/src/resolver/resolver.h
@@ -26,6 +26,7 @@
#include "src/scope_stack.h"
#include "src/sem/binding_point.h"
#include "src/sem/block_statement.h"
+#include "src/sem/function.h"
#include "src/sem/struct.h"
#include "src/utils/unique_vector.h"
@@ -112,6 +113,7 @@
std::vector<const ast::ReturnStatement*> return_statements;
sem::Type* return_type = nullptr;
std::string return_type_name;
+ std::array<sem::WorkgroupDimension, 3> workgroup_size;
// List of transitive calls this function makes
UniqueVector<FunctionInfo*> transitive_calls;
diff --git a/src/resolver/resolver_test.cc b/src/resolver/resolver_test.cc
index e81be39..109b31a 100644
--- a/src/resolver/resolver_test.cc
+++ b/src/resolver/resolver_test.cc
@@ -32,6 +32,7 @@
#include "src/ast/switch_statement.h"
#include "src/ast/unary_op_expression.h"
#include "src/ast/variable_decl_statement.h"
+#include "src/ast/workgroup_decoration.h"
#include "src/resolver/resolver_test_helper.h"
#include "src/sem/call.h"
#include "src/sem/function.h"
@@ -887,6 +888,40 @@
EXPECT_TRUE(func_sem->ReturnType()->Is<sem::F32>());
}
+TEST_F(ResolverTest, Function_WorkgroupSize_NotSet) {
+ auto* func = Func("main", ast::VariableList{}, ty.void_(), {}, {});
+
+ EXPECT_TRUE(r()->Resolve()) << r()->error();
+
+ auto* func_sem = Sem().Get(func);
+ ASSERT_NE(func_sem, nullptr);
+
+ EXPECT_EQ(func_sem->workgroup_size()[0].value, 1u);
+ EXPECT_EQ(func_sem->workgroup_size()[1].value, 1u);
+ EXPECT_EQ(func_sem->workgroup_size()[2].value, 1u);
+ EXPECT_EQ(func_sem->workgroup_size()[0].overridable_const, nullptr);
+ EXPECT_EQ(func_sem->workgroup_size()[1].overridable_const, nullptr);
+ EXPECT_EQ(func_sem->workgroup_size()[2].overridable_const, nullptr);
+}
+
+TEST_F(ResolverTest, Function_WorkgroupSize_Literals) {
+ auto* func = Func("main", ast::VariableList{}, ty.void_(), {},
+ {Stage(ast::PipelineStage::kCompute),
+ create<ast::WorkgroupDecoration>(8, 2, 3)});
+
+ EXPECT_TRUE(r()->Resolve()) << r()->error();
+
+ auto* func_sem = Sem().Get(func);
+ ASSERT_NE(func_sem, nullptr);
+
+ EXPECT_EQ(func_sem->workgroup_size()[0].value, 8u);
+ EXPECT_EQ(func_sem->workgroup_size()[1].value, 2u);
+ EXPECT_EQ(func_sem->workgroup_size()[2].value, 3u);
+ EXPECT_EQ(func_sem->workgroup_size()[0].overridable_const, nullptr);
+ EXPECT_EQ(func_sem->workgroup_size()[1].overridable_const, nullptr);
+ EXPECT_EQ(func_sem->workgroup_size()[2].overridable_const, nullptr);
+}
+
TEST_F(ResolverTest, Expr_MemberAccessor_Struct) {
auto* st = Structure("S", {Member("first_member", ty.i32()),
Member("second_member", ty.f32())});
diff --git a/src/sem/function.cc b/src/sem/function.cc
index 0c5e844..f33fc07 100644
--- a/src/sem/function.cc
+++ b/src/sem/function.cc
@@ -46,14 +46,16 @@
std::vector<const Variable*> referenced_module_vars,
std::vector<const Variable*> local_referenced_module_vars,
std::vector<const ast::ReturnStatement*> return_statements,
- std::vector<Symbol> ancestor_entry_points)
+ std::vector<Symbol> ancestor_entry_points,
+ std::array<WorkgroupDimension, 3> workgroup_size)
: Base(return_type, GetParameters(parameters)),
declaration_(declaration),
parameters_(std::move(parameters)),
referenced_module_vars_(std::move(referenced_module_vars)),
local_referenced_module_vars_(std::move(local_referenced_module_vars)),
return_statements_(std::move(return_statements)),
- ancestor_entry_points_(std::move(ancestor_entry_points)) {}
+ ancestor_entry_points_(std::move(ancestor_entry_points)),
+ workgroup_size_(std::move(workgroup_size)) {}
Function::~Function() = default;
diff --git a/src/sem/function.h b/src/sem/function.h
index a29e318..94cc02e 100644
--- a/src/sem/function.h
+++ b/src/sem/function.h
@@ -15,6 +15,7 @@
#ifndef SRC_SEM_FUNCTION_H_
#define SRC_SEM_FUNCTION_H_
+#include <array>
#include <utility>
#include <vector>
@@ -37,6 +38,16 @@
class Variable;
+/// WorkgroupDimension describes the size of a single dimension of an entry
+/// point's workgroup size.
+struct WorkgroupDimension {
+ /// The size of this dimension.
+ uint32_t value;
+ /// A pipeline-overridable constant that overrides the size, or nullptr if
+ /// this dimension is not overridable.
+ const ast::Variable* overridable_const = nullptr;
+};
+
/// Function holds the semantic information for function nodes.
class Function : public Castable<Function, CallTarget> {
public:
@@ -53,13 +64,15 @@
/// @param return_statements the function return statements
/// variables
/// @param ancestor_entry_points the ancestor entry points
+ /// @param workgroup_size the workgroup size
Function(ast::Function* declaration,
Type* return_type,
std::vector<const Variable*> parameters,
std::vector<const Variable*> referenced_module_vars,
std::vector<const Variable*> local_referenced_module_vars,
std::vector<const ast::ReturnStatement*> return_statements,
- std::vector<Symbol> ancestor_entry_points);
+ std::vector<Symbol> ancestor_entry_points,
+ std::array<WorkgroupDimension, 3> workgroup_size);
/// Destructor
~Function() override;
@@ -148,6 +161,11 @@
/// @returns true if `sym` is an ancestor entry point of this function
bool HasAncestorEntryPoint(Symbol sym) const;
+ /// @returns the workgroup size {x, y, z} for the function.
+ const std::array<WorkgroupDimension, 3>& workgroup_size() const {
+ return workgroup_size_;
+ }
+
private:
VariableBindings ReferencedSamplerVariablesImpl(ast::SamplerKind kind) const;
VariableBindings ReferencedSampledTextureVariablesImpl(
@@ -159,6 +177,7 @@
std::vector<const Variable*> const local_referenced_module_vars_;
std::vector<const ast::ReturnStatement*> const return_statements_;
std::vector<Symbol> const ancestor_entry_points_;
+ std::array<WorkgroupDimension, 3> workgroup_size_;
};
} // namespace sem
diff --git a/src/writer/hlsl/generator_impl.cc b/src/writer/hlsl/generator_impl.cc
index e894106..b516533 100644
--- a/src/writer/hlsl/generator_impl.cc
+++ b/src/writer/hlsl/generator_impl.cc
@@ -1989,12 +1989,19 @@
make_indent(out);
current_ep_sym_ = func->symbol();
+ auto* func_sem = builder_.Sem().Get(func);
if (func->pipeline_stage() == ast::PipelineStage::kCompute) {
- uint32_t x = 0;
- uint32_t y = 0;
- uint32_t z = 0;
- std::tie(x, y, z) = func->workgroup_size();
+ auto wgsize = func_sem->workgroup_size();
+ if (wgsize[0].overridable_const || wgsize[1].overridable_const ||
+ wgsize[2].overridable_const) {
+ // TODO(crbug.com/tint/713): Handle overridable constants.
+ TINT_UNIMPLEMENTED(builder_.Diagnostics())
+ << "pipeline-overridable workgroup sizes are not implemented";
+ }
+ uint32_t x = wgsize[0].value;
+ uint32_t y = wgsize[1].value;
+ uint32_t z = wgsize[2].value;
out << "[numthreads(" << std::to_string(x) << ", " << std::to_string(y)
<< ", " << std::to_string(z) << ")]" << std::endl;
make_indent(out);
diff --git a/src/writer/spirv/builder.cc b/src/writer/spirv/builder.cc
index 22f4a6a..60c159c 100644
--- a/src/writer/spirv/builder.cc
+++ b/src/writer/spirv/builder.cc
@@ -435,23 +435,30 @@
}
bool Builder::GenerateExecutionModes(ast::Function* func, uint32_t id) {
+ auto* func_sem = builder_.Sem().Get(func);
+
// WGSL fragment shader origin is upper left
if (func->pipeline_stage() == ast::PipelineStage::kFragment) {
push_execution_mode(
spv::Op::OpExecutionMode,
{Operand::Int(id), Operand::Int(SpvExecutionModeOriginUpperLeft)});
} else if (func->pipeline_stage() == ast::PipelineStage::kCompute) {
- uint32_t x = 0;
- uint32_t y = 0;
- uint32_t z = 0;
- std::tie(x, y, z) = func->workgroup_size();
+ auto& wgsize = func_sem->workgroup_size();
+ if (wgsize[0].overridable_const || wgsize[1].overridable_const ||
+ wgsize[2].overridable_const) {
+ // TODO(crbug.com/tint/713): Handle overridable constants.
+ TINT_UNIMPLEMENTED(builder_.Diagnostics())
+ << "pipeline-overridable workgroup sizes are not implemented";
+ }
+ uint32_t x = wgsize[0].value;
+ uint32_t y = wgsize[1].value;
+ uint32_t z = wgsize[2].value;
push_execution_mode(
spv::Op::OpExecutionMode,
{Operand::Int(id), Operand::Int(SpvExecutionModeLocalSize),
Operand::Int(x), Operand::Int(y), Operand::Int(z)});
}
- auto* func_sem = builder_.Sem().Get(func);
for (auto builtin : func_sem->ReferencedBuiltinVariables()) {
if (builtin.second->value() == ast::Builtin::kFragDepth) {
push_execution_mode(
