[spirv] Support for #pragma pack_matrix(). (#2112)

The front-end has taken care of annotating the AST nodes with
the right attribute on QualTypes. We need to read that information
properly and use it in the SPIR-V backend.

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -709,6 +709,17 @@ bool isRowMajorMatrix(const SpirvCodeGenOptions &spvOptions, QualType type) {
   if (hlsl::HasHLSLMatOrientation(type, &attrRowMajor))
     return !attrRowMajor;
 
+  // If it is a templated type the attribute may have been applied to the
+  // underlying type. For example: StructuredBuffer<row_major float2x3>
+  if (const auto *tst = dyn_cast<clang::TemplateSpecializationType>(type)) {
+    if (tst->getNumArgs() >= 1) {
+      auto args = tst->getArgs();
+      auto templateArgument = args[0];
+      auto templateArgumentType = templateArgument.getAsType();
+      return isRowMajorMatrix(spvOptions, templateArgumentType);
+    }
+  }
+
   return !spvOptions.defaultRowMajor;
 }
 

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -515,35 +515,33 @@ const SpirvType *LowerTypeVisitor::lowerResourceType(QualType type,
       rule = getCodeGenOptions().sBufferLayoutRule;
     }
 
+    // Get the underlying resource type.
     const auto s = hlsl::GetHLSLResourceResultType(type);
-    const auto *structType =
-        lowerType(s, rule, /*isRowMajor*/ llvm::None, srcLoc);
-    std::string structName;
-    const auto innerType = hlsl::GetHLSLResourceResultType(type);
-    if (innerType->isStructureType())
-      structName = innerType->getAs<RecordType>()->getDecl()->getName();
-    else
-      structName = getAstTypeName(innerType);
 
+    // If the underlying type is a matrix, check majorness.
+    llvm::Optional<bool> isRowMajor = llvm::None;
+    if (isMxNMatrix(s))
+      isRowMajor = isRowMajorMatrix(spvOptions, type);
+
+    // Lower the underlying type.
+    const auto *structType = lowerType(s, rule, isRowMajor, srcLoc);
+
+    // Calculate memory alignment for the resource.
     uint32_t size = 0, stride = 0;
-    std::tie(std::ignore, size) = alignmentCalc.getAlignmentAndSize(
-        s, rule, /*isRowMajor*/ llvm::None, &stride);
+    std::tie(std::ignore, size) =
+        alignmentCalc.getAlignmentAndSize(s, rule, isRowMajor, &stride);
 
     // We have a runtime array of structures. So:
     // The stride of the runtime array is the size of the struct.
     const auto *raType = spvContext.getRuntimeArrayType(structType, size);
     const bool isReadOnly = (name == "StructuredBuffer");
 
-    // Attach majorness and stride decorations if this is a
-    // *StructuredBuffer<matrix>.
-    llvm::Optional<bool> isRowMajor = llvm::None;
+    // Attach matrix stride decorations if this is a *StructuredBuffer<matrix>.
     llvm::Optional<uint32_t> matrixStride = llvm::None;
-    if (isMxNMatrix(s)) {
-      isRowMajor = isRowMajorMatrix(spvOptions, s);
+    if (isMxNMatrix(s))
       matrixStride = stride;
-    }
 
-    const std::string typeName = "type." + name.str() + "." + structName;
+    const std::string typeName = "type." + name.str() + "." + getAstTypeName(s);
     const auto *valType = spvContext.getStructType(
         {StructType::FieldInfo(raType, /*name*/ "", /*offset*/ 0, matrixStride,
                                isRowMajor)},

tools/clang/test/CodeGenSPIRV/pragma.pack_matrix.hlsl

@@ -0,0 +1,59 @@
+// Run: %dxc -T ps_6_0 -E main
+
+
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat2v3float 0 ColMajor
+#pragma pack_matrix(row_major)
+StructuredBuffer<float2x3> ROSB1;
+
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v2float 0 RowMajor
+#pragma pack_matrix(column_major)
+RWStructuredBuffer<float3x2> RWSB1;
+
+// CHECK: OpMemberDecorate %type_AppendStructuredBuffer_mat4v3float 0 ColMajor
+#pragma pack_matrix(row_major)
+ AppendStructuredBuffer<float4x3> ASB1;
+
+// CHECK: OpMemberDecorate %type_ConsumeStructuredBuffer_mat3v4float 0 RowMajor
+#pragma pack_matrix(column_major)
+ConsumeStructuredBuffer<float3x4> CSB1;
+
+#pragma pack_matrix(row_major)
+struct S {
+// CHECK: OpMemberDecorate %S 0 ColMajor
+               float2x3 sMat0;
+// CHECK: OpMemberDecorate %S 1 ColMajor
+     row_major float2x3 sMat1;
+// CHECK: OpMemberDecorate %S 2 RowMajor
+  column_major float2x3 sMat2;
+};
+RWStructuredBuffer<S> RWSB_S;
+
+#pragma pack_matrix(column_major)
+struct T {
+// CHECK: OpMemberDecorate %T 0 RowMajor
+               float2x3 tMat0;
+// CHECK: OpMemberDecorate %T 1 ColMajor
+  row_major    float2x3 tMat1;
+// CHECK: OpMemberDecorate %T 2 RowMajor
+  column_major float2x3 tMat2;
+};
+RWStructuredBuffer<T> RWSB_T;
+
+// CHECK: OpMemberDecorate %type__Globals 0 ColMajor
+#pragma pack_matrix(row_major)
+float2x2 globalMat_1;
+
+// CHECK: OpMemberDecorate %type__Globals 1 RowMajor
+#pragma pack_matrix(column_major)
+float2x2 globalMat_2;
+
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat3v3float 0 ColMajor
+StructuredBuffer<row_major float3x3> ExplicitOrientationRowMajor;
+
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat4v4float 0 RowMajor
+StructuredBuffer<column_major float4x4> ExplicitOrientationColMajor;
+
+float4 main() : SV_Target {
+  return 0.xxxx;
+}
+

tools/clang/test/CodeGenSPIRV/type.matrix.majorness.zpc.hlsl

@@ -17,21 +17,18 @@
 // CHECK: OpMemberDecorate %type_ConsumeStructuredBuffer_mat3v4float 0 RowMajor
 ConsumeStructuredBuffer<float3x4> CSB1;
 
-// NOTE: The parsed AST does not convey the majorness information for
-// the following cases right now.
-/*
-       StructuredBuffer<row_major float2x3> ROSB2;
-     RWStructuredBuffer<row_major float3x2> RWSB2;
- AppendStructuredBuffer<row_major float4x3> ASB2;
-ConsumeStructuredBuffer<row_major float3x4> CSB2;
+// NOTE: -Zpc does not override explicit matrix orientation specified for the cases below.
 
-       StructuredBuffer<column_major float2x3> ROSB3;
-     RWStructuredBuffer<column_major float3x2> RWSB3;
- AppendStructuredBuffer<column_major float4x3> ASB3;
-ConsumeStructuredBuffer<column_major float3x4> CSB3;
-*/
+// CHECK: OpDecorate %_runtimearr_mat4v4float ArrayStride 64
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat4v4float 0 MatrixStride 16
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat4v4float 0 ColMajor
+     StructuredBuffer<row_major float4x4> ROSB2;
+
+// CHECK: OpDecorate %_runtimearr_mat3v3float ArrayStride 48
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v3float 0 MatrixStride 16
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v3float 0 RowMajor
+RWStructuredBuffer<column_major float3x3> RWSB2;
 
 float4 main() : SV_Target {
-    return ROSB1[0][0][0] // + ROSB2[0][0][0] + ROSB3[0][0][0]
-        ;
+  return ROSB1[0][0][0] + ROSB2[0][0][0] + RWSB2[0][0][0];
 }

tools/clang/test/CodeGenSPIRV/type.matrix.majorness.zpr.hlsl

@@ -45,31 +45,33 @@ struct U {
 RWStructuredBuffer<U> MySBuffer;
 
 // CHECK: OpDecorate %_runtimearr_mat2v3float ArrayStride 32
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat2v3float 0 MatrixStride 16
 // CHECK: OpMemberDecorate %type_StructuredBuffer_mat2v3float 0 ColMajor
        StructuredBuffer<float2x3> ROSB1;
 // CHECK: OpDecorate %_runtimearr_mat3v2float ArrayStride 24
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v2float 0 MatrixStride 8
 // CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v2float 0 ColMajor
      RWStructuredBuffer<float3x2> RWSB1;
 // CHECK: OpDecorate %_runtimearr_mat4v3float ArrayStride 64
+// CHECK: OpMemberDecorate %type_AppendStructuredBuffer_mat4v3float 0 MatrixStride 16
 // CHECK: OpMemberDecorate %type_AppendStructuredBuffer_mat4v3float 0 ColMajor
  AppendStructuredBuffer<float4x3> ASB1;
 // CHECK: OpDecorate %_runtimearr_mat3v4float ArrayStride 48
+// CHECK: OpMemberDecorate %type_ConsumeStructuredBuffer_mat3v4float 0 MatrixStride 16
 // CHECK: OpMemberDecorate %type_ConsumeStructuredBuffer_mat3v4float 0 ColMajor
 ConsumeStructuredBuffer<float3x4> CSB1;
 
-// NOTE: The parsed AST does not convey the majorness information for
-// the following cases right now.
-/*
-       StructuredBuffer<row_major float2x3> ROSB2;
-     RWStructuredBuffer<row_major float3x2> RWSB2;
- AppendStructuredBuffer<row_major float4x3> ASB2;
-ConsumeStructuredBuffer<row_major float3x4> CSB2;
+// NOTE: -Zpr does not override explicit matrix orientation specified for the cases below.
 
-       StructuredBuffer<column_major float2x3> ROSB3;
-     RWStructuredBuffer<column_major float3x2> RWSB3;
- AppendStructuredBuffer<column_major float4x3> ASB3;
-ConsumeStructuredBuffer<column_major float3x4> CSB3;
-*/
+// CHECK: OpDecorate %_runtimearr_mat3v4float_0 ArrayStride 64
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat3v4float 0 MatrixStride 16
+// CHECK: OpMemberDecorate %type_StructuredBuffer_mat3v4float 0 RowMajor
+  StructuredBuffer<column_major float3x4> ROSB2;
+
+// CHECK: OpDecorate %_runtimearr_mat3v3float ArrayStride 48
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v3float 0 MatrixStride 16
+// CHECK: OpMemberDecorate %type_RWStructuredBuffer_mat3v3float 0 ColMajor
+RWStructuredBuffer<row_major float3x3> RWSB2;
 
 float3 main() : A {
   return MySBuffer[0].mat1[1][1];

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1886,4 +1886,7 @@ TEST_F(FileTest, DecorationNoContractionStageVars) {
   runFileTest("decoration.no-contraction.stage-vars.hlsl");
 }
 
+// For pragmas
+TEST_F(FileTest, PragmaPackMatrix) { runFileTest("pragma.pack_matrix.hlsl"); }
+
 } // namespace