[D3D12] Implement `chromium_experimental_subgroup_size_control`

This patch implements `chromium_experimental_subgroup_size_control`
on D3D12 backend. This feature requires at least Shader Model 6.6
for the HLSL attribute `[WaveSize]`.

As D3D12 runtime only accept the values between
`D3D12_FEATURE_DATA_D3D12_OPTIONS1::waveLaneCountMin` and
`D3D12_FEATURE_DATA_D3D12_OPTIONS1::waveLaneCountMax` as `[WaveSize]`,
while currently on D3D12 backend we don't always use these two values
as `AdapterInfo::subgroupMinSize` or `AdapterInfo::subgroupMaxSize`,
we introduce two new limitations (`minExplicitComputeSubgroupSize` and
`maxExplicitComputeSubgroupSize`) to provide valid values for
`[WaveSize]`.

Bug: 463721943
Test: dawn_end2end_tests
Change-Id: Ib86a6e2034957b5e1ab5091203d862d981247363
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/284855
Reviewed-by: James Price <jrprice@google.com>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Shao, Jiawei <jiawei.shao@intel.com>
diff --git a/docs/tint/extensions/chromium_experimental_subgroup_size_control.md b/docs/tint/extensions/chromium_experimental_subgroup_size_control.md
index e73a904..b3d75de 100644
--- a/docs/tint/extensions/chromium_experimental_subgroup_size_control.md
+++ b/docs/tint/extensions/chromium_experimental_subgroup_size_control.md
@@ -37,8 +37,8 @@
 ```
 - The parameter must be a const-expression or an override-expression that resolves to an `i32` or `u32`.
 - The parameter must be must be a power-of-two.
-- The parameter must be greater than or equal to the `subgroupMinSize` on the current `adapter`.
-- The parameter must be less than or equal to the `subgroupMaxSize` on the current `adapter`. We need to fix the D3D12 backend to report correct `subgroupMaxSize` on current `adapter`.
+- The parameter must be greater than or equal to the `minExplicitComputeSubgroupSize` (described later) on the current `adapter`.
+- The parameter must be less than or equal to the `maxExplicitComputeSubgroupSize` (described later) on the current `adapter`.
 - The total compute invocations per workgroup must be a multiple of the declared `subgroup_size`.
 
 # Example usage
@@ -55,7 +55,11 @@
 }
 ```
 
-This extension also adds a new limit `maxComputeWorkgroupSubgroups`, which is required by Vulkan `VK_EXT_subgroup_size_control` and defined in `VkPhysicalDeviceSubgroupSizeControlProperties.maxComputeWorkgroupSubgroups`. Note that there is no such limitation on D3D12.
+This extension also adds below limits in Dawn:
+
+1. `maxComputeWorkgroupSubgroups`
+
+Required by Vulkan `VK_EXT_subgroup_size_control` and defined in `VkPhysicalDeviceSubgroupSizeControlProperties.maxComputeWorkgroupSubgroups`. Note that there is no such limitation on D3D12.
 
 | Limit name | Type | Limit class | Default | Compatibility Mode Default |
 |------------|------|-------------|---------|----------------------------|
@@ -66,6 +70,20 @@
 workgroup_size.x * workgroup_size.y * workgroup_size.z <= subgroup_size * maxComputeWorkgroupSubgroups
 ```
 
+2. `minExplicitComputeSubgroupSize` and `maxExplicitComputeSubgroupSize` (in structure `wgpu::AdapterPropertiesExplicitComputeSubgroupSizeConfigs` that can be chained in `wgpu::AdapterInfo`)
+
+Required by both Vulkan and D3D12:
+- On D3D12 we should use `D3D12_FEATURE_DATA_D3D12_OPTIONS1::waveLaneCountMin` and `D3D12_FEATURE_DATA_D3D12_OPTIONS1::waveLaneCountMax`, which are not always used as `wgpu::AdapterInfo::subgroupMinSize` or `wgpu::AdapterInfo::subgroupMaxSize` in Dawn.
+  For example:
+  - On some Intel GPUs, it is possible to run some pixel shaders with wave lane count 8, while on that platform `waveLaneCountMin` is 16, meaning in compute shaders the wave lane count will always be at least 16 (A toggle ["d3d12_relax_min_subgroup_size_to_8"](https://issues.chromium.org/issues/381969450) has been added for this issue).
+  - Now Dawn always uses `128` as `SubgroupMaxSize` because in D3D12 document ["the WaveLaneCountMax queried from D3D12 API is not reliable and the meaning is unclear](https://github.com/Microsoft/DirectXShaderCompiler/wiki/Wave-Intrinsics#:~:text=UINT%20WaveLaneCountMax), while `waveLaneCountMax` is actually the maximum value that can be used as HLSL attribute `[WaveSize]`.
+- On Vulkan we should use `VkPhysicalDeviceSubgroupSizeControlPropertiesEXT::minSubgroupSize` and `VkPhysicalDeviceSubgroupSizeControlPropertiesEXT::maxSubgroupSize`.
+
+| Limit name | Type | Limit class | Default | Compatibility Mode Default |
+|------------|------|-------------|---------|----------------------------|
+| `minExplicitComputeSubgroupSize` | `GPUSize32` | `maximum` | - | Not Supported |
+| `maxExplicitComputeSubgroupSize` | `GPUSize32` | `maximum` | - | Not Supported |
+
 # References
 
 * [DirectX Specs HLSL Wave Size](https://microsoft.github.io/DirectX-Specs/d3d/HLSL_SM_6_6_WaveSize.html#allowed-wave-sizes)
diff --git a/src/dawn/dawn.json b/src/dawn/dawn.json
index 33f4810..e3e3f7b 100644
--- a/src/dawn/dawn.json
+++ b/src/dawn/dawn.json
@@ -4022,7 +4022,8 @@
             {"value": 78, "name": "shared texture memory D3D12 resource descriptor", "tags": ["dawn", "native"]},
             {"value": 79, "name": "request adapter options angle virtualization group", "tags": ["dawn", "native"]},
             {"value": 80, "name": "resource table limits", "tags": ["dawn"]},
-            {"value": 81, "name": "pipeline layout resource table", "tags": ["dawn"]}
+            {"value": 81, "name": "pipeline layout resource table", "tags": ["dawn"]},
+            {"value": 82, "name": "adapter properties explicit compute subgroup size configs", "tags": ["dawn"]}
         ]
     },
     "texture": {
@@ -4739,5 +4740,15 @@
             {"name": "config count", "type": "size_t"},
             {"name": "configs", "type": "subgroup matrix config", "annotation": "const*", "length": "config count"}
         ]
+    },
+    "adapter properties explicit compute subgroup size configs": {
+        "category": "structure",
+        "chained": "out",
+        "chain roots": ["adapter info"],
+        "tags": ["dawn"],
+        "members": [
+            {"name": "min explicit compute subgroup size", "type": "uint32_t"},
+            {"name": "max explicit compute subgroup size", "type": "uint32_t"}
+        ]
     }
 }
diff --git a/src/dawn/native/Adapter.cpp b/src/dawn/native/Adapter.cpp
index 12485c0..19b2064 100644
--- a/src/dawn/native/Adapter.cpp
+++ b/src/dawn/native/Adapter.cpp
@@ -150,6 +150,11 @@
         hadError |= mInstance->ConsumedError(
             DAWN_VALIDATION_ERROR("Feature ChromiumExperimentalSubgroupMatrix is not available."));
     }
+    if (unpacked.Has<AdapterPropertiesExplicitComputeSubgroupSizeConfigs>() &&
+        !mSupportedFeatures.IsEnabled(wgpu::FeatureName::ChromiumExperimentalSubgroupSizeControl)) {
+        hadError |= mInstance->ConsumedError(DAWN_VALIDATION_ERROR(
+            "Feature ChromiumExperimentalExplicitComputeSubgroupSize is not available."));
+    }
     if (hadError) {
         return wgpu::Status::Error;
     }
@@ -159,6 +164,11 @@
     }
 
     mPhysicalDevice->PopulateBackendProperties(unpacked, mTogglesState);
+    if (auto* explicitSubgroupSizeConfigs =
+            unpacked.Get<AdapterPropertiesExplicitComputeSubgroupSizeConfigs>()) {
+        DAWN_ASSERT(IsPowerOfTwo(explicitSubgroupSizeConfigs->minExplicitComputeSubgroupSize));
+        DAWN_ASSERT(IsPowerOfTwo(explicitSubgroupSizeConfigs->maxExplicitComputeSubgroupSize));
+    }
 
     // Allocate space for all strings.
     size_t allocSize = mPhysicalDevice->GetVendorName().length() +
@@ -192,6 +202,9 @@
         info->subgroupMinSize = std::min(info->subgroupMinSize, 8u);
     }
 
+    DAWN_ASSERT(IsPowerOfTwo(info->subgroupMaxSize));
+    DAWN_ASSERT(IsPowerOfTwo(info->subgroupMinSize));
+
     return wgpu::Status::Success;
 }
 
diff --git a/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp b/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
index 27c38f9..123f520 100644
--- a/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
+++ b/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
@@ -196,6 +196,11 @@
     if (mDeviceInfo.supportsWaveOps) {
         EnableFeature(Feature::Subgroups);
     }
+
+    // SubgroupSizeControl feature requires SM >= 6.6 for HLSL attribute `[WaveSize]`.
+    if (mDeviceInfo.highestSupportedShaderModel >= 66) {
+        EnableFeature(Feature::ChromiumExperimentalSubgroupSizeControl);
+    }
 #endif
 
     D3D12_FEATURE_DATA_FORMAT_SUPPORT bgra8unormFormatInfo = {};
@@ -951,6 +956,13 @@
         // Report highest supported shader model version, instead of actual applied version.
         d3dProperties->shaderModel = GetDeviceInfo().highestSupportedShaderModel;
     }
+    if (auto* explicitComputeSubgroupSizeConfigs =
+            info.Get<AdapterPropertiesExplicitComputeSubgroupSizeConfigs>()) {
+        explicitComputeSubgroupSizeConfigs->minExplicitComputeSubgroupSize =
+            mDeviceInfo.waveLaneCountMin;
+        explicitComputeSubgroupSizeConfigs->maxExplicitComputeSubgroupSize =
+            mDeviceInfo.waveLaneCountMax;
+    }
 }
 
 }  // namespace dawn::native::d3d12
diff --git a/src/dawn/native/null/DeviceNull.cpp b/src/dawn/native/null/DeviceNull.cpp
index ad47151..ca229dc 100644
--- a/src/dawn/native/null/DeviceNull.cpp
+++ b/src/dawn/native/null/DeviceNull.cpp
@@ -128,6 +128,11 @@
     if (auto* d3dProperties = info.Get<AdapterPropertiesD3D>()) {
         d3dProperties->shaderModel = 0;
     }
+    if (auto* explicitComputeSubgroupSizeConfigs =
+            info.Get<AdapterPropertiesExplicitComputeSubgroupSizeConfigs>()) {
+        explicitComputeSubgroupSizeConfigs->minExplicitComputeSubgroupSize = 4;
+        explicitComputeSubgroupSizeConfigs->maxExplicitComputeSubgroupSize = 128;
+    }
 }
 
 FeatureValidationResult PhysicalDevice::ValidateFeatureSupportedWithTogglesImpl(
diff --git a/src/dawn/tests/end2end/SubgroupsTests.cpp b/src/dawn/tests/end2end/SubgroupsTests.cpp
index 95197ec..019c66d 100644
--- a/src/dawn/tests/end2end/SubgroupsTests.cpp
+++ b/src/dawn/tests/end2end/SubgroupsTests.cpp
@@ -930,5 +930,97 @@
                             SubgroupOpDataType::I32,
                         });
 
+class SubgroupSizeControlTests : public DawnTest {
+  protected:
+    std::vector<wgpu::FeatureName> GetRequiredFeatures() override {
+        // Always require related features if available.
+        std::vector<wgpu::FeatureName> requiredFeatures;
+        if (SupportsFeatures({wgpu::FeatureName::Subgroups,
+                              wgpu::FeatureName::ChromiumExperimentalSubgroupSizeControl})) {
+            mSupportsSubgroupSizeControl = true;
+            requiredFeatures.push_back(wgpu::FeatureName::Subgroups);
+            requiredFeatures.push_back(wgpu::FeatureName::ChromiumExperimentalSubgroupSizeControl);
+        }
+        return requiredFeatures;
+    }
+
+    bool SupportSubgroupSizeControl() const { return mSupportsSubgroupSizeControl; }
+
+    void DoTest(uint32_t subgroupSize) {
+        DAWN_ASSERT(IsPowerOfTwo(subgroupSize));
+
+        std::stringstream code;
+        code << R"(
+enable subgroups;
+enable chromium_experimental_subgroup_size_control;
+
+override kSubgroupSize : u32;
+
+@group(0) @binding(0)
+var<storage, read_write> output: u32;
+
+@compute @workgroup_size(kSubgroupSize) @subgroup_size(kSubgroupSize)
+fn main(@builtin(subgroup_size) sg_size : u32) {
+    if (subgroupElect()) {
+        output = sg_size;
+    }
+})";
+        wgpu::ShaderModule csModule = utils::CreateShaderModule(device, code.str().c_str());
+
+        wgpu::ConstantEntry entry = {nullptr, "kSubgroupSize", static_cast<double>(subgroupSize)};
+        wgpu::ComputePipelineDescriptor csDesc;
+        csDesc.compute.module = csModule;
+        csDesc.compute.constantCount = 1;
+        csDesc.compute.constants = &entry;
+        auto pipeline = device.CreateComputePipeline(&csDesc);
+
+        uint32_t outputBufferSizeInBytes = sizeof(uint32_t);
+        wgpu::BufferDescriptor outputBufferDesc;
+        outputBufferDesc.size = outputBufferSizeInBytes;
+        outputBufferDesc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc;
+        wgpu::Buffer outputBuffer = device.CreateBuffer(&outputBufferDesc);
+
+        wgpu::BindGroup bindGroup = utils::MakeBindGroup(device, pipeline.GetBindGroupLayout(0),
+                                                         {
+                                                             {0, outputBuffer},
+                                                         });
+
+        wgpu::CommandEncoder encoder = device.CreateCommandEncoder();
+        wgpu::ComputePassEncoder pass = encoder.BeginComputePass();
+        pass.SetPipeline(pipeline);
+        pass.SetBindGroup(0, bindGroup);
+        pass.DispatchWorkgroups(1);
+        pass.End();
+        wgpu::CommandBuffer commands = encoder.Finish();
+        queue.Submit(1, &commands);
+
+        EXPECT_BUFFER_U32_EQ(subgroupSize, outputBuffer, 0);
+    }
+
+  private:
+    bool mSupportsSubgroupSizeControl = false;
+};
+
+// Test all the values that are between `minExplicitComputeSubgroupSize` and
+// `maxExplicitComputeSubgroupSize` and are a power of 2 can be used as WGSL attribute
+// `@subgroup_size` and the value of the WGSL builtin `subgroup_size` exactly matches the value of
+// the WGSL attribute `@subgroup_size`.
+TEST_P(SubgroupSizeControlTests, TestAllSubgroupSizes) {
+    DAWN_TEST_UNSUPPORTED_IF(!SupportSubgroupSizeControl());
+
+    wgpu::AdapterInfo info;
+    wgpu::AdapterPropertiesExplicitComputeSubgroupSizeConfigs subgroupSizeConfigs;
+    info.nextInChain = &subgroupSizeConfigs;
+    adapter.GetInfo(&info);
+
+    ASSERT_TRUE(IsPowerOfTwo(subgroupSizeConfigs.minExplicitComputeSubgroupSize));
+    for (uint32_t subgroupSize = subgroupSizeConfigs.minExplicitComputeSubgroupSize;
+         subgroupSize <= subgroupSizeConfigs.maxExplicitComputeSubgroupSize; subgroupSize *= 2) {
+        DoTest(subgroupSize);
+    }
+}
+
+DAWN_INSTANTIATE_TEST(SubgroupSizeControlTests, D3D12Backend(), MetalBackend(), VulkanBackend());
+
 }  // anonymous namespace
 }  // namespace dawn
diff --git a/src/dawn/tests/unittests/wire/WireInstanceTests.cpp b/src/dawn/tests/unittests/wire/WireInstanceTests.cpp
index a18211d..743377b 100644
--- a/src/dawn/tests/unittests/wire/WireInstanceTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireInstanceTests.cpp
@@ -251,11 +251,19 @@
     fakePowerProperties.chain.sType = WGPUSType_DawnAdapterPropertiesPowerPreference;
     fakePowerProperties.powerPreference = WGPUPowerPreference::WGPUPowerPreference_LowPower;
 
+    WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs fakeExplicitComputeSubgroupSizeConfigs =
+        {};
+    fakeExplicitComputeSubgroupSizeConfigs.chain.sType =
+        WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs;
+    fakeExplicitComputeSubgroupSizeConfigs.minExplicitComputeSubgroupSize = 8;
+    fakeExplicitComputeSubgroupSizeConfigs.maxExplicitComputeSubgroupSize = 32;
+
     std::initializer_list<WGPUFeatureName> fakeFeaturesList = {
         WGPUFeatureName_AdapterPropertiesMemoryHeaps,
         WGPUFeatureName_AdapterPropertiesD3D,
         WGPUFeatureName_AdapterPropertiesVk,
         WGPUFeatureName_ChromiumExperimentalSubgroupMatrix,
+        WGPUFeatureName_ChromiumExperimentalSubgroupSizeControl,
     };
     WGPUSupportedFeatures fakeFeatures = {fakeFeaturesList.size(), std::data(fakeFeaturesList)};
 
@@ -297,6 +305,11 @@
                                 *reinterpret_cast<WGPUDawnAdapterPropertiesPowerPreference*>(
                                     chain) = fakePowerProperties;
                                 break;
+                            case WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs:
+                                *reinterpret_cast<
+                                    WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs*>(
+                                    chain) = fakeExplicitComputeSubgroupSizeConfigs;
+                                break;
                             default:
                                 ADD_FAILURE() << "Unexpected chain";
                                 return WGPUStatus_Error;
@@ -387,6 +400,18 @@
                 adapter.GetInfo(reinterpret_cast<wgpu::AdapterInfo*>(&info));
                 // Expect them to match.
                 EXPECT_EQ(powerProperties.powerPreference, fakePowerProperties.powerPreference);
+
+                // Get the explicit compute subgroup size properties
+                WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs subgroupSizeConfigs = {};
+                subgroupSizeConfigs.chain.sType =
+                    WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs;
+                info.nextInChain = &subgroupSizeConfigs.chain;
+                adapter.GetInfo(reinterpret_cast<wgpu::AdapterInfo*>(&info));
+                // Expect them to match
+                EXPECT_EQ(subgroupSizeConfigs.minExplicitComputeSubgroupSize,
+                          fakeExplicitComputeSubgroupSizeConfigs.minExplicitComputeSubgroupSize);
+                EXPECT_EQ(subgroupSizeConfigs.maxExplicitComputeSubgroupSize,
+                          fakeExplicitComputeSubgroupSizeConfigs.maxExplicitComputeSubgroupSize);
             }));
 
         FlushCallbacks();
diff --git a/src/dawn/wire/client/Adapter.cpp b/src/dawn/wire/client/Adapter.cpp
index cbdb0177..2f10046 100644
--- a/src/dawn/wire/client/Adapter.cpp
+++ b/src/dawn/wire/client/Adapter.cpp
@@ -202,6 +202,16 @@
                 mPowerProperties.powerPreference = powerProperties->powerPreference;
                 break;
             }
+            case WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs: {
+                auto* subgroupSizeConfigs =
+                    reinterpret_cast<WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs*>(
+                        chain);
+                mExplicitComputeSubgroupSizeConfigs.minExplicitComputeSubgroupSize =
+                    subgroupSizeConfigs->minExplicitComputeSubgroupSize;
+                mExplicitComputeSubgroupSizeConfigs.maxExplicitComputeSubgroupSize =
+                    subgroupSizeConfigs->maxExplicitComputeSubgroupSize;
+                break;
+            }
             default:
                 DAWN_UNREACHABLE();
                 break;
@@ -260,6 +270,20 @@
                 powerProperties->powerPreference = mPowerProperties.powerPreference;
                 break;
             }
+            case WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs: {
+                if (!APIHasFeature(WGPUFeatureName_ChromiumExperimentalSubgroupSizeControl)) {
+                    return WGPUStatus_Error;
+                }
+                auto* explicitComputeSubgroupSizeConfigs =
+                    reinterpret_cast<WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs*>(
+                        chain);
+                explicitComputeSubgroupSizeConfigs->minExplicitComputeSubgroupSize =
+                    mExplicitComputeSubgroupSizeConfigs.minExplicitComputeSubgroupSize;
+                explicitComputeSubgroupSizeConfigs->maxExplicitComputeSubgroupSize =
+                    mExplicitComputeSubgroupSizeConfigs.maxExplicitComputeSubgroupSize;
+                break;
+            }
+
             default:
                 break;
         }
diff --git a/src/dawn/wire/client/Adapter.h b/src/dawn/wire/client/Adapter.h
index 391da45..000c70d 100644
--- a/src/dawn/wire/client/Adapter.h
+++ b/src/dawn/wire/client/Adapter.h
@@ -80,6 +80,7 @@
     WGPUAdapterPropertiesVk mVkProperties;
     std::vector<WGPUSubgroupMatrixConfig> mSubgroupMatrixConfigs;
     WGPUDawnAdapterPropertiesPowerPreference mPowerProperties;
+    WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs mExplicitComputeSubgroupSizeConfigs;
 };
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/server/ServerInstance.cpp b/src/dawn/wire/server/ServerInstance.cpp
index bae3ab2..ac3f29a 100644
--- a/src/dawn/wire/server/ServerInstance.cpp
+++ b/src/dawn/wire/server/ServerInstance.cpp
@@ -127,6 +127,16 @@
     *propertiesChain = &powerProperties.chain;
     propertiesChain = &(*propertiesChain)->next;
 
+    // Query AdapterPropertiesExplicitComputeSubgroupSizeConfigs if the feature is supported.
+    WGPUAdapterPropertiesExplicitComputeSubgroupSizeConfigs explicitComputeSubgroupSizeConfigs = {};
+    explicitComputeSubgroupSizeConfigs.chain.sType =
+        WGPUSType_AdapterPropertiesExplicitComputeSubgroupSizeConfigs;
+    if (mProcs.adapterHasFeature(adapter,
+                                 WGPUFeatureName_ChromiumExperimentalSubgroupSizeControl)) {
+        *propertiesChain = &explicitComputeSubgroupSizeConfigs.chain;
+        propertiesChain = &(*propertiesChain)->next;
+    }
+
     mProcs.adapterGetInfo(adapter, &info);
     cmd.info = &info;