Add support for KHR_compute_shader_derivatives (#7249)

Add support for KHR_compute_shader_derivatives
- DirectxShaderCompiler already supports `NV_compute_shader_derivatives`
which is functionality identical to `KHR_compute_shader_derivatives`
- The KHR extension will be used by default instead of the NV one
following the same approach as the RT extension.
- We currently explain this in a comment in
`tools/clang/lib/SPIRV/FeatureManager.cpp`
`FeatureManager::enabledByDefault`.
- Check commit introducing RT for more info
04a84f0

Fixes #7179

Author: iagoCL

docs/SPIR-V.rst

@@ -315,6 +315,7 @@ Supported extensions
 * SPV_KHR_fragment_shader_barycentric
 * SPV_KHR_physical_storage_buffer
 * SPV_KHR_vulkan_memory_model
+* SPV_KHR_compute_shader_derivatives
 * SPV_NV_compute_shader_derivatives
 * SPV_KHR_maximal_reconvergence
 * SPV_KHR_float_controls

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -59,6 +59,7 @@ enum class Extension {
   KHR_physical_storage_buffer,
   KHR_vulkan_memory_model,
   NV_compute_shader_derivatives,
+  KHR_compute_shader_derivatives,
   KHR_fragment_shader_barycentric,
   KHR_maximal_reconvergence,
   KHR_float_controls,

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -852,6 +852,12 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
           spv::Capability::FragmentShaderShadingRateInterlockEXT,
       });
 
+  addExtensionAndCapabilitiesIfEnabled(
+      Extension::KHR_compute_shader_derivatives,
+      {
+          spv::Capability::ComputeDerivativeGroupQuadsKHR,
+          spv::Capability::ComputeDerivativeGroupLinearKHR,
+      });
   addExtensionAndCapabilitiesIfEnabled(
       Extension::NV_compute_shader_derivatives,
       {

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -215,6 +215,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_KHR_physical_storage_buffer",
             Extension::KHR_physical_storage_buffer)
       .Case("SPV_KHR_vulkan_memory_model", Extension::KHR_vulkan_memory_model)
+      .Case("SPV_KHR_compute_shader_derivatives",
+            Extension::KHR_compute_shader_derivatives)
       .Case("SPV_NV_compute_shader_derivatives",
             Extension::NV_compute_shader_derivatives)
       .Case("SPV_KHR_fragment_shader_barycentric",
@@ -283,6 +285,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_KHR_physical_storage_buffer";
   case Extension::KHR_vulkan_memory_model:
     return "SPV_KHR_vulkan_memory_model";
+  case Extension::KHR_compute_shader_derivatives:
+    return "SPV_KHR_compute_shader_derivatives";
   case Extension::NV_compute_shader_derivatives:
     return "SPV_NV_compute_shader_derivatives";
   case Extension::KHR_fragment_shader_barycentric:
@@ -370,6 +374,10 @@ bool FeatureManager::enabledByDefault(Extension ext) {
     // KHR_ray_tracing and NV_ray_tracing are mutually exclusive so enable only
     // KHR extension by default
   case Extension::NV_ray_tracing:
+    return false;
+    // KHR_compute_shader_derivatives and NV_compute_shader_derivatives are
+    // mutually exclusive so enable only KHR extension by default
+  case Extension::NV_compute_shader_derivatives:
     return false;
     // Enabling EXT_demote_to_helper_invocation changes the code generation
     // behavior for the 'discard' statement. Therefore we will only enable it if

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -15049,6 +15049,10 @@ void SpirvEmitter::addDerivativeGroupExecutionMode() {
   // to 2D quad rules. Using derivative operations in any numthreads
   // configuration not matching either of these is invalid and will produce an
   // error.
+  static_assert(spv::ExecutionMode::DerivativeGroupQuadsNV ==
+                spv::ExecutionMode::DerivativeGroupQuadsKHR);
+  static_assert(spv::ExecutionMode::DerivativeGroupLinearNV ==
+                spv::ExecutionMode::DerivativeGroupLinearKHR);
   spv::ExecutionMode em = spv::ExecutionMode::DerivativeGroupQuadsNV;
   if (numThreads[0] % 4 == 0 && numThreads[1] == 1 && numThreads[2] == 1) {
     em = spv::ExecutionMode::DerivativeGroupLinearNV;

tools/clang/test/CodeGenSPIRV/ddx.compute.khr.hlsl

@@ -0,0 +1,29 @@
+// RUN: %dxc -T cs_6_6 -E main -fspv-extension=SPV_KHR_compute_shader_derivatives -fcgl  %s -spirv  2>&1 | FileCheck %s
+
+// CHECK: OpCapability ComputeDerivativeGroupQuadsKHR
+// CHECK: OpExtension "SPV_KHR_compute_shader_derivatives"
+// CHECK: OpExecutionMode %main DerivativeGroupQuadsKHR
+
+
+SamplerState ss : register(s2);
+SamplerComparisonState scs;
+
+RWStructuredBuffer<uint> o;
+Texture1D        <float>  t1;
+
+[numthreads(2,2,1)]
+void main(uint3 id : SV_GroupThreadID)
+{
+    // CHECK: OpDPdx %float %float_0_5
+    o[0] = ddx(0.5);
+    // CHECK: OpDPdxCoarse %float %float_0_5
+    o[1] = ddx_coarse(0.5);
+    // CHECK: OpDPdy %float %float_0_5
+    o[2] = ddy(0.5);
+    // CHECK: OpDPdyCoarse %float %float_0_5
+    o[3] = ddy_coarse(0.5);
+    // CHECK: OpDPdxFine %float %float_0_5
+    o[4] = ddx_fine(0.5);
+    // CHECK: OpDPdyFine %float %float_0_5
+    o[5] = ddy_fine(0.5);
+}

tools/clang/test/CodeGenSPIRV/texture.calculate.lod.compute.linear.khr.hlsl

@@ -0,0 +1,23 @@
+// RUN: %dxc -T cs_6_6 -E main -fspv-extension=SPV_KHR_compute_shader_derivatives -fcgl  %s -spirv  2>&1 | FileCheck %s --check-prefix=CHECK
+// RUN: %dxc -T cs_6_6 -E main -fspv-extension=SPV_KHR_compute_shader_derivatives %s -spirv  2>&1 | FileCheck %s --check-prefix=CHECK
+
+// CHECK: OpCapability ComputeDerivativeGroupLinearKHR
+// CHECK: OpExtension "SPV_KHR_compute_shader_derivatives"
+// CHECK: OpExecutionMode %main DerivativeGroupLinearKHR
+
+SamplerState ss : register(s2);
+SamplerComparisonState scs;
+
+RWStructuredBuffer<uint> o;
+Texture1D        <float>  t1;
+
+[numthreads(16,1,1)]
+void main(uint3 id : SV_GroupThreadID)
+{
+    //CHECK:          [[t1:%[0-9]+]] = OpLoad %type_1d_image %t1
+    //CHECK-NEXT:    [[ss1:%[0-9]+]] = OpLoad %type_sampler %ss
+    //CHECK-NEXT:    [[si1:%[0-9]+]] = OpSampledImage %type_sampled_image [[t1]] [[ss1]]
+    //CHECK-NEXT: [[query1:%[0-9]+]] = OpImageQueryLod %v2float [[si1]] %float_0_5
+    //CHECK-NEXT:        {{%[0-9]+}} = OpCompositeExtract %float [[query1]] 0
+    o[0] = t1.CalculateLevelOfDetail(ss, 0.5);
+}