[spirv] Add support for SV_ShadingRate as a PS input (#2226)

* Add support for SV_ShadingRate as a PS input

Solves a portion of #2208.

* Fix small typo in error message

docs/SPIR-V.rst

@@ -278,6 +278,7 @@ Supported extensions
 * SPV_KHR_shader_draw_parameters
 * SPV_EXT_descriptor_indexing
 * SPV_EXT_fragment_fully_covered
+* SPV_EXT_fragment_invocation_density
 * SPV_EXT_shader_stencil_support
 * SPV_AMD_shader_explicit_vertex_parameter
 * SPV_GOOGLE_hlsl_functionality1
@@ -1388,6 +1389,8 @@ some system-value (SV) semantic strings will be translated into SPIR-V
 |                           +-------------+--------------------------+-----------------------+-----------------------------+
 |                           | PSIn        | ``ViewIndex``            | N/A                   | ``MultiView``               |
 +---------------------------+-------------+--------------------------+-----------------------+-----------------------------+
+| SV_ShadingRate            | PSIn        | ``FragSizeEXT``          | N/A                   | ``FragmentDensityEXT``      |
++---------------------------+-------------+--------------------------+-----------------------+-----------------------------+
 
 For entities (function parameters, function return values, struct fields) with
 the above SV semantic strings attached, SPIR-V variables of the

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

@@ -36,6 +36,7 @@ enum class Extension {
   KHR_post_depth_coverage,
   EXT_descriptor_indexing,
   EXT_fragment_fully_covered,
+  EXT_fragment_invocation_density,
   EXT_shader_stencil_export,
   EXT_shader_viewport_index_layer,
   AMD_gpu_shader_half_float,

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -322,6 +322,13 @@ bool CapabilityVisitor::visit(SpirvDecoration *decor) {
     case spv::BuiltIn::BaryCoordPullModelAMD: {
       addExtension(Extension::AMD_shader_explicit_vertex_parameter,
                    "SV_Barycentrics", loc);
+      break;
+    }
+    case spv::BuiltIn::FragSizeEXT: {
+      addExtension(Extension::EXT_fragment_invocation_density,
+                   "SV_ShadingRate", loc);
+      addCapability(spv::Capability::FragmentDensityEXT);
+      break;
     }
     default:
       break;

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1710,6 +1710,8 @@ 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(sigPoint->GetKind(), semanticKind,
                                 semanticToUse->index, invocationId, value,
@@ -1747,6 +1749,9 @@ 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;
@@ -1919,6 +1924,24 @@ 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),
+            y, /* SourceLocation */ {});
+      }
 
       // Reciprocate SV_Position.w if requested
       if (semanticKind == hlsl::Semantic::Kind::Position)
@@ -2732,6 +2755,23 @@ SpirvVariable *DeclResultIdMapper::createSpirvStageVar(
     return spvBuilder.addStageBuiltinVar(type, sc, BuiltIn::FullyCoveredEXT,
                                          isPrecise, srcLoc);
   }
+  // 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.
+  case hlsl::Semantic::Kind::ShadingRate: {
+    switch (sigPointKind) {
+    case hlsl::SigPoint::Kind::PSIn:
+      stageVar->setIsSpirvBuiltin();
+      return spvBuilder.addStageBuiltinVar(type, sc, BuiltIn::FragSizeEXT,
+                                           isPrecise, srcLoc);
+    default:
+      emitError("semantic ShadingRate currently unsupported in non-PS shader"
+                " stages",
+                srcLoc);
+      break;
+    }
+    break;
+  }
   default:
     emitError("semantic %0 unimplemented", srcLoc)
         << stageVar->getSemanticStr();

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -103,6 +103,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_EXT_descriptor_indexing", Extension::EXT_descriptor_indexing)
       .Case("SPV_EXT_fragment_fully_covered",
             Extension::EXT_fragment_fully_covered)
+      .Case("SPV_EXT_fragment_invocation_density",
+            Extension::EXT_fragment_invocation_density)
       .Case("SPV_EXT_shader_stencil_export",
             Extension::EXT_shader_stencil_export)
       .Case("SPV_EXT_shader_viewport_index_layer",
@@ -136,6 +138,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_EXT_descriptor_indexing";
   case Extension::EXT_fragment_fully_covered:
     return "SPV_EXT_fragment_fully_covered";
+  case Extension::EXT_fragment_invocation_density:
+    return "SPV_EXT_fragment_invocation_density";
   case Extension::EXT_shader_stencil_export:
     return "SPV_EXT_shader_stencil_export";
   case Extension::EXT_shader_viewport_index_layer:

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

@@ -0,0 +1,11 @@
+// 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.vs-error.hlsl

@@ -0,0 +1,7 @@
+// 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/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1982,5 +1982,11 @@ TEST_F(FileTest, VulkanShaderRecordBufferNVOffset) {
   // Checks the behavior of [[vk::offset]] with [[vk::shader_record_nv]]
   runFileTest("vk.shader-record-nv.offset.hlsl");
 }
+TEST_F(FileTest, VulkanShadingRate) {
+  runFileTest("vk.shading-rate.hlsl");
+}
+TEST_F(FileTest, VulkanShadingRateError) {
+  runFileTest("vk.shading-rate.vs-error.hlsl", Expect::Failure);
+}
 
 } // namespace