Vulkan: Use Zero Initialize Workgroup Memory extension if possible

This patch adds a toggle to enable/disable workgroup memory
initialization with OpConstantNull according to the Vulkan extension
VK_KHR_zero_initialize_workgroup_memory. This toggle is by default
enabled when VK_KHR_zero_initialize_workgroup_memory is supported by
the Vulkan driver.

BUG=dawn:1302
TEST=dawn_end2end_tests

Change-Id: Ie04484c2d0944ead082bd22a436b1c52bc7d93bb
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/82400
Reviewed-by: Austin Eng <enga@chromium.org>
Commit-Queue: Jiawei Shao <jiawei.shao@intel.com>

src/dawn/native/Toggles.cpp

@@ -241,6 +241,11 @@ namespace dawn::native {
              {"disable_timestamp_query_conversion",
               "Resolve timestamp queries into ticks instead of nanoseconds.",
               "https://crbug.com/dawn/1305"}},
+            {Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension,
+             {"use_vulkan_zero_initialize_workgroup_memory_extension",
+              "Initialize workgroup memory with OpConstantNull on Vulkan when the Vulkan extension "
+              "VK_KHR_zero_initialize_workgroup_memory is supported.",
+              "https://crbug.com/dawn/1302"}},
 
             // Dummy comment to separate the }} so it is clearer what to copy-paste to add a toggle.
         }};

src/dawn/native/Toggles.h

@@ -63,6 +63,7 @@ namespace dawn::native {
         FxcOptimizations,
         RecordDetailedTimingInTraceEvents,
         DisableTimestampQueryConversion,
+        VulkanUseZeroInitializeWorkgroupMemoryExtension,
 
         EnumCount,
         InvalidEnum = EnumCount,

src/dawn/native/vulkan/DeviceVk.cpp

@@ -97,6 +97,10 @@ namespace dawn::native::vulkan {
         // the decision if it is not applicable.
         ApplyDepth24PlusS8Toggle();
 
+        // The environment can only request to use VK_KHR_zero_initialize_workgroup_memory when the
+        // extension is available. Override the decision if it is no applicable.
+        ApplyUseZeroInitializeWorkgroupMemoryExtensionToggle();
+
         return DeviceBase::Initialize(Queue::Create(this));
     }
 
@@ -345,6 +349,20 @@ namespace dawn::native::vulkan {
             mComputeSubgroupSize = FindComputeSubgroupSize();
         }
 
+        if (mDeviceInfo.HasExt(DeviceExt::ZeroInitializeWorkgroupMemory)) {
+            ASSERT(usedKnobs.HasExt(DeviceExt::ZeroInitializeWorkgroupMemory));
+
+            usedKnobs.zeroInitializeWorkgroupMemoryFeatures.sType =
+                VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ZERO_INITIALIZE_WORKGROUP_MEMORY_FEATURES_KHR;
+
+            // Always allow initializing workgroup memory with OpConstantNull when available.
+            // Note that the driver still won't initialize workgroup memory unless the workgroup
+            // variable is explicitly initialized with OpConstantNull.
+            usedKnobs.zeroInitializeWorkgroupMemoryFeatures.shaderZeroInitializeWorkgroupMemory =
+                VK_TRUE;
+            featuresChain.Add(&usedKnobs.zeroInitializeWorkgroupMemoryFeatures);
+        }
+
         if (mDeviceInfo.features.samplerAnisotropy == VK_TRUE) {
             usedKnobs.features.samplerAnisotropy = VK_TRUE;
         }
@@ -488,6 +506,7 @@ namespace dawn::native::vulkan {
         fn.GetDeviceQueue(mVkDevice, mQueueFamily, 0, &mQueue);
     }
 
+    // Note that this function is called before mDeviceInfo is initialized.
     void Device::InitTogglesFromDriver() {
         // TODO(crbug.com/dawn/857): tighten this workaround when this issue is fixed in both
         // Vulkan SPEC and drivers.
@@ -495,6 +514,10 @@ namespace dawn::native::vulkan {
 
         // By default try to use D32S8 for Depth24PlusStencil8
         SetToggle(Toggle::VulkanUseD32S8, true);
+
+        // By default try to initialize workgroup memory with OpConstantNull according to the Vulkan
+        // extension VK_KHR_zero_initialize_workgroup_memory.
+        SetToggle(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension, true);
     }
 
     void Device::ApplyDepth24PlusS8Toggle() {
@@ -513,6 +536,12 @@ namespace dawn::native::vulkan {
         }
     }
 
+    void Device::ApplyUseZeroInitializeWorkgroupMemoryExtensionToggle() {
+        if (!mDeviceInfo.HasExt(DeviceExt::ZeroInitializeWorkgroupMemory)) {
+            ForceSetToggle(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension, false);
+        }
+    }
+
     VulkanFunctions* Device::GetMutableFunctions() {
         return const_cast<VulkanFunctions*>(&fn);
     }

src/dawn/native/vulkan/DeviceVk.h

@@ -152,6 +152,7 @@ namespace dawn::native::vulkan {
         uint32_t FindComputeSubgroupSize() const;
         void InitTogglesFromDriver();
         void ApplyDepth24PlusS8Toggle();
+        void ApplyUseZeroInitializeWorkgroupMemoryExtensionToggle();
 
         void DestroyImpl() override;
         MaybeError WaitForIdleForDestruction() override;

src/dawn/native/vulkan/ShaderModuleVk.cpp

@@ -204,6 +204,8 @@ namespace dawn::native::vulkan {
         tint::writer::spirv::Options options;
         options.emit_vertex_point_size = true;
         options.disable_workgroup_init = GetDevice()->IsToggleEnabled(Toggle::DisableWorkgroupInit);
+        options.use_zero_initialize_workgroup_memory_extension =
+            GetDevice()->IsToggleEnabled(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension);
 
         std::vector<uint32_t> spirv;
         {

src/dawn/native/vulkan/VulkanExtensions.cpp

@@ -24,6 +24,7 @@ namespace dawn::native::vulkan {
 
     static constexpr uint32_t VulkanVersion_1_1 = VK_MAKE_VERSION(1, 1, 0);
     static constexpr uint32_t VulkanVersion_1_2 = VK_MAKE_VERSION(1, 2, 0);
+    static constexpr uint32_t VulkanVersion_1_3 = VK_MAKE_VERSION(1, 3, 0);
     static constexpr uint32_t NeverPromoted = std::numeric_limits<uint32_t>::max();
 
     // A static array for InstanceExtInfo that can be indexed with InstanceExts.
@@ -149,6 +150,9 @@ namespace dawn::native::vulkan {
         {DeviceExt::ImageFormatList, "VK_KHR_image_format_list", VulkanVersion_1_2},
         {DeviceExt::ShaderFloat16Int8, "VK_KHR_shader_float16_int8", VulkanVersion_1_2},
 
+        {DeviceExt::ZeroInitializeWorkgroupMemory, "VK_KHR_zero_initialize_workgroup_memory",
+         VulkanVersion_1_3},
+
         {DeviceExt::ExternalMemoryFD, "VK_KHR_external_memory_fd", NeverPromoted},
         {DeviceExt::ExternalMemoryDmaBuf, "VK_EXT_external_memory_dma_buf", NeverPromoted},
         {DeviceExt::ExternalMemoryZirconHandle, "VK_FUCHSIA_external_memory", NeverPromoted},
@@ -276,6 +280,10 @@ namespace dawn::native::vulkan {
                     hasDependencies = icdVersion >= VulkanVersion_1_1;
                     break;
 
+                case DeviceExt::ZeroInitializeWorkgroupMemory:
+                    hasDependencies = HasDep(DeviceExt::GetPhysicalDeviceProperties2);
+                    break;
+
                 case DeviceExt::EnumCount:
                     UNREACHABLE();
             }

src/dawn/native/vulkan/VulkanExtensions.h

@@ -89,6 +89,9 @@ namespace dawn::native::vulkan {
         ImageFormatList,
         ShaderFloat16Int8,
 
+        // Promoted to 1.3
+        ZeroInitializeWorkgroupMemory,
+
         // External* extensions
         ExternalMemoryFD,
         ExternalMemoryDmaBuf,

src/dawn/native/vulkan/VulkanInfo.h

@@ -51,6 +51,8 @@ namespace dawn::native::vulkan {
         VkPhysicalDeviceShaderFloat16Int8FeaturesKHR shaderFloat16Int8Features;
         VkPhysicalDevice16BitStorageFeaturesKHR _16BitStorageFeatures;
         VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroupSizeControlFeatures;
+        VkPhysicalDeviceZeroInitializeWorkgroupMemoryFeaturesKHR
+            zeroInitializeWorkgroupMemoryFeatures;
 
         bool HasExt(DeviceExt ext) const;
         DeviceExtSet extensions;

src/dawn/tests/end2end/ComputeSharedMemoryTests.cpp

@@ -201,4 +201,5 @@ DAWN_INSTANTIATE_TEST(ComputeSharedMemoryTests,
                       MetalBackend(),
                       OpenGLBackend(),
                       OpenGLESBackend(),
-                      VulkanBackend());
+                      VulkanBackend(),
+                      VulkanBackend({}, {"use_vulkan_zero_initialize_workgroup_memory_extension"}));