[spirv] Add VK_KHR_variable_rate_fragment_shading SPIR-V support (#3664)

Change the mapping of SV_ShadingRate from FragSizeEXT (VK_EXT_fragment_density_map) to PrimitiveShadingRateKHR/ShadingRateKHR (VK_KHR_variable_rate_fragment_shading).

docs/SPIR-V.rst

@@ -283,7 +283,7 @@ Supported extensions
 * SPV_KHR_shader_draw_parameters
 * SPV_EXT_descriptor_indexing
 * SPV_EXT_fragment_fully_covered
-* SPV_EXT_fragment_invocation_density
+* SPV_KHR_fragment_shading_rate
 * SPV_EXT_shader_stencil_support
 * SPV_AMD_shader_explicit_vertex_parameter
 * SPV_GOOGLE_hlsl_functionality1
@@ -1495,7 +1495,13 @@ some system-value (SV) semantic strings will be translated into SPIR-V
 |                           +-------------+----------------------------------------+-----------------------+-----------------------------+
 |                           | MSIn        | ``ViewIndex``                          | N/A                   | ``MultiView``               |
 +---------------------------+-------------+----------------------------------------+-----------------------+-----------------------------+
-| SV_ShadingRate            | PSIn        | ``FragSizeEXT``                        | N/A                   | ``FragmentDensityEXT``      |
+|                           | VSOut       | ``PrimitiveShadingRateKHR``            | N/A                   | ``FragmentShadingRate``     |
+|                           +-------------+----------------------------------------+-----------------------+-----------------------------+
+|                           | GSOut       | ``PrimitiveShadingRateKHR``            | N/A                   | ``FragmentShadingRate``     |
+| SV_ShadingRate            +-------------+----------------------------------------+-----------------------+-----------------------------+
+|                           | PSIn        | ``ShadingRateKHR``                     | N/A                   | ``FragmentShadingRate``     |
+|                           +-------------+----------------------------------------+-----------------------+-----------------------------+
+|                           | MSOut       | ``PrimitiveShadingRateKHR``            | N/A                   | ``FragmentShadingRate``     |
 +---------------------------+-------------+----------------------------------------+-----------------------+-----------------------------+
 
 For entities (function parameters, function return values, struct fields) with

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

@@ -31,6 +31,7 @@ enum class Extension {
   KHR = 0,
   KHR_16bit_storage,
   KHR_device_group,
+  KHR_fragment_shading_rate,
   KHR_non_semantic_info,
   KHR_multiview,
   KHR_shader_draw_parameters,

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -350,10 +350,10 @@ bool CapabilityVisitor::visit(SpirvDecoration *decor) {
                    "SV_Barycentrics", loc);
       break;
     }
-    case spv::BuiltIn::FragSizeEXT: {
-      addExtension(Extension::EXT_fragment_invocation_density, "SV_ShadingRate",
-                   loc);
-      addCapability(spv::Capability::FragmentDensityEXT);
+    case spv::BuiltIn::ShadingRateKHR:
+    case spv::BuiltIn::PrimitiveShadingRateKHR: {
+      addExtension(Extension::KHR_fragment_shading_rate, "SV_ShadingRate", loc);
+      addCapability(spv::Capability::FragmentShadingRateKHR);
       break;
     }
     default:

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -2221,8 +2221,6 @@ bool DeclResultIdMapper::createStageVars(
     // * SV_DispatchThreadID, SV_GroupThreadID, and SV_GroupID are allowed to be
     //   uint, uint2, or uint3, but the corresponding builtins
     //   (GlobalInvocationId, LocalInvocationId, WorkgroupId) must be a uint3.
-    // * SV_ShadingRate is a uint value, but the builtin it corresponds to is a
-    //   int2.
 
     if (glPerVertex.tryToAccess(sigPointKind, semanticKind,
                                 semanticToUse->index, invocationId, value,
@@ -2264,9 +2262,6 @@ bool DeclResultIdMapper::createStageVars(
           hlsl::IsHLSLVecType(type) ? hlsl::GetHLSLVecElementType(type) : type,
           3);
       break;
-    case hlsl::Semantic::Kind::ShadingRate:
-      evalType = astContext.getExtVectorType(astContext.IntTy, 2);
-      break;
     default:
       // Only the semantic kinds mentioned above are handled.
       break;
@@ -2505,25 +2500,6 @@ bool DeclResultIdMapper::createStageVars(
               astContext.getExtVectorType(srcVecElemType, 2), *value, *value,
               {0, 1}, thisSemantic.loc);
       }
-      // Special handling of SV_ShadingRate, which is a bitpacked enum value,
-      // but SPIR-V's FragSizeEXT uses an int2. We build the enum value from
-      // the separate axis values.
-      else if (semanticKind == hlsl::Semantic::Kind::ShadingRate) {
-        // From the D3D12 functional spec for Variable-Rate Shading.
-        // #define D3D12_MAKE_COARSE_SHADING_RATE(x,y) ((x) << 2 | (y))
-        const auto x = spvBuilder.createCompositeExtract(
-            astContext.IntTy, *value, {0}, thisSemantic.loc);
-        const auto y = spvBuilder.createCompositeExtract(
-            astContext.IntTy, *value, {1}, thisSemantic.loc);
-        const auto constTwo =
-            spvBuilder.getConstantInt(astContext.IntTy, llvm::APInt(32, 2));
-        *value = spvBuilder.createBinaryOp(
-            spv::Op::OpBitwiseOr, astContext.UnsignedIntTy,
-            spvBuilder.createBinaryOp(spv::Op::OpShiftLeftLogical,
-                                      astContext.IntTy, x, constTwo,
-                                      thisSemantic.loc),
-            y, thisSemantic.loc);
-      }
 
       // Reciprocate SV_Position.w if requested
       if (semanticKind == hlsl::Semantic::Kind::Position)
@@ -3472,16 +3448,30 @@ SpirvVariable *DeclResultIdMapper::createSpirvStageVar(
   }
   // According to DXIL spec, the ShadingRate SV can only be used by GSOut,
   // VSOut, or PSIn. According to Vulkan spec, the FragSizeEXT BuiltIn can only
-  // be used as PSIn.
+  // be used as VSOut, GSOut, MSOut or PSIn.
   case hlsl::Semantic::Kind::ShadingRate: {
+    QualType checkType = type->getAs<ReferenceType>()
+                             ? type->getAs<ReferenceType>()->getPointeeType()
+                             : type;
+    QualType scalarTy;
+    if (!isScalarType(checkType, &scalarTy) || !scalarTy->isIntegerType()) {
+      emitError("semantic ShadingRate must be interger scalar type", srcLoc);
+    }
+
     switch (sigPointKind) {
     case hlsl::SigPoint::Kind::PSIn:
       stageVar->setIsSpirvBuiltin();
-      return spvBuilder.addStageBuiltinVar(type, sc, BuiltIn::FragSizeEXT,
+      return spvBuilder.addStageBuiltinVar(type, sc, BuiltIn::ShadingRateKHR,
                                            isPrecise, srcLoc);
+    case hlsl::SigPoint::Kind::VSOut:
+    case hlsl::SigPoint::Kind::GSOut:
+    case hlsl::SigPoint::Kind::MSOut:
+      stageVar->setIsSpirvBuiltin();
+      return spvBuilder.addStageBuiltinVar(
+          type, sc, BuiltIn::PrimitiveShadingRateKHR, isPrecise, srcLoc);
     default:
-      emitError("semantic ShadingRate currently unsupported in non-PS shader"
-                " stages",
+      emitError("semantic ShadingRate must be used only for PSIn, VSOut, "
+                "GSOut, MSOut",
                 srcLoc);
       break;
     }

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -140,6 +140,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_NV_ray_tracing", Extension::NV_ray_tracing)
       .Case("SPV_NV_mesh_shader", Extension::NV_mesh_shader)
       .Case("SPV_KHR_ray_query", Extension::KHR_ray_query)
+      .Case("SPV_KHR_fragment_shading_rate",
+            Extension::KHR_fragment_shading_rate)
       .Default(Extension::Unknown);
 }
 
@@ -189,6 +191,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_NV_mesh_shader";
   case Extension::KHR_ray_query:
     return "SPV_KHR_ray_query";
+  case Extension::KHR_fragment_shading_rate:
+    return "SPV_KHR_fragment_shading_rate";
   default:
     break;
   }

tools/clang/test/CodeGenSPIRV/vk.shading-rate.hlsl

@@ -1,11 +0,0 @@
-// Run: %dxc -T ps_6_4 -E main
-
-float4 main(uint rate : SV_ShadingRate) : SV_TARGET {
-// CHECK:   OpDecorate [[r:%\d+]] BuiltIn FragSizeEXT
-  return float4(rate, 0, 0, 0);
-// CHECK:   [[v:%\d+]] = OpLoad %v2int [[r]]
-// CHECK:   [[x:%\d+]] = OpCompositeExtract %int [[v]] 0
-// CHECK:   [[y:%\d+]] = OpCompositeExtract %int [[v]] 1
-// CHECK:  [[xs:%\d+]] = OpShiftLeftLogical %int [[x]] %int_2
-// CHECK:     {{%\d+}} = OpBitwiseOr %uint [[xs]] [[y]]
-}

tools/clang/test/CodeGenSPIRV/vk.shading-rate.ps.hlsl

@@ -0,0 +1,9 @@
+// Run: %dxc -T ps_6_4 -E main
+
+float4 main(uint rate : SV_ShadingRate) : SV_TARGET {
+// CHECK:   OpCapability FragmentShadingRateKHR
+// CHECK:   OpExtension "SPV_KHR_fragment_shading_rate"
+// CHECK:   OpDecorate [[r:%\d+]] BuiltIn ShadingRateKHR
+// CHECK:   [[r:%\d+]] = OpVariable %_ptr_Input_uint Input
+  return float4(rate, 0, 0, 0);
+}

tools/clang/test/CodeGenSPIRV/vk.shading-rate.vs-error.hlsl

@@ -1,7 +0,0 @@
-// Run: %dxc -T vs_6_4 -E main
-
-void main(out uint rate : SV_ShadingRate) {
-    rate = 0;
-}
-
-// CHECK:  :3:27: error: semantic ShadingRate currently unsupported in non-PS shader stages

tools/clang/test/CodeGenSPIRV/vk.shading-rate.vs.hlsl

@@ -0,0 +1,9 @@
+// Run: %dxc -T vs_6_4 -E main
+
+void main(out uint rate : SV_ShadingRate) {
+// CHECK:   OpCapability FragmentShadingRateKHR
+// CHECK:   OpExtension "SPV_KHR_fragment_shading_rate"
+// CHECK:   OpDecorate [[r:%\d+]] BuiltIn PrimitiveShadingRateKHR
+// CHECK:   [[r:%\d+]] = OpVariable %_ptr_Output_uint Output
+    rate = 0;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -2371,9 +2371,11 @@ TEST_F(FileTest, VulkanShaderRecordBufferEXTOffset) {
   // Checks the behavior of [[vk::offset]] with [[vk::shader_record_ext]]
   runFileTest("vk.shader-record-ext.offset.hlsl");
 }
-TEST_F(FileTest, VulkanShadingRate) { runFileTest("vk.shading-rate.hlsl"); }
-TEST_F(FileTest, VulkanShadingRateError) {
-  runFileTest("vk.shading-rate.vs-error.hlsl", Expect::Failure);
+TEST_F(FileTest, VulkanShadingRateVs) {
+  runFileTest("vk.shading-rate.vs.hlsl");
+}
+TEST_F(FileTest, VulkanShadingRatePs) {
+  runFileTest("vk.shading-rate.ps.hlsl");
 }
 
 // === MeshShading NV examples ===