[ir] Validate attribute subgroup_size in IR validator

This patch adds the attribute `subgroup_size` and all the related
validation in Tint IR validator.
- `subgroup_size` must be a constant or override expression
- `subgroup_size` must be `i32` or `u32`.
- `subgroup_size` must be greater than 0.
- `subgroup_size` must be a power of 2.

Bug: 463721943
Change-Id: I9de573e6253de3f66d0ee4809674d5f49881671a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/280415
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: Shao, Jiawei <jiawei.shao@intel.com>

Author: Jiawei-Shao

src/tint/cmd/bench/enums_core_bench.cc

@@ -249,6 +249,9 @@ enum attribute {
 
   // chromium_internal_input_attachments
   input_attachment_index
+
+  // chromium_experimental_subgroup_size_control
+  subgroup_size
 }
 
 // These are parameter usages which show up in other def files but not in core.def.

src/tint/lang/core/core.def

@@ -1163,6 +1163,9 @@ Attribute ParseAttribute(std::string_view str) {
     if (str == "size") {
         return Attribute::kSize;
     }
+    if (str == "subgroup_size") {
+        return Attribute::kSubgroupSize;
+    }
     if (str == "vertex") {
         return Attribute::kVertex;
     }
@@ -1207,6 +1210,8 @@ std::string_view ToString(Attribute value) {
             return "must_use";
         case Attribute::kSize:
             return "size";
+        case Attribute::kSubgroupSize:
+            return "subgroup_size";
         case Attribute::kVertex:
             return "vertex";
         case Attribute::kWorkgroupSize:

src/tint/lang/core/enums.cc

@@ -687,6 +687,7 @@ enum class Attribute : uint8_t {
     kLocation,
     kMustUse,
     kSize,
+    kSubgroupSize,
     kVertex,
     kWorkgroupSize,
 };
@@ -726,6 +727,7 @@ constexpr std::string_view kAttributeStrings[] = {
     "location",
     "must_use",
     "size",
+    "subgroup_size",
     "vertex",
     "workgroup_size",
 };

src/tint/lang/core/enums.h

@@ -1083,6 +1083,7 @@ static constexpr AttributeCase kValidAttributeCases[] = {
     {"location", Attribute::kLocation},
     {"must_use", Attribute::kMustUse},
     {"size", Attribute::kSize},
+    {"subgroup_size", Attribute::kSubgroupSize},
     {"vertex", Attribute::kVertex},
     {"workgroup_size", Attribute::kWorkgroupSize},
 };
@@ -1136,12 +1137,15 @@ static constexpr AttributeCase kInvalidAttributeCases[] = {
     {"dzIp", Attribute::kUndefined},
     {"ize", Attribute::kUndefined},
     {"LN", Attribute::kUndefined},
-    {"r", Attribute::kUndefined},
-    {"vxxGteqqR", Attribute::kUndefined},
-    {"GGerteS", Attribute::kUndefined},
-    {"oqkccr8up_size", Attribute::kUndefined},
-    {"workgroup_sze", Attribute::kUndefined},
-    {"woppkgroup_sie", Attribute::kUndefined},
+    {"bgoupsiz", Attribute::kUndefined},
+    {"sxxbgrup_RRGizqq", Attribute::kUndefined},
+    {"GGubgrSup_size", Attribute::kUndefined},
+    {"88ccte", Attribute::kUndefined},
+    {"vrtex", Attribute::kUndefined},
+    {"perttx", Attribute::kUndefined},
+    {"workgr00up_siqFF55", Attribute::kUndefined},
+    {"workgroupsize", Attribute::kUndefined},
+    {"wokgroup_sie", Attribute::kUndefined},
 };
 
 using AttributeParseTest = testing::TestWithParam<AttributeCase>;

src/tint/lang/core/enums_test.cc

@@ -283,6 +283,11 @@ struct Decoder {
                                      Value(wg_size_in.z()));
         }
 
+        if (fn_in.has_subgroup_size()) {
+            uint32_t subgroup_size_in = fn_in.subgroup_size();
+            fn_out->SetSubgroupSize(Value(subgroup_size_in));
+        }
+
         Vector<FunctionParam*, 8> params_out;
         for (auto param_in : fn_in.parameters()) {
             auto* param_out = ValueAs<FunctionParam>(param_in);

src/tint/lang/core/ir/binary/decode.cc

@@ -150,6 +150,9 @@ struct Encoder {
             wg_size_out.set_y(Value((*wg_size_in)[1]));
             wg_size_out.set_z(Value((*wg_size_in)[2]));
         }
+        if (auto subgroup_size_in = fn_in->SubgroupSize()) {
+            fn_out->set_subgroup_size(Value(*subgroup_size_in));
+        }
         for (auto* param_in : fn_in->Params()) {
             fn_out->add_parameters(Value(param_in));
         }

src/tint/lang/core/ir/binary/encode.cc

@@ -343,6 +343,12 @@ void Disassembler::EmitFunction(const Function* func) {
         EmitValue(arr[2]);
         out_ << ")";
     }
+    if (func->SubgroupSize()) {
+        auto subgroup_size = func->SubgroupSize().value();
+        out_ << " " << StyleAttribute("@subgroup_size") << "(";
+        EmitValue(subgroup_size);
+        out_ << ")";
+    }
 
     out_ << " " << StyleKeyword("func") << "(";
 

src/tint/lang/core/ir/disassembler.cc

@@ -134,6 +134,26 @@ class Function : public Castable<Function, Value> {
         }};
     }
 
+    /// Sets the subgroup size
+    /// @param size the new subgroup size
+    void SetSubgroupSize(Value* size) { subgroup_size_ = size; }
+
+    /// @returns the subgroup size information
+    std::optional<Value*> SubgroupSize() const { return subgroup_size_; }
+
+    /// @returns the subgroup size information as `uint32_t` values. Note, this requires the value
+    /// to be constant.
+    std::optional<uint32_t> SubgroupSizeAsConst() const {
+        if (!subgroup_size_.has_value()) {
+            return std::nullopt;
+        }
+
+        auto* value = subgroup_size_.value()->As<core::ir::Constant>();
+        TINT_ASSERT(value);
+
+        return value->Value()->ValueAs<uint32_t>();
+    }
+
     /// @param type the return type for the function
     void SetReturnType(const core::type::Type* type) { return_.type = type; }
 
@@ -224,6 +244,7 @@ class Function : public Castable<Function, Value> {
   private:
     PipelineStage pipeline_stage_ = PipelineStage::kUndefined;
     std::optional<std::array<Value*, 3>> workgroup_size_;
+    std::optional<Value*> subgroup_size_;
 
     struct {
         const core::type::Type* type = nullptr;

src/tint/lang/core/ir/function.h

@@ -1205,6 +1205,10 @@ class Validator {
     /// @param func the function to validate
     void CheckWorkgroupSize(const Function* func);
 
+    /// Validates the subgroup_size attribute for a given function
+    /// @param func the function to validate
+    void CheckSubgroupSize(const Function* func);
+
     /// Validates the specific function as a vertex entry point
     /// @param ep the function to validate
     void CheckPositionPresentForVertexOutput(const Function* ep);
@@ -2700,6 +2704,8 @@ void Validator::CheckFunction(const Function* func) {
 
     CheckWorkgroupSize(func);
 
+    CheckSubgroupSize(func);
+
     if (func->Stage() == Function::PipelineStage::kCompute) {
         if (DAWN_UNLIKELY(func->ReturnType() && !func->ReturnType()->Is<core::type::Void>())) {
             AddError(func) << "compute entry point must not have a return type, found "
@@ -3099,6 +3105,69 @@ void Validator::CheckWorkgroupSize(const Function* func) {
     }
 }
 
+void Validator::CheckSubgroupSize(const Function* func) {
+    // @subgroup_size is optional
+    if (!func->SubgroupSize().has_value()) {
+        return;
+    }
+
+    if (!func->IsCompute()) {
+        AddError(func) << "@subgroup_size only valid on compute entry point";
+        return;
+    }
+
+    auto subgroup_size = func->SubgroupSize().value();
+    if (!subgroup_size->Type()) {
+        AddError(func) << "a @subgroup_size param is missing a type";
+        return;
+    }
+
+    auto* ty = subgroup_size->Type();
+    if (!ty->IsAnyOf<core::type::I32, core::type::U32>()) {
+        AddError(func) << "@subgroup_size param must be an 'i32' or 'u32', received " << NameOf(ty);
+        return;
+    }
+
+    if (auto* c = subgroup_size->As<ir::Constant>()) {
+        int64_t value = c->Value()->ValueAs<int64_t>();
+        if (value <= 0) {
+            AddError(func) << "@subgroup_size param must be greater than 0";
+            return;
+        }
+
+        if (!IsPowerOfTwo<int64_t>(value)) {
+            AddError(func) << "@subgroup_size param must be a power of 2";
+            return;
+        }
+
+        return;
+    }
+
+    if (!capabilities_.Contains(Capability::kAllowOverrides)) {
+        AddError(func) << "@subgroup_size param is not a constant value, and IR capability "
+                          "'kAllowOverrides' is not set";
+        return;
+    }
+
+    if (auto* r = subgroup_size->As<ir::InstructionResult>()) {
+        if (!r->Instruction()) {
+            AddError(func) << "instruction for @subgroup_size param is not defined";
+            return;
+        }
+
+        if (r->Instruction()->Block() != mod_.root_block) {
+            AddError(func) << "@subgroup_size param defined by non-module scope value";
+            return;
+        }
+
+        if (r->Instruction()->Is<core::ir::Override>()) {
+            return;
+        }
+    }
+
+    AddError(func) << "@subgroup_size must be an InstructionResult or a Constant";
+}
+
 void Validator::CheckPositionPresentForVertexOutput(const Function* ep) {
     if (IsPositionPresent(ep->ReturnAttributes(), ep->ReturnType())) {
         return;