Make the location map run per entry point (microsoft#6688)

The code that adds the input and output decoration in the entry points
inputs and outputs assumes that there is a single entry point in the
module. When using the `lib` profile that is not true.

This commit modifies the code so that it groups the stage variables by
entry point, and runs the current code on each group separably.

I hesitate to make this change because it will change the locations for
code that currently works, and will force users to update their
applications accordingly. Or they could modify their shaders
to use explicit locations attributes. Neither is great.

However, the advantage is that this allows the implicit locations to
match what would happen if the shader were compiled individually. It
also makes the locations more predictable because change in another
shader would change all shader after it. This is a better design, and
worth the breakage.

Fixes microsoft#6678
Fixes microsoft#5213

Author: s-perron

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -2143,10 +2143,8 @@ bool DeclResultIdMapper::assignLocations(
   return true;
 }
 
-bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
-  if (!checkSemanticDuplication(forInput))
-    return false;
-
+bool DeclResultIdMapper::finalizeStageIOLocationsForASingleEntryPoint(
+    bool forInput, ArrayRef<StageVar> functionStageVars) {
   // Returns false if the given StageVar is an input/output variable without
   // explicit location assignment. Otherwise, returns true.
   const auto locAssigned = [forInput, this](const StageVar &v) {
@@ -2166,11 +2164,12 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
 
   // If we have explicit location specified for all input/output variables,
   // use them instead assign by ourselves.
-  if (std::all_of(stageVars.begin(), stageVars.end(), locAssigned)) {
+  if (std::all_of(functionStageVars.begin(), functionStageVars.end(),
+                  locAssigned)) {
     LocationSet locSet;
     bool noError = true;
 
-    for (const auto &var : stageVars) {
+    for (const auto &var : functionStageVars) {
       // Skip builtins & those stage variables we are not handling for this call
       if (var.isSpirvBuitin() || var.hasLocOrBuiltinDecorateAttr() ||
           forInput != isInputStorageClass(var)) {
@@ -2189,7 +2188,7 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
           emitError("stage %select{output|input}0 location #%1 already "
                     "consumed by semantic '%2'",
                     attrLoc)
-              << forInput << l << stageVars[idx].getSemanticStr();
+              << forInput << l << functionStageVars[idx].getSemanticStr();
           noError = false;
         }
 
@@ -2213,7 +2212,7 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
   std::vector<const StageVar *> vars;
   LocationSet locSet;
 
-  for (const auto &var : stageVars) {
+  for (const auto &var : functionStageVars) {
     if (var.isSpirvBuitin() || var.hasLocOrBuiltinDecorateAttr() ||
         forInput != isInputStorageClass(var)) {
       continue;
@@ -2286,6 +2285,29 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
   return assignLocations(vars, nextLocs, &stageVariableLocationInfo);
 }
 
+llvm::DenseMap<const SpirvFunction *, SmallVector<StageVar, 8>>
+DeclResultIdMapper::getStageVarsPerFunction() {
+  llvm::DenseMap<const SpirvFunction *, SmallVector<StageVar, 8>> result;
+  for (const auto &var : stageVars) {
+    result[var.getEntryPoint()].push_back(var);
+  }
+  return result;
+}
+
+bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
+  if (!checkSemanticDuplication(forInput))
+    return false;
+
+  auto stageVarPerFunction = getStageVarsPerFunction();
+  for (const auto &functionStageVars : stageVarPerFunction) {
+    if (!finalizeStageIOLocationsForASingleEntryPoint(
+            forInput, functionStageVars.getSecond())) {
+      return false;
+    }
+  }
+  return true;
+}
+
 namespace {
 /// A class for maintaining the binding number shift requested for descriptor
 /// sets.

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -627,6 +627,22 @@ class DeclResultIdMapper {
       llvm::DenseSet<StageVariableLocationInfo, StageVariableLocationInfo>
           *stageVariableLocationInfo);
 
+  /// \brief Returns a map that divides all of the shader stage variables into
+  /// separate vectors for each entry point.
+  llvm::DenseMap<const SpirvFunction *, SmallVector<StageVar, 8>>
+  getStageVarsPerFunction();
+
+  /// \brief Decorates all stage variables in `functionStageVars` with proper
+  /// location and returns true on success.
+  ///
+  /// It is assumed that all variables in `functionStageVars` belong to the same
+  /// entry point.
+  ///
+  /// This method will write the location assignment into the module under
+  /// construction.
+  bool finalizeStageIOLocationsForASingleEntryPoint(
+      bool forInput, ArrayRef<StageVar> functionStageVars);
+
   /// \brief Decorates all stage input (if forInput is true) or output (if
   /// forInput is false) variables with proper location and returns true on
   /// success.

tools/clang/test/CodeGenSPIRV/attribute.location.lib.hlsl

@@ -0,0 +1,38 @@
+// RUN: %dxc -T lib_6_6 %s -spirv | FileCheck %s
+
+
+// The inputs to the vertex shader should be at location 0 and 1 to match
+// the location attributes.
+// CHECK-DAG: OpDecorate %in_var_Position Location 0
+// CHECK-DAG: OpDecorate %in_var_Color Location 1
+struct Vertex {
+    [[vk::location(0)]] float3 position : Position;
+    [[vk::location(1)]] float4 color : Color;
+};
+
+// The "color" output to the vertex shader should match the "color" input to the
+// pixel shader. The should both be at location 1 to match the attribute.
+// CHECK-DAG: OpDecorate %out_var_Color Location 1
+// CHECK-DAG: OpDecorate %in_var_Color_0 Location 1
+struct Fragment {
+    [[vk::location(0)]] float4 sv_position : SV_Position;
+    [[vk::location(1)]] float4 color : Color;
+};
+
+[shader("vertex")]
+Fragment vsmain(in Vertex vertex) {
+    Fragment o;
+    
+    o.sv_position = float4(vertex.position, 1);
+    o.color = vertex.color;
+
+    return o;
+}
+
+// This location should set the location of the pixel shader output to location 0.
+// CHECK-DAG: OpDecorate %out_var_SV_Target0 Location 0
+[[vk::location(0)]]
+[shader("pixel")]
+float4 psmain(in Fragment fragment) : SV_Target0 {
+    return fragment.color;
+}

tools/clang/test/CodeGenSPIRV/implicit.location.lib.hlsl

@@ -0,0 +1,46 @@
+// RUN: %dxc -T lib_6_6 %s -spirv | FileCheck %s
+
+// The inputs for the vertex shader are 0 and 1. vId is a builtin and does not
+// get a location.
+// CHECK-DAG: OpDecorate %in_var_COLOR0 Location 0
+// CHECK-DAG: OpDecorate %in_var_COLOR1 Location 1
+struct VIN {
+  uint vId : SV_VertexID;
+  float2 col0 : COLOR0;
+  float2 col1 : COLOR1;
+};
+
+// The output for the vertex shader and the input for the pixel shader should
+// both be 0. Position is a builtin, and each shader should restart their
+// numbering at 0.
+// CHECK-DAG: OpDecorate %in_var_COLOR2 Location 0
+// CHECK-DAG: OpDecorate %out_var_COLOR2 Location 0
+struct V2P
+{
+    float4 Pos : SV_Position;
+    float2 Uv : COLOR2;
+};
+
+#define PI (3.14159f)
+
+[shader("vertex")]
+V2P VSMain(VIN vIn)
+{
+    float2 uv = vIn.col0 + vIn.col1;
+    V2P vsOut;
+    vsOut.Uv = float2(uv.x, 1.0 - uv.y);
+    vsOut.Pos = float4((2.0 * uv) - 1.0, 0.0, 1.0);
+    return vsOut;
+}
+
+[[vk::binding(0)]] Texture2D<float3> source_tex;
+[[vk::binding(1)]] SamplerState samp;
+
+// The location for the sole output from the pixel shader should be 0. The
+// numbering should restart at 0 for each shader.
+// CHECK-DAG: OpDecorate %out_var_SV_Target0 Location 0
+[shader("pixel")]
+float4 PSMain(V2P psIn) : SV_Target0
+{
+    return float4(source_tex.Sample(samp, psIn.Uv), 1.0);
+}