[Spirv] Fix for validation error VUID-StandaloneSpirv-Flat-04744 (#5276)

generated SPIR-V is invalid: [VUID-StandaloneSpirv-Flat-04744] Fragment OpEntryPoint operand 4 with Input interfaces with integer or float type must have a Flat decoration for Entry Point id 1.
  %SubgroupLocalInvocationId = OpVariable %_ptr_Input_uint Input

Author: rajivmg

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -3435,12 +3435,14 @@ SpirvVariable *DeclResultIdMapper::getBuiltinVar(spv::BuiltIn builtIn,
   if (builtInVar != builtinToVarMap.end()) {
     return builtInVar->second;
   }
+  bool mayNeedFlatDecoration = false;
   spv::StorageClass sc = spv::StorageClass::Max;
   // Valid builtins supported
   switch (builtIn) {
   case spv::BuiltIn::SubgroupSize:
   case spv::BuiltIn::SubgroupLocalInvocationId:
     needsLegalization = true;
+    mayNeedFlatDecoration = true;
     LLVM_FALLTHROUGH;
   case spv::BuiltIn::HitTNV:
   case spv::BuiltIn::RayTmaxNV:
@@ -3481,6 +3483,9 @@ SpirvVariable *DeclResultIdMapper::getBuiltinVar(spv::BuiltIn builtIn,
   // Create a dummy StageVar for this builtin variable
   auto var = spvBuilder.addStageBuiltinVar(type, sc, builtIn,
                                            /*isPrecise*/ false, loc);
+  if (mayNeedFlatDecoration && spvContext.isPS()) {
+    spvBuilder.decorateFlat(var, loc);
+  }
 
   const hlsl::SigPoint *sigPoint =
       hlsl::SigPoint::GetSigPoint(hlsl::SigPointFromInputQual(

tools/clang/test/CodeGenSPIRV/sm6.wave-get-lane-count.ps.hlsl

@@ -0,0 +1,19 @@
+// RUN: %dxc -T ps_6_0 -E main -fspv-target-env=vulkan1.1
+
+// CHECK: ; Version: 1.3
+
+// CHECK: OpCapability GroupNonUniform
+
+// CHECK: OpEntryPoint Fragment
+// CHECK-SAME: %SubgroupSize
+
+// CHECK: OpDecorate %SubgroupSize BuiltIn SubgroupSize
+// CHECK: OpDecorate %SubgroupSize Flat
+
+// CHECK: %SubgroupSize = OpVariable %_ptr_Input_uint Input
+
+float4 main() : SV_Target {
+// CHECK: OpLoad %uint %SubgroupSize
+    uint laneCount = WaveGetLaneCount();
+    return float4(laneCount, 1, 1, 1);
+}

tools/clang/test/CodeGenSPIRV/sm6.wave-get-lane-index.ps.hlsl

@@ -0,0 +1,19 @@
+// RUN: %dxc -T ps_6_0 -E main -fspv-target-env=vulkan1.1
+
+// CHECK: ; Version: 1.3
+
+// CHECK: OpCapability GroupNonUniform
+
+// CHECK: OpEntryPoint Fragment
+// CHECK-SAME: %SubgroupLocalInvocationId
+
+// CHECK: OpDecorate %SubgroupLocalInvocationId BuiltIn SubgroupLocalInvocationId
+// CHECK: OpDecorate %SubgroupLocalInvocationId Flat
+
+// CHECK: %SubgroupLocalInvocationId = OpVariable %_ptr_Input_uint Input
+
+float4 main() : SV_Target {
+// CHECK: OpLoad %uint %SubgroupLocalInvocationId
+    uint laneIndex = WaveGetLaneIndex();
+    return float4(laneIndex, 1, 1, 1);
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1541,6 +1541,12 @@ TEST_F(FileTest, SM6WaveGetLaneCount) {
 TEST_F(FileTest, SM6WaveGetLaneIndex) {
   runFileTest("sm6.wave-get-lane-index.hlsl");
 }
+TEST_F(FileTest, SM6WaveGetLaneCountPS) {
+  runFileTest("sm6.wave-get-lane-count.ps.hlsl");
+}
+TEST_F(FileTest, SM6WaveGetLaneIndexPS) {
+  runFileTest("sm6.wave-get-lane-index.ps.hlsl");
+}
 TEST_F(FileTest, SM6WaveBuiltInNoDuplicate) {
   runFileTest("sm6.wave.builtin.no-dup.hlsl");
 }