[SPIR-V] Fix usage of indices in subfunctions (microsoft#7242)

The parameters tagged with indices are linked to a builtin. Because
their layout is different between HLSL and SPIR-V, there is a common
mechanism to handle those 'stage I/O variables'.

Usually, a local variable with the correct HLSL layout is created, and
when required, the value is copied in and copied out in the entrypoint
wrapper. Then, a function-scoped pointer is passed to sub-functions.

The issue is that `indices` marks an array which is also shared across
invocations. Meaning we cannot simple copy-in/copy-out. We are only
allowed to write to the indices touched by the shader.

This required pushing the handling to the assignment expression
handling: when a value is assigned to such builtin, the layout
transformation is done, and the builtin written to.

Issue was how to find back the Builtin from an assignment: the code
assumed the ParmDecl of the entrypoint was the only way to access this
variable, but nothing prevents the user to pass this indice array to
another function.
The simple solution is to move this out of the generic map, and have a
new field which stored the SpirvVariable we created, and allow any HLSL
function to access this as soon as the HLSLIndices attribute is found.

Fixes microsoft#7009

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -860,7 +860,7 @@ bool DeclResultIdMapper::createStageOutputVar(const DeclaratorDecl *decl,
         QualType arrayType = astContext.getConstantArrayType(
             type, llvm::APInt(32, arraySize), clang::ArrayType::Normal, 0);
 
-        stageVarInstructions[cast<DeclaratorDecl>(decl)] =
+        msOutIndicesBuiltin =
             getBuiltinVar(builtinID, arrayType, decl->getLocation());
       } else {
         // For NV_mesh_shader, the built type is PrimitiveIndicesNV
@@ -871,7 +871,7 @@ bool DeclResultIdMapper::createStageOutputVar(const DeclaratorDecl *decl,
             astContext.UnsignedIntTy, llvm::APInt(32, arraySize),
             clang::ArrayType::Normal, 0);
 
-        stageVarInstructions[cast<DeclaratorDecl>(decl)] =
+        msOutIndicesBuiltin =
             getBuiltinVar(builtinID, arrayType, decl->getLocation());
       }
 

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -559,6 +559,11 @@ class DeclResultIdMapper {
     return value;
   }
 
+  SpirvVariable *getMSOutIndicesBuiltin() {
+    assert(msOutIndicesBuiltin && "Variable usage before decl parsing.");
+    return msOutIndicesBuiltin;
+  }
+
   /// Decorate with spirv intrinsic attributes with lamda function variable
   /// check
   void decorateWithIntrinsicAttrs(
@@ -1014,6 +1019,25 @@ class DeclResultIdMapper {
   /// creating that stage variable, so that we don't need to query them again
   /// for reading and writing.
   llvm::DenseMap<const ValueDecl *, SpirvVariable *> stageVarInstructions;
+
+  /// Special case for the Indices builtin:
+  /// - this builtin has a different layout in HLSL & SPIR-V, meaning it
+  /// requires
+  ///   the same kind of handling as classic stageVarInstructions:
+  ///   -> load into a HLSL compatible tmp
+  ///   -> write back into the SPIR-V compatible layout.
+  /// - but the builtin is shared across invocations (not only lanes).
+  ///   -> we must only write/read from the indices requested by the user.
+  /// - the variable can be passed to other functions as a out param
+  ///   -> we cannot copy-in/copy-out because shared across invocations.
+  ///   -> we cannot pass a simple pointer: layout differences between
+  ///   HLSL/SPIR-V.
+  ///
+  /// All this means we must keep track of the builtin, and each assignment to
+  /// this will have to handle the layout differences. The easiest solution is
+  /// to keep this builtin global to the module if present.
+  SpirvVariable *msOutIndicesBuiltin = nullptr;
+
   /// Vector of all defined resource variables.
   llvm::SmallVector<ResourceVar, 8> resourceVars;
   /// Mapping from {RW|Append|Consume}StructuredBuffers to their

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -8133,17 +8133,21 @@ void SpirvEmitter::assignToMSOutIndices(
   if (indices.size() > 1) {
     vecComponent = indices.back();
   }
-  auto *var = declIdMapper.getStageVarInstruction(decl);
-  const auto *varTypeDecl = astContext.getAsConstantArrayType(decl->getType());
-  QualType varType = varTypeDecl->getElementType();
+  SpirvVariable *var = declIdMapper.getMSOutIndicesBuiltin();
+
   uint32_t numVertices = 1;
-  if (!isVectorType(varType, nullptr, &numVertices)) {
-    assert(isScalarType(varType));
-  }
-  QualType valueType = value->getAstResultType();
   uint32_t numValues = 1;
-  if (!isVectorType(valueType, nullptr, &numValues)) {
-    assert(isScalarType(valueType));
+  {
+    const auto *varTypeDecl =
+        astContext.getAsConstantArrayType(decl->getType());
+    QualType varType = varTypeDecl->getElementType();
+    if (!isVectorType(varType, nullptr, &numVertices)) {
+      assert(isScalarType(varType));
+    }
+    QualType valueType = value->getAstResultType();
+    if (!isVectorType(valueType, nullptr, &numValues)) {
+      assert(isScalarType(valueType));
+    }
   }
 
   const auto loc = decl->getLocation();
@@ -8190,7 +8194,10 @@ void SpirvEmitter::assignToMSOutIndices(
       assert(numValues == numVertices);
       if (extMesh) {
         // create accesschain for Primitive*IndicesEXT[vertIndex].
-        auto *ptr = spvBuilder.createAccessChain(varType, var, vertIndex, loc);
+        const ConstantArrayType *CAT =
+            astContext.getAsConstantArrayType(var->getAstResultType());
+        auto *ptr = spvBuilder.createAccessChain(CAT->getElementType(), var,
+                                                 vertIndex, loc);
         // finally create store for Primitive*IndicesEXT[vertIndex] = value.
         spvBuilder.createStore(ptr, value, loc);
       } else {

tools/clang/test/CodeGenSPIRV/meshshading.nv.triangle.indices.out.hlsl

@@ -0,0 +1,65 @@
+// RUN: %dxc -T ms_6_5 -E outie -fcgl  %s -spirv | FileCheck %s
+// RUN: %dxc -T ms_6_5 -E innie -fcgl  %s -spirv | FileCheck %s
+
+// CHECK-DAG: [[v4_n05_05_0_1:%[0-9]+]] = OpConstantComposite %v4float %float_n0_5 %float_0_5 %float_0 %float_1
+// CHECK-DAG:  [[v4_05_05_0_1:%[0-9]+]] = OpConstantComposite %v4float %float_0_5 %float_0_5 %float_0 %float_1
+// CHECK-DAG:  [[v4_0_n05_0_1:%[0-9]+]] = OpConstantComposite %v4float %float_0 %float_n0_5 %float_0 %float_1
+// CHECK-DAG:  [[v3_1_0_0:%[0-9]+]] = OpConstantComposite %v3float %float_1 %float_0 %float_0
+// CHECK-DAG:  [[v3_0_1_0:%[0-9]+]] = OpConstantComposite %v3float %float_0 %float_1 %float_0
+// CHECK-DAG:  [[v3_0_0_1:%[0-9]+]] = OpConstantComposite %v3float %float_0 %float_0 %float_1
+// CHECK-DAG:  [[u3_0_1_2:%[0-9]+]] = OpConstantComposite %v3uint %uint_0 %uint_1 %uint_2
+
+// CHECK-DAG:  OpDecorate [[indices:%[0-9]+]] BuiltIn PrimitiveIndicesNV
+
+struct MeshOutput {
+  float4 position : SV_Position;
+  float3 color : COLOR0;
+};
+
+[outputtopology("triangle")]
+[numthreads(1, 1, 1)]
+void innie(out indices uint3 triangles[1], out vertices MeshOutput verts[3]) {
+    SetMeshOutputCounts(3, 2);
+
+    triangles[0] = uint3(0, 1, 2);
+// CHECK: [[off:%[0-9]+]] = OpIMul %uint %uint_0 %uint_3
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_uint [[indices]] [[off]]
+// CHECK: [[tmp:%[0-9]+]] = OpCompositeExtract %uint [[u3_0_1_2]] 0
+// CHECK:                   OpStore [[ptr]] [[tmp]]
+// CHECK: [[idx:%[0-9]+]] = OpIAdd %uint [[off]] %uint_1
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_uint [[indices]] [[idx]]
+// CHECK: [[tmp:%[0-9]+]] = OpCompositeExtract %uint [[u3_0_1_2]] 1
+// CHECK:                   OpStore [[ptr]] [[tmp]]
+// CHECK: [[idx:%[0-9]+]] = OpIAdd %uint [[off]] %uint_2
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_uint [[indices]] [[idx]]
+// CHECK: [[tmp:%[0-9]+]] = OpCompositeExtract %uint [[u3_0_1_2]] 2
+// CHECK:                   OpStore [[ptr]] [[tmp]]
+
+    verts[0].position = float4(-0.5, 0.5, 0.0, 1.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v4float %gl_Position %int_0
+// CHECK:                   OpStore [[ptr]] [[v4_n05_05_0_1]]
+    verts[0].color = float3(1.0, 0.0, 0.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v3float %out_var_COLOR0 %int_0
+// CHECK:                   OpStore [[ptr]] [[v3_1_0_0]]
+
+    verts[1].position = float4(0.5, 0.5, 0.0, 1.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v4float %gl_Position %int_1
+// CHECK:                   OpStore [[ptr]] [[v4_05_05_0_1]]
+    verts[1].color = float3(0.0, 1.0, 0.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v3float %out_var_COLOR0 %int_1
+// CHECK:                   OpStore [[ptr]] [[v3_0_1_0]]
+
+    verts[2].position = float4(0.0, -0.5, 0.0, 1.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v4float %gl_Position %int_2
+// CHECK:                   OpStore [[ptr]] [[v4_0_n05_0_1]]
+    verts[2].color = float3(0.0, 0.0, 1.0);
+// CHECK: [[ptr:%[0-9]+]] = OpAccessChain %_ptr_Output_v3float %out_var_COLOR0 %int_2
+// CHECK:                   OpStore [[ptr]] [[v3_0_0_1]]
+
+}
+
+[outputtopology("triangle")]
+[numthreads(1, 1, 1)]
+void outie(out indices uint3 triangles[1], out vertices MeshOutput verts[3]) {
+	innie(triangles, verts);
+}