[metal][msl] Support subgroup_size_control on metal * Trivially support the extension when min subgroup size == max subgroup size * Validate compile Fix: 529959695 Change-Id: I08101012e218b45101a7c6f61cfd43d8bd63ce0c Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/321775 Reviewed-by: dan sinclair <dsinclair@chromium.org> Reviewed-by: Corentin Wallez <cwallez@chromium.org> Commit-Queue: Alan Baker <alanbaker@google.com>

diff --git a/src/dawn/native/metal/PhysicalDeviceMTL.mm b/src/dawn/native/metal/PhysicalDeviceMTL.mm
index 9b5279c5..c236051 100644
--- a/src/dawn/native/metal/PhysicalDeviceMTL.mm
+++ b/src/dawn/native/metal/PhysicalDeviceMTL.mm

@@ -735,6 +735,11 @@
     if (([*mDevice supportsFamily:MTLGPUFamilyApple6] ||
          [*mDevice supportsFamily:MTLGPUFamilyMac2])) {
         EnableFeature(Feature::Subgroups);
+        // Apple doesn't support selecting a subgroup size, but if there is only one possible size
+        // we can trivially enable the feature.
+        if (mSubgroupMinSize == mSubgroupMaxSize) {
+            EnableFeature(Feature::SubgroupSizeControl);
+        }
     }
 
     if ([*mDevice supportsFamily:MTLGPUFamilyApple7]) {

diff --git a/src/dawn/native/metal/ShaderModuleMTL.mm b/src/dawn/native/metal/ShaderModuleMTL.mm
index a8f09b9..4b4967b 100644
--- a/src/dawn/native/metal/ShaderModuleMTL.mm
+++ b/src/dawn/native/metal/ShaderModuleMTL.mm

@@ -63,6 +63,7 @@
     X(UnsafeUnserializedValue<ShaderModuleBase::ScopedUseTintProgram>, inputProgram) \
     X(LimitsForCompilationRequest, limits)                                           \
     X(UnsafeUnserializedValue<LimitsForCompilationRequest>, adapterSupportedLimits)  \
+    X(uint32_t, minSubgroupSize)                                                     \
     X(uint32_t, maxSubgroupSize)                                                     \
     X(bool, usesSubgroupMatrix)                                                      \
     X(bool, useStrictMath)                                                           \
@@ -369,6 +370,7 @@
     req.limits = LimitsForCompilationRequest::Create(device->GetLimits().v1);
     req.adapterSupportedLimits = UnsafeUnserializedValue(
         LimitsForCompilationRequest::Create(device->GetAdapter()->GetLimits().v1));
+    req.minSubgroupSize = device->GetAdapter()->GetPhysicalDevice()->GetSubgroupMinSize();
     req.maxSubgroupSize = device->GetAdapter()->GetPhysicalDevice()->GetSubgroupMaxSize();
 
     CacheResult<MslCompilation> mslCompilation;
@@ -418,6 +420,15 @@
                                 ValidateComputeStageWorkgroupSize(
                                     result->workgroup_info, r.usesSubgroupMatrix, r.maxSubgroupSize,
                                     r.limits, r.adapterSupportedLimits.UnsafeGetValue()));
+
+                if (result->workgroup_info.subgroup_size.has_value()) {
+                    uint32_t explicitSubgroupSize = result->workgroup_info.subgroup_size.value();
+                    DAWN_INVALID_IF(
+                        explicitSubgroupSize < r.minSubgroupSize ||
+                            explicitSubgroupSize > r.maxSubgroupSize,
+                        "The subgroup_size attribute (%u) is not in the allowed range ([%u, %u]).",
+                        explicitSubgroupSize, r.minSubgroupSize, r.maxSubgroupSize);
+                }
             }
 
             auto msl = std::move(result->msl);

diff --git a/src/dawn/tests/end2end/ArchTierLimitsExhaustive.cpp b/src/dawn/tests/end2end/ArchTierLimitsExhaustive.cpp
index 5ca6fc9..8bb5c63 100644
--- a/src/dawn/tests/end2end/ArchTierLimitsExhaustive.cpp
+++ b/src/dawn/tests/end2end/ArchTierLimitsExhaustive.cpp

@@ -449,7 +449,7 @@
 AddDevice({1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1}, "Vulkan_Radeon_RX_5500_XT");
 
 // Apple
-AddDevice({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1}, "Metal_Apple_M2");
+AddDevice({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1}, "Metal_Apple_M2");
 
 // ARM
 

diff --git a/src/tint/lang/msl/writer/printer/printer.cc b/src/tint/lang/msl/writer/printer/printer.cc
index f027707..2faddee 100644
--- a/src/tint/lang/msl/writer/printer/printer.cc
+++ b/src/tint/lang/msl/writer/printer/printer.cc

@@ -375,6 +375,13 @@
                     result_.workgroup_info.y = wg_size[1];
                     result_.workgroup_info.z = wg_size[2];
 
+                    // Store the subgroup size information away to return from the generator when
+                    // the `@subgroup_size` attribute is used.
+                    const auto const_sg_size = func->SubgroupSizeAsConst();
+                    if (const_sg_size.has_value()) {
+                        result_.workgroup_info.subgroup_size = const_sg_size;
+                    }
+
                     break;
                 }
                 case core::ir::Function::PipelineStage::kFragment:

diff --git a/src/tint/lang/msl/writer/writer.cc b/src/tint/lang/msl/writer/writer.cc
index 89a826c..029450a 100644
--- a/src/tint/lang/msl/writer/writer.cc
+++ b/src/tint/lang/msl/writer/writer.cc

@@ -102,11 +102,6 @@
             continue;
         }
 
-        // Check `@subgroup_size` attribute.
-        if (f->SubgroupSize().has_value()) {
-            return Failure("subgroup_size attribute is not supported by the MSL backend");
-        }
-
         // Check input attributes.
         for (auto* param : f->Params()) {
             if (auto* str = param->Type()->As<core::type::Struct>()) {