[SPIRV] Add the derivative group execution mode only on shader types that allow it. (microsoft#7628)

DXC allows user to use decrivative instruction in shader models that do
not allow it, but they must be dead code that will be removed. However,
when we see a derivative instruction in the SPIR-V backend that is not
in a pixel shader we assume it need the DerivativeGroup execution mode,
and we fail when we try to add it to a vertex shader.

To allow out implementation to match DXIL, we will not assume we can add
the execution mode. We will only add it for shader that we know can use
is, and skip the other.

If the derivative instruction is not removed during optimizations, there
will be a validation error.

While fixing this, we observed another bug that is fixed at the same
time since they are closely related. The TaskNV and TaskEXT shader types
do not have the same id, and the SPV_KHR_compute_shader_derivatives does
not work with the NV mesh shader extension. That was fixed up.

Fixes microsoft#7478

Author: s-perron

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -4399,9 +4399,7 @@ SpirvEmitter::processTextureLevelOfDetail(const CXXMemberCallExpr *expr,
       spvBuilder.createImageQuery(spv::Op::OpImageQueryLod, queryResultType,
                                   expr->getExprLoc(), sampledImage, coordinate);
 
-  if (spvContext.isCS() || spvContext.isNode()) {
-    addDerivativeGroupExecutionMode();
-  }
+  addDerivativeGroupExecutionMode();
   // The first component of the float2 contains the mipmap array layer.
   // The second component of the float2 represents the unclamped lod.
   return spvBuilder.createCompositeExtract(astContext.FloatTy, query,
@@ -5780,9 +5778,7 @@ SpirvEmitter::processTextureSampleGather(const CXXMemberCallExpr *expr,
 
   const auto retType = expr->getDirectCallee()->getReturnType();
   if (isSample) {
-    if (spvContext.isCS() || spvContext.isNode()) {
-      addDerivativeGroupExecutionMode();
-    }
+    addDerivativeGroupExecutionMode();
     return createImageSample(retType, imageType, image, sampler, coordinate,
                              /*compareVal*/ nullptr, /*bias*/ nullptr,
                              /*lod*/ nullptr, std::make_pair(nullptr, nullptr),
@@ -5870,9 +5866,9 @@ SpirvEmitter::processTextureSampleBiasLevel(const CXXMemberCallExpr *expr,
 
   const auto retType = expr->getDirectCallee()->getReturnType();
 
-  if (!lod && (spvContext.isCS() || spvContext.isNode())) {
+  if (!lod)
     addDerivativeGroupExecutionMode();
-  }
+
   return createImageSample(
       retType, imageType, image, sampler, coordinate,
       /*compareVal*/ nullptr, bias, lod, std::make_pair(nullptr, nullptr),
@@ -5992,9 +5988,7 @@ SpirvEmitter::processTextureSampleCmp(const CXXMemberCallExpr *expr) {
   const auto retType = expr->getDirectCallee()->getReturnType();
   const auto imageType = imageExpr->getType();
 
-  if (spvContext.isCS()) {
-    addDerivativeGroupExecutionMode();
-  }
+  addDerivativeGroupExecutionMode();
 
   return createImageSample(
       retType, imageType, image, sampler, coordinate, compareVal,
@@ -6047,9 +6041,7 @@ SpirvEmitter::processTextureSampleCmpBias(const CXXMemberCallExpr *expr) {
   const auto retType = expr->getDirectCallee()->getReturnType();
   const auto imageType = imageExpr->getType();
 
-  if (spvContext.isCS()) {
-    addDerivativeGroupExecutionMode();
-  }
+  addDerivativeGroupExecutionMode();
 
   return createImageSample(
       retType, imageType, image, sampler, coordinate, compareVal, bias,
@@ -9782,8 +9774,7 @@ SpirvInstruction *SpirvEmitter::processDerivativeIntrinsic(
   QualType returnType = arg->getAstResultType();
   assert(isFloatOrVecOfFloatType(returnType));
 
-  if (!spvContext.isPS())
-    addDerivativeGroupExecutionMode();
+  addDerivativeGroupExecutionMode();
   needsLegalization = true;
 
   QualType B32Type = astContext.FloatTy;
@@ -12512,8 +12503,7 @@ SpirvInstruction *SpirvEmitter::processIntrinsicUsingSpirvInst(
     case spv::Op::OpFwidth:
     case spv::Op::OpFwidthFine:
     case spv::Op::OpFwidthCoarse:
-      if (spvContext.isCS() || spvContext.isNode())
-        addDerivativeGroupExecutionMode();
+      addDerivativeGroupExecutionMode();
       needsLegalization = true;
       break;
     default:
@@ -15771,8 +15761,29 @@ bool SpirvEmitter::spirvToolsValidate(std::vector<uint32_t> *mod,
   return tools.Validate(mod->data(), mod->size(), options);
 }
 
+static bool canUseDerivativeGroupExecutionMode(SpirvContext::ShaderModelKind sm,
+                                               bool usingEXTMeshShader) {
+  switch (sm) {
+  case SpirvContext::ShaderModelKind::Compute:
+  case SpirvContext::ShaderModelKind::Node:
+    return true;
+
+  // The KHR extension that allows derivative instruction in mesh and task
+  // (amplification) shader does not work with SPV_NV_mesh_shader extesion.
+  case SpirvContext::ShaderModelKind::Mesh:
+  case SpirvContext::ShaderModelKind::Amplification:
+    return usingEXTMeshShader;
+  default:
+    return false;
+  }
+}
+
 void SpirvEmitter::addDerivativeGroupExecutionMode() {
-  assert(spvContext.isCS());
+  bool usingEXTMeshShader =
+      featureManager.isExtensionEnabled(Extension::EXT_mesh_shader);
+  SpirvContext::ShaderModelKind sm = spvContext.getCurrentShaderModelKind();
+  if (!canUseDerivativeGroupExecutionMode(sm, usingEXTMeshShader))
+    return;
 
   SpirvExecutionMode *numThreadsEm =
       cast<SpirvExecutionMode>(spvBuilder.getModule()->findExecutionMode(

tools/clang/test/CodeGenSPIRV/amplification_shader_derivative.hlsl

@@ -0,0 +1,28 @@
+// RUN: %dxc -T as_6_5 -E main -fspv-target-env=vulkan1.3 %s -spirv | FileCheck %s --check-prefix=VK13
+// RUN: %dxc -T as_6_5 -E main -fspv-target-env=vulkan1.1 -Vd %s -spirv | FileCheck %s --check-prefix=VK11
+
+// VK13-DAG: OpCapability ComputeDerivativeGroupLinearKHR
+// VK13-DAG: OpCapability DerivativeControl
+// VK13-DAG: OpCapability MeshShadingEXT
+// VK13-DAG: OpExtension "SPV_EXT_mesh_shader"
+// VK13-DAG: OpExtension "SPV_KHR_compute_shader_derivatives"
+// VK13: OpEntryPoint TaskEXT %main "main"
+// VK13: OpExecutionMode %main DerivativeGroupLinearKHR
+
+// VK11-DAG: OpExtension "SPV_NV_mesh_shader"
+// VK11: OpEntryPoint TaskNV %main "main"
+// VK11-NOT: OpExecutionMode %main DerivativeGroup
+
+struct AmplificationPayload
+{
+    float4 value;
+};
+
+groupshared AmplificationPayload payload;
+
+[numthreads(4, 1, 1)]
+void main(in uint tid : SV_GroupThreadID, in uint gtid : SV_GroupID)
+{
+    payload.value = ddx_coarse(float4(tid, 0, 0, 0));
+    DispatchMesh(1,1,1, payload);
+}

tools/clang/test/CodeGenSPIRV/mesh_shader_derivative.hlsl

@@ -0,0 +1,34 @@
+// RUN: %dxc -T ms_6_5 -E main -fspv-target-env=vulkan1.3 %s -spirv | FileCheck %s --check-prefix=VK13
+// RUN: %dxc -T ms_6_5 -E main -fspv-target-env=vulkan1.1 -Vd %s -spirv | FileCheck %s --check-prefix=VK11
+
+// VK13-DAG: OpCapability ComputeDerivativeGroupLinearKHR
+// VK13-DAG: OpCapability DerivativeControl
+// vk13-DAG: OpCapability MeshShadingEXT
+// VK13-DAG: OpExtension "SPV_EXT_mesh_shader"
+// VK13-DAG: OpExtension "SPV_KHR_compute_shader_derivatives"
+// VK13: OpEntryPoint MeshEXT %main "main"
+// VK13: OpExecutionMode %main DerivativeGroupLinearKHR
+
+// VK11-DAG: OpExtension "SPV_NV_mesh_shader"
+// VK11: OpEntryPoint MeshNV %main "main"
+// VK11-NOT: OpExecutionMode %main DerivativeGroup
+
+struct VSOut
+{
+    float4 pos : SV_Position;
+};
+
+[numthreads(4, 1, 1)]
+[outputtopology("triangle")]
+void main(in uint tid : SV_GroupThreadID, out vertices VSOut verts[3], out indices uint3 tris[1])
+{
+    SetMeshOutputCounts(3, 1);
+
+    float4 val = ddx_coarse(float4(tid, 0, 0, 0));
+
+    verts[0].pos = val;
+    verts[1].pos = val + float4(0,1,0,0);
+    verts[2].pos = val + float4(1,0,0,0);
+
+    tris[0] = uint3(0,1,2);
+}

tools/clang/test/CodeGenSPIRV/vertex_shader_derivative_in_branch.hlsl

@@ -0,0 +1,23 @@
+// RUN: %dxc -T vs_6_0 -E main -DCOND=false -fspv-target-env=vulkan1.3 %s -spirv | FileCheck %s
+// CHECK-NOT: OpCapability DerivativeControl
+// CHECK-NOT: OpExtension "SPV_KHR_compute_shader_derivatives"
+
+// RUN: not %dxc -T vs_6_0 -E main -DCOND=true -fspv-target-env=vulkan1.3 %s -spirv 2>&1 | FileCheck %s -check-prefix=ERROR
+// ERROR: generated SPIR-V is invalid:
+// ERROR-NEXT: Derivative instructions require Fragment, GLCompute, MeshEXT or TaskEXT execution model: DPdx
+
+struct VSOut
+{
+    float4 pos : SV_Position;
+};
+
+VSOut main(float4 pos : POSITION)
+{
+    VSOut output;
+    output.pos = pos;
+    if (COND)
+    {
+        output.pos += ddx(pos);
+    }
+    return output;
+}