[tint] Support attribute `subgroup_size` in WGSL resolver This patch adds `SubgroupSizeAttribute` to WGSL AST and implements the validations on the attribute `subgroup_size` in WGSL resolver. Bug: 463721943 Change-Id: I48896a52e827f862b8b6036ca971d1d7a07f0221 Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/281155 Reviewed-by: James Price <jrprice@google.com> Commit-Queue: Shao, Jiawei <jiawei.shao@intel.com> Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/tint/lang/wgsl/ast/BUILD.bazel b/src/tint/lang/wgsl/ast/BUILD.bazel index c5b0c42..0a75e11 100644 --- a/src/tint/lang/wgsl/ast/BUILD.bazel +++ b/src/tint/lang/wgsl/ast/BUILD.bazel
@@ -103,6 +103,7 @@ "struct_member.cc", "struct_member_align_attribute.cc", "struct_member_size_attribute.cc", + "subgroup_size_attribute.cc", "switch_statement.cc", "templated_identifier.cc", "type_decl.cc", @@ -179,6 +180,7 @@ "struct_member.h", "struct_member_align_attribute.h", "struct_member_size_attribute.h", + "subgroup_size_attribute.h", "switch_statement.h", "templated_identifier.h", "traverse_expressions.h",
diff --git a/src/tint/lang/wgsl/ast/BUILD.cmake b/src/tint/lang/wgsl/ast/BUILD.cmake index 85031bf..492cf6c 100644 --- a/src/tint/lang/wgsl/ast/BUILD.cmake +++ b/src/tint/lang/wgsl/ast/BUILD.cmake
@@ -168,6 +168,8 @@ lang/wgsl/ast/struct_member_align_attribute.h lang/wgsl/ast/struct_member_size_attribute.cc lang/wgsl/ast/struct_member_size_attribute.h + lang/wgsl/ast/subgroup_size_attribute.cc + lang/wgsl/ast/subgroup_size_attribute.h lang/wgsl/ast/switch_statement.cc lang/wgsl/ast/switch_statement.h lang/wgsl/ast/templated_identifier.cc
diff --git a/src/tint/lang/wgsl/ast/BUILD.gn b/src/tint/lang/wgsl/ast/BUILD.gn index a9300b5..9961f96 100644 --- a/src/tint/lang/wgsl/ast/BUILD.gn +++ b/src/tint/lang/wgsl/ast/BUILD.gn
@@ -174,6 +174,8 @@ "struct_member_align_attribute.h", "struct_member_size_attribute.cc", "struct_member_size_attribute.h", + "subgroup_size_attribute.cc", + "subgroup_size_attribute.h", "switch_statement.cc", "switch_statement.h", "templated_identifier.cc",
diff --git a/src/tint/lang/wgsl/ast/builder.h b/src/tint/lang/wgsl/ast/builder.h index faee3fe..1644c68 100644 --- a/src/tint/lang/wgsl/ast/builder.h +++ b/src/tint/lang/wgsl/ast/builder.h
@@ -86,6 +86,7 @@ #include "src/tint/lang/wgsl/ast/struct.h" #include "src/tint/lang/wgsl/ast/struct_member_align_attribute.h" #include "src/tint/lang/wgsl/ast/struct_member_size_attribute.h" +#include "src/tint/lang/wgsl/ast/subgroup_size_attribute.h" #include "src/tint/lang/wgsl/ast/switch_statement.h" #include "src/tint/lang/wgsl/ast/templated_identifier.h" #include "src/tint/lang/wgsl/ast/type.h" @@ -3318,6 +3319,23 @@ Expr(std::forward<EXPR_Z>(z))); } + /// Creates an ast::SubgroupSizeAttribute + /// @param source the source information + /// @param subgroup_size the subgroup size value expression + /// @returns the subgroup size attribute pointer + template <typename EXPR> + const ast::SubgroupSizeAttribute* SubgroupSize(const Source& source, EXPR&& subgroup_size) { + return create<ast::SubgroupSizeAttribute>(source, std::forward<EXPR>(subgroup_size)); + } + + /// Creates an ast::SubgroupSizeAttribute + /// @param subgroup_size the subgroup size value expression + /// @returns the subgroup size attribute pointer + template <typename EXPR> + const ast::SubgroupSizeAttribute* SubgroupSize(EXPR&& subgroup_size) { + return SubgroupSize(source_, Expr(std::forward<EXPR>(subgroup_size))); + } + /// Passthrough overload /// @param name the diagnostic rule name /// @returns @p name
diff --git a/src/tint/lang/wgsl/ast/subgroup_size_attribute.cc b/src/tint/lang/wgsl/ast/subgroup_size_attribute.cc new file mode 100644 index 0000000..47ab8a2 --- /dev/null +++ b/src/tint/lang/wgsl/ast/subgroup_size_attribute.cc
@@ -0,0 +1,49 @@ +// Copyright 2025 The Dawn & Tint Authors +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +#include "src/tint/lang/wgsl/ast/subgroup_size_attribute.h" + +#include <string> + +#include "src/tint/lang/wgsl/ast/builder.h" + +TINT_INSTANTIATE_TYPEINFO(tint::ast::SubgroupSizeAttribute); + +namespace tint::ast { + +SubgroupSizeAttribute::SubgroupSizeAttribute(NodeID nid, + const Source& src, + const Expression* subgroup_size_) + : Base(nid, src), subgroup_size(subgroup_size_) {} + +SubgroupSizeAttribute::~SubgroupSizeAttribute() = default; + +std::string SubgroupSizeAttribute::Name() const { + return "subgroup_size"; +} + +} // namespace tint::ast
diff --git a/src/tint/lang/wgsl/ast/subgroup_size_attribute.h b/src/tint/lang/wgsl/ast/subgroup_size_attribute.h new file mode 100644 index 0000000..e841192 --- /dev/null +++ b/src/tint/lang/wgsl/ast/subgroup_size_attribute.h
@@ -0,0 +1,65 @@ +// Copyright 2025 The Dawn & Tint Authors +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +#ifndef SRC_TINT_LANG_WGSL_AST_SUBGROUP_SIZE_ATTRIBUTE_H_ +#define SRC_TINT_LANG_WGSL_AST_SUBGROUP_SIZE_ATTRIBUTE_H_ + +#include <string> + +#include "src/tint/lang/wgsl/ast/attribute.h" + +// Forward declarations +namespace tint::ast { +class Expression; +} // namespace tint::ast + +namespace tint::ast { + +/// A subgroup_size attribute +class SubgroupSizeAttribute final : public Castable<SubgroupSizeAttribute, Attribute> { + public: + /// constructor + /// @param nid the unique node identifier + /// @param src the source of this node + /// @param subgroup_size_ the subgroup_size expression + SubgroupSizeAttribute(NodeID nid, const Source& src, const Expression* subgroup_size_); + + ~SubgroupSizeAttribute() override; + + /// @returns the subgroup_size value + const Expression* Value() const { return subgroup_size; } + + /// @returns the WGSL name for the attribute + std::string Name() const override; + + /// The subgroup_size expression. + const Expression* const subgroup_size; +}; + +} // namespace tint::ast + +#endif // SRC_TINT_LANG_WGSL_AST_SUBGROUP_SIZE_ATTRIBUTE_H_
diff --git a/src/tint/lang/wgsl/reader/parser/parser.cc b/src/tint/lang/wgsl/reader/parser/parser.cc index 9870164..252681c 100644 --- a/src/tint/lang/wgsl/reader/parser/parser.cc +++ b/src/tint/lang/wgsl/reader/parser/parser.cc
@@ -50,6 +50,7 @@ #include "src/tint/lang/wgsl/ast/loop_statement.h" #include "src/tint/lang/wgsl/ast/return_statement.h" #include "src/tint/lang/wgsl/ast/stage_attribute.h" +#include "src/tint/lang/wgsl/ast/subgroup_size_attribute.h" #include "src/tint/lang/wgsl/ast/switch_statement.h" #include "src/tint/lang/wgsl/ast/unary_op_expression.h" #include "src/tint/lang/wgsl/ast/var.h" @@ -3154,6 +3155,8 @@ return create<ast::WorkgroupAttribute>(t.source(), args[0], args.Length() > 1 ? args[1] : nullptr, args.Length() > 2 ? args[2] : nullptr); + case core::Attribute::kSubgroupSize: + return create<ast::SubgroupSizeAttribute>(t.source(), args[0]); default: return Failure::kNoMatch; }
diff --git a/src/tint/lang/wgsl/resolver/attribute_validation_test.cc b/src/tint/lang/wgsl/resolver/attribute_validation_test.cc index 148aa5a..6ba61fb 100644 --- a/src/tint/lang/wgsl/resolver/attribute_validation_test.cc +++ b/src/tint/lang/wgsl/resolver/attribute_validation_test.cc
@@ -71,6 +71,7 @@ kSize, kStageCompute, kWorkgroupSize, + kSubgroupSize, }; static std::ostream& operator<<(std::ostream& o, AttributeKind k) { switch (k) { @@ -106,6 +107,8 @@ return o << "@compute"; case AttributeKind::kWorkgroupSize: return o << "@workgroup_size"; + case AttributeKind::kSubgroupSize: + return o << "@subgroup_size"; } TINT_UNREACHABLE(); } @@ -195,6 +198,10 @@ TestParams{ {AttributeKind::kBinding, AttributeKind::kGroup}, "1:2 error: '@binding' is not valid for " + thing, + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + "1:2 error: '@subgroup_size' is not valid for " + thing, }}; } @@ -239,7 +246,9 @@ case AttributeKind::kStageCompute: return builder.Stage(source, ast::PipelineStage::kCompute); case AttributeKind::kWorkgroupSize: - return builder.create<ast::WorkgroupAttribute>(source, builder.Expr(1_i)); + return builder.WorkgroupSize(source, builder.Expr(1_i)); + case AttributeKind::kSubgroupSize: + return builder.SubgroupSize(source, builder.Expr(16_i)); } TINT_UNREACHABLE() << kind; } @@ -256,6 +265,10 @@ case AttributeKind::kInputAttachmentIndex: Enable(wgsl::Extension::kChromiumInternalInputAttachments); break; + case AttributeKind::kSubgroupSize: + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + break; default: break; } @@ -371,6 +384,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is only valid for compute stages)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is only valid for compute stages)", })); using NonVoidFunctionAttributeTest = TestWithParams; @@ -452,6 +469,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is only valid for compute stages)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is only valid for compute stages)", })); } // namespace FunctionTests @@ -535,6 +556,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for function parameters)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for function parameters)", })); using FunctionReturnTypeAttributeTest = TestWithParams; @@ -616,6 +641,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for non-entry point function return types)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for non-entry point function return types)", })); } // namespace FunctionInputAndOutputTests @@ -702,6 +731,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for function parameters)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for function parameters)", })); using FragmentShaderParameterAttributeTest = TestWithParams; @@ -799,6 +832,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for function parameters)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for function parameters)", })); using VertexShaderParameterAttributeTest = TestWithParams; @@ -902,6 +939,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for function parameters)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for function parameters)", })); using ComputeShaderReturnTypeAttributeTest = TestWithParams; @@ -987,6 +1028,10 @@ TestParams{ {AttributeKind::kWorkgroupSize}, R"(1:2 error: '@workgroup_size' is not valid for entry point return types)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for entry point return types)", })); using FragmentShaderReturnTypeAttributeTest = TestWithParams; @@ -1085,6 +1130,10 @@ TestParams{ {AttributeKind::kBinding, AttributeKind::kGroup}, R"(1:2 error: '@binding' is not valid for entry point return types)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for entry point return types)", })); using VertexShaderReturnTypeAttributeTest = TestWithParams; @@ -1183,6 +1232,10 @@ {AttributeKind::kLocation, AttributeKind::kLocation}, R"(3:4 error: duplicate location attribute 1:2 note: first attribute declared here)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for entry point return types)", })); } // namespace EntryPointInputAndOutputTests @@ -1267,6 +1320,10 @@ TestParams{ {AttributeKind::kBinding, AttributeKind::kGroup}, R"(1:2 error: '@binding' is not valid for 'struct' declarations)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for 'struct' declarations)", })); using StructMemberAttributeTest = TestWithParams; @@ -1360,6 +1417,10 @@ {AttributeKind::kAlign, AttributeKind::kAlign}, R"(3:4 error: duplicate align attribute 1:2 note: first attribute declared here)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for 'struct' members)", })); TEST_F(StructMemberAttributeTest, Align_Attribute_Const) { @@ -1615,6 +1676,10 @@ {AttributeKind::kBinding, AttributeKind::kGroup, AttributeKind::kBinding}, R"(5:6 error: duplicate binding attribute 1:2 note: first attribute declared here)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for module-scope 'var')", })); TEST_F(VariableAttributeTest, LocalVar) { @@ -1710,6 +1775,10 @@ TestParams{ {AttributeKind::kBinding, AttributeKind::kGroup}, R"(1:2 error: '@binding' is not valid for 'const' declaration)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for 'const' declaration)", })); using OverrideAttributeTest = TestWithParams; @@ -1788,6 +1857,10 @@ {AttributeKind::kId, AttributeKind::kId}, R"(3:4 error: duplicate id attribute 1:2 note: first attribute declared here)", + }, + TestParams{ + {AttributeKind::kSubgroupSize}, + R"(1:2 error: '@subgroup_size' is not valid for 'override' declaration)", })); using SwitchStatementAttributeTest = TestWithParams;
diff --git a/src/tint/lang/wgsl/resolver/dependency_graph.cc b/src/tint/lang/wgsl/resolver/dependency_graph.cc index 3e23026..0400b51 100644 --- a/src/tint/lang/wgsl/resolver/dependency_graph.cc +++ b/src/tint/lang/wgsl/resolver/dependency_graph.cc
@@ -64,6 +64,7 @@ #include "src/tint/lang/wgsl/ast/struct.h" #include "src/tint/lang/wgsl/ast/struct_member_align_attribute.h" #include "src/tint/lang/wgsl/ast/struct_member_size_attribute.h" +#include "src/tint/lang/wgsl/ast/subgroup_size_attribute.h" #include "src/tint/lang/wgsl/ast/switch_statement.h" #include "src/tint/lang/wgsl/ast/templated_identifier.h" #include "src/tint/lang/wgsl/ast/traverse_expressions.h" @@ -385,6 +386,7 @@ TraverseExpression(wg->y); TraverseExpression(wg->z); }, + [&](const ast::SubgroupSizeAttribute* sg) { TraverseExpression(sg->subgroup_size); }, [&](Default) { if (!attr->IsAnyOf<ast::BuiltinAttribute, ast::DiagnosticAttribute, ast::InterpolateAttribute, ast::InvariantAttribute,
diff --git a/src/tint/lang/wgsl/resolver/resolver.cc b/src/tint/lang/wgsl/resolver/resolver.cc index eb26470..5ea42fc 100644 --- a/src/tint/lang/wgsl/resolver/resolver.cc +++ b/src/tint/lang/wgsl/resolver/resolver.cc
@@ -889,6 +889,14 @@ func->SetWorkgroupSize(value.Get()); return true; }, + [&](const ast::SubgroupSizeAttribute* attr) { + auto value = SubgroupSizeAttribute(attr); + if (value != Success) { + return false; + } + func->SetSubgroupSize(value.Get()); + return true; + }, [&](Default) { ErrorInvalidAttribute(attribute, StyledText{} << "functions"); return false; @@ -3977,6 +3985,60 @@ return ws; } +tint::Result<uint32_t> Resolver::SubgroupSizeAttribute(const ast::SubgroupSizeAttribute* attr) { + auto value = attr->subgroup_size; + + auto err_bad_expr = [&]() { + AddError(attr->source) << style::Attribute("@subgroup_size") + << " argument must be a constant or override-expression of type " + << style::Type("abstract-integer") << ", " << style::Type("i32") + << " or " << style::Type("u32"); + }; + + const auto* expr = ValueExpression(value); + if (!expr) { + return Failure{}; + } + auto* type = expr->Type(); + if (!type->IsAnyOf<core::type::I32, core::type::U32, core::type::AbstractInt>()) { + err_bad_expr(); + return Failure{}; + } + + if (expr->Stage() != core::EvaluationStage::kConstant && + expr->Stage() != core::EvaluationStage::kOverride) { + err_bad_expr(); + return Failure{}; + } + + // If all arguments are abstract-integers, then materialize to i32. + if (type->Is<core::type::AbstractInt>()) { + type = b.create<core::type::I32>(); + } + + auto* materialized = Materialize(expr, type); + if (!materialized) { + return Failure{}; + } + + uint32_t subgroup_size = 0u; + if (auto* constant_value = materialized->ConstantValue()) { + if (constant_value->ValueAs<AInt>() < 1) { + AddError(attr->source) + << style::Attribute("@subgroup_size") << " argument must be at least 1"; + return Failure{}; + } + subgroup_size = constant_value->ValueAs<u32>(); + if (!IsPowerOfTwo(subgroup_size)) { + AddError(attr->source) + << style::Attribute("@subgroup_size") << " argument must be a power of 2"; + return Failure{}; + } + } + + return subgroup_size; +} + bool Resolver::DiagnosticAttribute(const ast::DiagnosticAttribute* attr) { return DiagnosticControl(attr->control); }
diff --git a/src/tint/lang/wgsl/resolver/resolver.h b/src/tint/lang/wgsl/resolver/resolver.h index f572551e..3e417ad 100644 --- a/src/tint/lang/wgsl/resolver/resolver.h +++ b/src/tint/lang/wgsl/resolver/resolver.h
@@ -440,6 +440,10 @@ /// @returns the workgroup size on success. tint::Result<sem::WorkgroupSize> WorkgroupAttribute(const ast::WorkgroupAttribute* attr); + /// Resolves the `@sugbroup_size` attribute @p attr + /// @returns the subgroup size on success. + tint::Result<uint32_t> SubgroupSizeAttribute(const ast::SubgroupSizeAttribute* attr); + /// Resolves the `@diagnostic` attribute @p attr /// @returns true on success, false on failure bool DiagnosticAttribute(const ast::DiagnosticAttribute* attr);
diff --git a/src/tint/lang/wgsl/resolver/subgroup_size_control_extension_test.cc b/src/tint/lang/wgsl/resolver/subgroup_size_control_extension_test.cc index fb35e5a..26929e2 100644 --- a/src/tint/lang/wgsl/resolver/subgroup_size_control_extension_test.cc +++ b/src/tint/lang/wgsl/resolver/subgroup_size_control_extension_test.cc
@@ -54,5 +54,142 @@ R"(error: extension 'chromium_experimental_subgroup_size_control' cannot be used without extension 'subgroups')"); } +TEST_F(SubgroupSizeControlExtensionTest, RequiresExtension) { + // Test without enabling the extension + Enable(wgsl::Extension::kSubgroups); + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(16_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ( + r()->error(), + R"(12:34 error: use of '@subgroup_size' requires enabling extension 'chromium_experimental_subgroup_size_control')"); +} + +TEST_F(SubgroupSizeControlExtensionTest, ValidWithExtension) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Expr(16_i)), + }); + + EXPECT_TRUE(r()->Resolve()) << r()->error(); +} + +TEST_F(SubgroupSizeControlExtensionTest, NotAnEntryPoint) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + SubgroupSize(Source{{12, 34}}, Expr(16_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), R"(12:34 error: '@subgroup_size' is only valid for compute stages)"); +} + +TEST_F(SubgroupSizeControlExtensionTest, NotAComputeShader) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kFragment), + SubgroupSize(Source{{12, 34}}, Expr(16_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), R"(12:34 error: '@subgroup_size' is only valid for compute stages)"); +} + +TEST_F(SubgroupSizeControlExtensionTest, InvalidValue_Negative) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(-1_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), R"(12:34 error: '@subgroup_size' argument must be at least 1)"); +} + +TEST_F(SubgroupSizeControlExtensionTest, InvalidValue_Zero) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(0_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), R"(12:34 error: '@subgroup_size' argument must be at least 1)"); +} + +TEST_F(SubgroupSizeControlExtensionTest, InvalidValue_NotAPowerOfTwo) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(15_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), R"(12:34 error: '@subgroup_size' argument must be a power of 2)"); +} + +TEST_F(SubgroupSizeControlExtensionTest, InvalidValue_Float) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(16.5_f)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ( + r()->error(), + R"(12:34 error: '@subgroup_size' argument must be a constant or override-expression of type 'abstract-integer', 'i32' or 'u32')"); +} + +TEST_F(SubgroupSizeControlExtensionTest, DuplicateAttribute) { + Enable(wgsl::Extension::kSubgroups); + Enable(wgsl::Extension::kChromiumExperimentalSubgroupSizeControl); + + Func("main", tint::Empty, ty.void_(), tint::Empty, + Vector{ + Stage(ast::PipelineStage::kCompute), + WorkgroupSize(1_i), + SubgroupSize(Source{{12, 34}}, Expr(16_i)), + SubgroupSize(Source{{56, 78}}, Expr(32_i)), + }); + + EXPECT_FALSE(r()->Resolve()); + EXPECT_EQ(r()->error(), + R"(56:78 error: duplicate subgroup_size attribute +12:34 note: first attribute declared here)"); +} + } // namespace } // namespace tint::resolver
diff --git a/src/tint/lang/wgsl/resolver/validator.cc b/src/tint/lang/wgsl/resolver/validator.cc index fe4d89d..4e6d6f5 100644 --- a/src/tint/lang/wgsl/resolver/validator.cc +++ b/src/tint/lang/wgsl/resolver/validator.cc
@@ -56,6 +56,7 @@ #include "src/tint/lang/wgsl/ast/id_attribute.h" #include "src/tint/lang/wgsl/ast/interpolate_attribute.h" #include "src/tint/lang/wgsl/ast/return_statement.h" +#include "src/tint/lang/wgsl/ast/subgroup_size_attribute.h" #include "src/tint/lang/wgsl/ast/switch_statement.h" #include "src/tint/lang/wgsl/ast/traverse_expressions.h" #include "src/tint/lang/wgsl/ast/variable_decl_statement.h" @@ -1358,6 +1359,22 @@ } return true; }, + [&](const ast::SubgroupSizeAttribute*) { + if (!enabled_extensions_.Contains( + wgsl::Extension::kChromiumExperimentalSubgroupSizeControl)) { + AddError(attr->source) + << "use of " << style::Attribute("@subgroup_size") + << " requires enabling extension " + << style::Code("chromium_experimental_subgroup_size_control"); + return false; + } + if (decl->PipelineStage() != ast::PipelineStage::kCompute) { + AddError(attr->source) << style::Attribute("@subgroup_size") + << " is only valid for compute stages"; + return false; + } + return true; + }, [&](Default) { return true; }); if (!ok) { return false;
diff --git a/src/tint/lang/wgsl/sem/function.h b/src/tint/lang/wgsl/sem/function.h index 740bf8c..94cca32 100644 --- a/src/tint/lang/wgsl/sem/function.h +++ b/src/tint/lang/wgsl/sem/function.h
@@ -84,6 +84,13 @@ workgroup_size_ = std::move(workgroup_size); } + /// Sets the subgroup size for the function. + /// @param subgroup_size the new subgroup size of the function + void SetSubgroupSize(uint32_t subgroup_size) { subgroup_size_ = subgroup_size; } + + /// @returns the subgroup size for the functions. + std::optional<uint32_t> SubgroupSize() const { return subgroup_size_; } + /// @returns all directly referenced global variables const UniqueVector<const GlobalVariable*, 4>& DirectlyReferencedGlobals() const { return directly_referenced_globals_; @@ -225,6 +232,7 @@ const ast::Function* const declaration_; sem::WorkgroupSize workgroup_size_; + std::optional<uint32_t> subgroup_size_; UniqueVector<const GlobalVariable*, 4> directly_referenced_globals_; UniqueVector<const GlobalVariable*, 8> transitively_referenced_globals_; UniqueVector<const Function*, 8> transitively_called_functions_;
diff --git a/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program.cc b/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program.cc index 704fde9..95ae411 100644 --- a/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program.cc +++ b/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program.cc
@@ -308,6 +308,10 @@ auto wgsize = fn->WorkgroupSize().value(); attrs.Push(b.Stage(ast::PipelineStage::kCompute)); attrs.Push(b.WorkgroupSize(Expr(wgsize[0]), Expr(wgsize[1]), Expr(wgsize[2]))); + if (fn->SubgroupSize().has_value()) { + auto sgsize = fn->SubgroupSize().value(); + attrs.Push(b.SubgroupSize(Expr(sgsize))); + } break; } case core::ir::Function::PipelineStage::kFragment: