[spirv-reader][ir] Support `WorkgroupSize` annotation.
Adds support for creating computer shader workgroup size information
from a `WorkgroupSize` annotated variable.
Bug: 42250952
Change-Id: I54d3112b658a4b9c62c3334bf8ced143aedb6482
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/245995
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: James Price <jrprice@google.com>
diff --git a/src/tint/lang/spirv/reader/parser/function_test.cc b/src/tint/lang/spirv/reader/parser/function_test.cc
index 2b53807..98a2aea 100644
--- a/src/tint/lang/spirv/reader/parser/function_test.cc
+++ b/src/tint/lang/spirv/reader/parser/function_test.cc
@@ -74,6 +74,101 @@
)");
}
+TEST_F(SpirvParserTest, WorkgroupSize_Constant) {
+ EXPECT_IR(R"(
+ OpCapability Shader
+ OpMemoryModel Logical GLSL450
+ OpEntryPoint GLCompute %main "main"
+ OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize
+ %void = OpTypeVoid
+ %uint = OpTypeInt 32 0
+ %v3uint = OpTypeVector %uint 3
+ %uint_3 = OpConstant %uint 3
+ %uint_5 = OpConstant %uint 5
+ %uint_7 = OpConstant %uint 7
+%gl_WorkGroupSize = OpConstantComposite %v3uint %uint_3 %uint_5 %uint_7
+ %ep_type = OpTypeFunction %void
+ %main = OpFunction %void None %ep_type
+ %main_start = OpLabel
+ OpReturn
+ OpFunctionEnd
+)",
+ R"(
+%main = @compute @workgroup_size(3u, 5u, 7u) func():void {
+ $B1: {
+ ret
+ }
+}
+)");
+}
+
+TEST_F(SpirvParserTest, WorkgroupSize_SpecConstant_Mixed) {
+ EXPECT_IR(R"(
+ OpCapability Shader
+ OpMemoryModel Logical GLSL450
+ OpEntryPoint GLCompute %main "main"
+ OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize
+ %void = OpTypeVoid
+ %uint = OpTypeInt 32 0
+ %v3uint = OpTypeVector %uint 3
+ %uint_3 = OpSpecConstant %uint 3
+ %uint_5 = OpConstant %uint 5
+ %uint_7 = OpSpecConstant %uint 7
+%gl_WorkGroupSize = OpSpecConstantComposite %v3uint %uint_3 %uint_5 %uint_7
+ %ep_type = OpTypeFunction %void
+ %main = OpFunction %void None %ep_type
+ %main_start = OpLabel
+ OpReturn
+ OpFunctionEnd
+)",
+ R"(
+$B1: { # root
+ %1:u32 = override 3u
+ %2:u32 = override 7u
+}
+
+%main = @compute @workgroup_size(%1, 5u, %2) func():void {
+ $B2: {
+ ret
+ }
+}
+)");
+}
+
+TEST_F(SpirvParserTest, WorkgroupSize_SpecConstant) {
+ EXPECT_IR(R"(
+ OpCapability Shader
+ OpMemoryModel Logical GLSL450
+ OpEntryPoint GLCompute %main "main"
+ OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize
+ %void = OpTypeVoid
+ %uint = OpTypeInt 32 0
+ %v3uint = OpTypeVector %uint 3
+ %uint_3 = OpSpecConstant %uint 3
+ %uint_5 = OpSpecConstant %uint 5
+ %uint_7 = OpSpecConstant %uint 7
+%gl_WorkGroupSize = OpSpecConstantComposite %v3uint %uint_3 %uint_5 %uint_7
+ %ep_type = OpTypeFunction %void
+ %main = OpFunction %void None %ep_type
+ %main_start = OpLabel
+ OpReturn
+ OpFunctionEnd
+)",
+ R"(
+$B1: { # root
+ %1:u32 = override 3u
+ %2:u32 = override 5u
+ %3:u32 = override 7u
+}
+
+%main = @compute @workgroup_size(%1, %2, %3) func():void {
+ $B2: {
+ ret
+ }
+}
+)");
+}
+
TEST_F(SpirvParserTest, FragmentShader) {
EXPECT_IR(R"(
OpCapability Shader
diff --git a/src/tint/lang/spirv/reader/parser/parser.cc b/src/tint/lang/spirv/reader/parser/parser.cc
index b7c6cb2..6a82af2 100644
--- a/src/tint/lang/spirv/reader/parser/parser.cc
+++ b/src/tint/lang/spirv/reader/parser/parser.cc
@@ -1183,6 +1183,39 @@
// first one found.
break;
}
+ } else if (func->IsCompute()) {
+ // Search for `WorkgroupSize` decorated Ids
+ for (const spvtools::opt::Instruction& inst :
+ spirv_context_->module()->annotations()) {
+ if (inst.opcode() != spv::Op::OpDecorate ||
+ inst.GetSingleWordInOperand(1) != uint32_t(spv::Decoration::BuiltIn) ||
+ inst.GetSingleWordInOperand(2) != uint32_t(spv::BuiltIn::WorkgroupSize)) {
+ continue;
+ }
+ uint32_t id = inst.GetSingleWordInOperand(0);
+
+ Vector<core::ir::Value*, 3> args;
+ if (auto* c = SpvConstant(id)) {
+ auto* vals = c->AsVectorConstant();
+ TINT_ASSERT(vals);
+
+ for (auto& el : vals->GetComponents()) {
+ args.Push(b_.Constant(Constant(el)));
+ }
+ } else {
+ TINT_ASSERT(spec_composites_.contains(id));
+
+ auto info = spec_composites_[id];
+ TINT_ASSERT(info.args.Length() == 3);
+
+ for (auto arg : info.args) {
+ args.Push(Value(arg));
+ }
+ }
+ func->SetWorkgroupSize(args[0], args[1], args[2]);
+
+ break;
+ }
}
}
