[SPIRV] Implement WaveOpsIncludeHelperLanes (microsoft#7806)

This commit implements the `[WaveOpsIncludeHelperLanes]` attribute for
the SPIR-V backend.

When this attribute is applied to a shader entry point, the compiler
will:
- Add the `SPV_KHR_maximal_reconvergence` and `SPV_KHR_quad_control`
extensions.
- Add the `QuadControlKHR` capability.
- Add the `MaximallyReconvergesKHR` and `RequireFullQuadsKHR` execution
modes to the entry point.

This change also includes new tests to verify the correctness of the
implementation and updates the documentation.

Fixes microsoft#7334

---------

Co-authored-by: Nathan Gauër <[email protected]>

Author: s-perron

docs/SPIR-V.rst

@@ -285,6 +285,21 @@ Right now the following ``<builtin>`` are supported:
 Please see Vulkan spec. `15.9. Built-In Variables <https://registry.khronos.org/vulkan/specs/latest/html/vkspec.html#interfaces-builtin-variables>`_
 for detailed explanation of these builtins.
 
+Helper Lane Support
+~~~~~~~~~~~~~~~~~~~
+
+Shader Model 6.7 introduces the `[WaveOpsIncludeHelperLanes]` attribute. When this
+attribute is applied to a shader entry point, the SPIR-V backend will:
+
+1. Add the ``SPV_KHR_maximal_reconvergence`` and ``SPV_KHR_quad_control``
+   extensions to the module.
+2. Add the ``QuadControlKHR`` capability.
+3. Add the ``MaximallyReconvergesKHR`` and ``RequireFullQuadsKHR`` execution modes
+   to the entry point.
+
+This ensures that helper lanes are included in wave operations, which is the
+behavior required by the HLSL specification.
+
 Supported extensions
 ~~~~~~~~~~~~~~~~~~~~
 

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -862,6 +862,32 @@ void SpirvEmitter::HandleTranslationUnit(ASTContext &context) {
                                 SourceLocation());
   }
 
+  for (const FunctionInfo *entryInfo : workQueue) {
+    if (!entryInfo->isEntryFunction)
+      continue;
+
+    if (entryInfo->shaderModelKind != hlsl::ShaderModel::Kind::Pixel) {
+      continue;
+    }
+
+    const auto *funcDecl = entryInfo->funcDecl;
+    if (!funcDecl->hasAttr<HLSLWaveOpsIncludeHelperLanesAttr>())
+      continue;
+
+    spvBuilder.requireExtension("SPV_KHR_maximal_reconvergence",
+                                funcDecl->getLocation());
+    spvBuilder.requireExtension("SPV_KHR_quad_control",
+                                funcDecl->getLocation());
+    spvBuilder.requireCapability(spv::Capability::QuadControlKHR,
+                                 funcDecl->getLocation());
+    spvBuilder.addExecutionMode(entryInfo->entryFunction,
+                                spv::ExecutionMode::MaximallyReconvergesKHR, {},
+                                funcDecl->getLocation());
+    spvBuilder.addExecutionMode(entryInfo->entryFunction,
+                                spv::ExecutionMode::RequireFullQuadsKHR, {},
+                                funcDecl->getLocation());
+  }
+
   // For Vulkan 1.2 and later, add SignedZeroInfNanPreserve when -Gis is
   // provided to preserve NaN/Inf and signed zeros.
   if (spirvOptions.IEEEStrict) {

tools/clang/test/CodeGenSPIRV/wave-ops-include-helper-lanes-lib.hlsl

@@ -0,0 +1,35 @@
+// RUN: %dxc -T lib_6_7 -spirv %s | FileCheck %s
+
+// CHECK: OpCapability QuadControlKHR
+// CHECK-DAG: OpExtension "SPV_KHR_maximal_reconvergence"
+// CHECK-DAG: OpExtension "SPV_KHR_quad_control"
+
+// CHECK: OpEntryPoint Fragment %ps_main1 "ps_main1"
+// CHECK: OpEntryPoint Fragment %ps_main2 "ps_main2"
+// CHECK: OpEntryPoint Fragment %ps_main3 "ps_main3"
+
+// CHECK-DAG: OpExecutionMode %ps_main1 MaximallyReconvergesKHR
+// CHECK-DAG: OpExecutionMode %ps_main1 RequireFullQuadsKHR
+
+// CHECK-NOT: OpExecutionMode %ps_main2 MaximallyReconvergesKHR
+// CHECK-NOT: OpExecutionMode %ps_main2 RequireFullQuadsKHR
+
+// CHECK-DAG: OpExecutionMode %ps_main3 MaximallyReconvergesKHR
+// CHECK-DAG: OpExecutionMode %ps_main3 RequireFullQuadsKHR
+
+[WaveOpsIncludeHelperLanes]
+[shader("pixel")]
+void ps_main1() : SV_Target0
+{
+}
+
+[shader("pixel")]
+void ps_main2() : SV_Target0
+{
+}
+
+[WaveOpsIncludeHelperLanes]
+[shader("pixel")]
+void ps_main3() : SV_Target0
+{
+}

tools/clang/test/CodeGenSPIRV/wave-ops-include-helper-lanes-not-present.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -T vs_6_7 -E main -spirv %s | FileCheck %s
+//
+// CHECK-NOT: OpCapability QuadControlKHR
+// CHECK-NOT: OpExtension "SPV_KHR_maximal_reconvergence"
+// CHECK-NOT: OpExtension "SPV_KHR_quad_control"
+//
+// CHECK-NOT: OpExecutionMode %main MaximallyReconvergesKHR
+// CHECK-NOT: OpExecutionMode %main RequireFullQuadsKHR
+
+[WaveOpsIncludeHelperLanes]
+[shader("vertex")]
+float4 main() : SV_Position
+{
+    return float4(0.0, 0.0, 0.0, 0.0);
+}

tools/clang/test/CodeGenSPIRV/wave-ops-include-helper-lanes.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -T ps_6_7 -E main -spirv %s | FileCheck %s
+
+// CHECK: OpCapability QuadControlKHR
+// CHECK-DAG: OpExtension "SPV_KHR_maximal_reconvergence"
+// CHECK-DAG: OpExtension "SPV_KHR_quad_control"
+
+// CHECK: OpExecutionMode %main MaximallyReconvergesKHR
+// CHECK: OpExecutionMode %main RequireFullQuadsKHR
+
+[WaveOpsIncludeHelperLanes]
+float4 main(float4 pos : SV_Position) : SV_Target
+{
+    return pos;
+}