Apply #pragma pack_matrix at the right token-level granularity (#1806)

This ensures that #pragma pack_matrix is applied as close to as possible to the same token-level granularity as FXC. Also fix a crash when combining `uniform` or `out` with `row_major` or `column_major`.

tools/clang/include/clang/AST/HlslTypes.h

@@ -371,8 +371,8 @@ bool IsHLSLVecMatType(clang::QualType);
 bool IsHLSLVecType(clang::QualType type);
 bool IsHLSLMatType(clang::QualType type);
 clang::QualType GetElementTypeOrType(clang::QualType type);
-bool HasHLSLMatOrientation(clang::QualType type, bool *pIsRowMajor);
-bool HasHLSLUNormSNorm(clang::QualType type, bool *pIsSNorm);
+bool HasHLSLMatOrientation(clang::QualType type, bool *pIsRowMajor = nullptr);
+bool HasHLSLUNormSNorm(clang::QualType type, bool *pIsSNorm = nullptr);
 bool IsHLSLInputPatchType(clang::QualType type);
 bool IsHLSLOutputPatchType(clang::QualType type);
 bool IsHLSLPointStreamType(clang::QualType type);

tools/clang/include/clang/Sema/DeclSpec.h

@@ -347,8 +347,14 @@ private:
   /*TSCS*/unsigned ThreadStorageClassSpec : 2;
   unsigned SCS_extern_in_linkage_spec : 1;
 
-  // default matrix orientation specifier
-  unsigned DefaultMatrixOrientationColumnMajor : 1; // HLSL change
+  // HLSL Change Start
+  // Whether the default matrix pack is defined at the point
+  // of the declaration. This is false when rewriting
+  // and no #pragma pack_matrix have been encountered yet.
+  unsigned HasDefaultMatrixPack : 1;
+  // Default matrix pack at the point of the declaration
+  unsigned DefaultMatrixPackRowMajor : 1;
+  // HLSL Change End
 
   // type-specifier
   /*TSW*/unsigned TypeSpecWidth : 2;
@@ -434,6 +440,8 @@ public:
     : StorageClassSpec(SCS_unspecified),
       ThreadStorageClassSpec(TSCS_unspecified),
       SCS_extern_in_linkage_spec(false),
+      HasDefaultMatrixPack(false), // HLSL Change
+      DefaultMatrixPackRowMajor(false), // HLSL Change
       TypeSpecWidth(TSW_unspecified),
       TypeSpecComplex(TSC_unspecified),
       TypeSpecSign(TSS_unspecified),
@@ -467,8 +475,15 @@ public:
   }
 
   // HLSL changes begin
-  void SetDefaultMatrixOrientation(bool Value) {
-    DefaultMatrixOrientationColumnMajor = Value;
+  bool TryGetDefaultMatrixPackRowMajor(bool& rowMajor) const {
+    if (!HasDefaultMatrixPack) return false;
+    rowMajor = DefaultMatrixPackRowMajor;
+    return true;
+  }
+
+  void SetDefaultMatrixPackRowMajor(bool Value) {
+    HasDefaultMatrixPack = true;
+    DefaultMatrixPackRowMajor = Value;
   }
   // HLSL changes end
 

tools/clang/include/clang/Sema/Sema.h

@@ -333,10 +333,12 @@ public:
   LangOptions::PragmaMSPointersToMembersKind
       MSPointerToMemberRepresentationMethod;
 
-  // HLSL Change Begin - pragma pack_matrix.
-  // Add both row/col to identify the default case which no pragma.
-  bool PackMatrixRowMajorPragmaOn = false; // True when \#pragma pack_matrix(row_major) on.
-  bool PackMatrixColMajorPragmaOn = false; // True when \#pragma pack_matrix(column_major) on.
+  // HLSL Change Begin
+  // The HLSL rewriter doesn't define a default matrix pack,
+  // so we must preserve the lack of annotations to avoid changing semantics.
+  bool HasDefaultMatrixPack = false;
+  // Uses of #pragma pack_matrix change the default pack.
+  bool DefaultMatrixPackRowMajor = false;
   // HLSL Change End.
 
   enum PragmaVtorDispKind {
@@ -1361,7 +1363,6 @@ public:
   };
 
 private:
-  bool GetAttributedTypeWithMatrixPackingInfo(QualType& T, QualType *RetTy); // HLSL change
   bool RequireCompleteTypeImpl(SourceLocation Loc, QualType T,
                            TypeDiagnoser &Diagnoser);
 

tools/clang/lib/Parse/ParseDecl.cpp

@@ -2202,21 +2202,6 @@ Parser::DeclGroupPtrTy Parser::ParseDeclGroup(ParsingDeclSpec &DS,
     return DeclGroupPtrTy();
   }
 
-  // HLSL changes begin
-  // Initialize the default matrix orientation.
-  // TODO: Take into account global matrix orientation when initializing default
-  // matrix orientation. That would also mean adding an accessor method to read
-  // default matrix orientation and use it in all applicable places.
-  if (!Parser::Actions.PackMatrixRowMajorPragmaOn &&
-      !Parser::Actions.PackMatrixColMajorPragmaOn) {
-    DS.SetDefaultMatrixOrientation(true);
-  } else if (Parser::Actions.PackMatrixRowMajorPragmaOn) {
-    DS.SetDefaultMatrixOrientation(false);
-  } else if (Parser::Actions.PackMatrixColMajorPragmaOn) {
-    DS.SetDefaultMatrixOrientation(true);
-  }
-  // HLSL changes end
-
   // HLSL Change Starts: change global variables that will be in constant buffer to be constant by default
   // Global variables that are groupshared, static, or typedef
   // will not be part of constant buffer and therefore should not be const by default.
@@ -3493,6 +3478,13 @@ void Parser::ParseDeclarationSpecifiers(DeclSpec &DS,
       if (DS.hasTypeSpecifier() && DS.hasTagDefinition())
         goto DoneWithDeclSpec;
 
+      // HLSL Change Starts
+      // Remember the current state of the default matrix orientation,
+      // since it can change between any two tokens with #pragma pack_matrix
+      if (Parser::Actions.HasDefaultMatrixPack)
+        DS.SetDefaultMatrixPackRowMajor(Parser::Actions.DefaultMatrixPackRowMajor);
+      // HLSL Change Ends
+
       if (Tok.getAnnotationValue()) {
         ParsedType T = getTypeAnnotation(Tok);
         isInvalid = DS.SetTypeSpecType(DeclSpec::TST_typename, Loc, PrevSpec,
@@ -3603,12 +3595,20 @@ void Parser::ParseDeclarationSpecifiers(DeclSpec &DS,
           Actions.isCurrentClassName(*Tok.getIdentifierInfo(), getCurScope()) &&
           isConstructorDeclarator(/*Unqualified*/true))
         goto DoneWithDeclSpec;
-      // HLSL Change start - modify TypeRep for unsigned vectors/matrix
+
+      // HLSL Change Starts
+      // Modify TypeRep for unsigned vectors/matrix
       QualType qt = TypeRep.get();
       QualType newType = ApplyTypeSpecSignToParsedType(&Actions, qt, DS.getTypeSpecSign(), Loc);
       isInvalid = DS.SetTypeSpecType(DeclSpec::TST_typename, Loc, PrevSpec,
                                      DiagID, ParsedType::make(newType), Policy);
-      // HLSL Change end
+
+      // Remember the current state of the default matrix orientation,
+      // since it can change between any two tokens with #pragma pack_matrix
+      if (Parser::Actions.HasDefaultMatrixPack)
+        DS.SetDefaultMatrixPackRowMajor(Parser::Actions.DefaultMatrixPackRowMajor);
+      // HLSL Change Ends
+
       if (isInvalid)
         break;
 

tools/clang/lib/Sema/SemaAttr.cpp

@@ -266,13 +266,10 @@ void Sema::ActOnPragmaPack(PragmaPackKind Kind, IdentifierInfo *Name,
 }
 
 void Sema::ActOnPragmaPackMatrix(bool bRowMajor, SourceLocation PragmaLoc) {
-  if (bRowMajor) {
-    PackMatrixRowMajorPragmaOn = true;
-    PackMatrixColMajorPragmaOn = false;
-  } else {
-    PackMatrixRowMajorPragmaOn = false;
-    PackMatrixColMajorPragmaOn = true;
-  }
+  // Once we've encountered one #pragma pack_matrix, we have a well-defined
+  // default orientation for the rest of the program, even if we're rewriting.
+  HasDefaultMatrixPack = true;
+  DefaultMatrixPackRowMajor = bRowMajor;
 }
 
 void Sema::ActOnPragmaMSStruct(PragmaMSStructKind Kind) { 

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -11279,56 +11279,6 @@ void Sema::TransferUnusualAttributes(Declarator &D, NamedDecl *NewDecl) {
         D.UnusualAnnotations.size()));
     D.UnusualAnnotations.clear();
   }
-  // pragma pack_matrix.
-  // Do this for struct member also.
-  if (ValueDecl *VD = dyn_cast<ValueDecl>(NewDecl)) {
-    QualType Ty = VD->getType();
-    QualType EltTy = Ty;
-    while (EltTy->isArrayType()) {
-      EltTy = EltTy->getAsArrayTypeUnsafe()->getElementType();
-    }
-    if (hlsl::IsHLSLMatType(EltTy)) {
-      bool bRowMajor = false;
-      if (!hlsl::HasHLSLMatOrientation(EltTy, &bRowMajor)) {
-        if (PackMatrixColMajorPragmaOn || PackMatrixRowMajorPragmaOn) {
-          // Add major.
-          QualType NewEltTy = Context.getAttributedType(
-              PackMatrixRowMajorPragmaOn
-                  ? AttributedType::attr_hlsl_row_major
-                  : AttributedType::attr_hlsl_column_major,
-              EltTy, EltTy);
-
-          QualType NewTy = NewEltTy;
-          if (Ty->isArrayType()) {
-            // Build new array type.
-            SmallVector<const ArrayType *, 2> arrayTys;
-            while (EltTy->isArrayType()) {
-              const ArrayType *AT = EltTy->getAsArrayTypeUnsafe();
-              arrayTys.emplace_back(AT);
-            }
-            for (auto rit = arrayTys.rbegin(); rit != arrayTys.rend(); rit++) {
-              // Create array type with NewTy.
-              const ArrayType *AT = *rit;
-              if (const ConstantArrayType *CAT =
-                      dyn_cast<ConstantArrayType>(AT)) {
-                NewTy = Context.getConstantArrayType(
-                    NewTy, CAT->getSize(), CAT->getSizeModifier(),
-                    CAT->getIndexTypeCVRQualifiers());
-              } else if (const IncompleteArrayType *IAT =
-                             dyn_cast<IncompleteArrayType>(AT)) {
-                NewTy = Context.getIncompleteArrayType(NewTy, IAT->getSizeModifier(),
-                    IAT->getIndexTypeCVRQualifiers());
-              } else {
-                DXASSERT(false, "");
-              }
-            }
-          }
-          // Update Type.
-          VD->setType(NewTy);
-        }
-      }
-    }
-  }
 }
 
 /// Checks whether a usage attribute is compatible with those seen so far and

tools/clang/lib/Sema/SemaType.cpp

@@ -4302,24 +4302,6 @@ static TypeSourceInfo *GetFullTypeForDeclarator(TypeProcessingState &state,
   return S.GetTypeSourceInfoForDeclarator(D, T, TInfo);
 }
 
-// HLSL changes begin
-bool Sema::GetAttributedTypeWithMatrixPackingInfo(QualType& T, QualType *RetTy) {
-  bool isRowMajorAttributed;
-  if (getLangOpts().HLSL && hlsl::IsHLSLMatType(T) && !hlsl::HasHLSLMatOrientation(T, &isRowMajorAttributed)) {
-    if (PackMatrixColMajorPragmaOn || PackMatrixRowMajorPragmaOn) {
-      *RetTy = Context.getAttributedType(
-        PackMatrixRowMajorPragmaOn
-        ? AttributedType::attr_hlsl_row_major
-        : AttributedType::attr_hlsl_column_major,
-        T, T);
-      return true;
-    }
-  }
-
-  return false;
-}
-// HLSL changes end
-
 /// GetTypeForDeclarator - Convert the type for the specified
 /// declarator to Type instances.
 ///
@@ -4338,9 +4320,20 @@ TypeSourceInfo *Sema::GetTypeForDeclarator(Declarator &D, Scope *S) {
     inferARCWriteback(state, T);
 
   // HLSL changes begin
-  QualType NewT;
-  if (GetAttributedTypeWithMatrixPackingInfo(T, &NewT))
-    T = NewT;
+  // If there is no explicit pack orientation on matrix types, but there is file-level
+  // default orientation set by #pragma pack_matrix, apply it here.
+  // There is no default if rewriting (in the absence of #pragma pack_matrix), since
+  // it is agnostic to default orientation and we want to preserve the lack of annotation.
+  // For codegen, it'd be nice to annotate everything here, but it causes error
+  // messages to have pack orientation added to types, so we handle it through
+  // the codegen option's default packing orientation flag.
+  bool defaultRowMajor;
+  if (getLangOpts().HLSL && hlsl::IsHLSLMatType(T) && !hlsl::HasHLSLMatOrientation(T)
+    && D.getDeclSpec().TryGetDefaultMatrixPackRowMajor(defaultRowMajor)) {
+    AttributedType::Kind AttributeKind = defaultRowMajor
+      ? AttributedType::attr_hlsl_row_major : AttributedType::attr_hlsl_column_major;
+    T = Context.getAttributedType(AttributeKind, T, T);
+  }
   // HLSL changes end
 
   return GetFullTypeForDeclarator(state, T, ReturnTypeInfo);
@@ -4538,11 +4531,9 @@ static void fillAttributedTypeLoc(AttributedTypeLoc TL,
 
   // HLSL changes begin
   // Don't fill the location info for matrix orientation attributes
-  if (!attrs && !DeclAttrs) {
-    if (TL.getAttrKind() == AttributedType::attr_hlsl_row_major ||
-        TL.getAttrKind() == AttributedType::attr_hlsl_column_major)
-      return;
-  }
+  if (TL.getAttrKind() == AttributedType::attr_hlsl_row_major ||
+      TL.getAttrKind() == AttributedType::attr_hlsl_column_major)
+    return;
   // HLSL changes end
 
   // DeclAttrs and attrs cannot be both empty.

tools/clang/test/CodeGenHLSL/quick-test/matrix_orientation_pragma_granularity.hlsl

@@ -0,0 +1,31 @@
+// RUN: %dxc /T vs_6_0 /E main > %s | FileCheck %s
+
+// Tests the exact place at which #pragma pack_matrix takes effect.
+
+#pragma pack_matrix(column_major)
+
+typedef
+#pragma pack_matrix(row_major)
+int2x2
+// With FXC, we could place the #pragma pack_matrix(column_major) here
+// and still get the type be row_major. This not easy to replicate
+// in DXC because the parser looks ahead one token to see if the
+// type is followed by '::', which causes the execution of a #pragma
+// following the 'int2x2' one token early, but it is highly unlikely that this
+// would be a backwards compatibility issue.
+i22
+#pragma pack_matrix(column_major)
+;
+
+void main(out i22 mat : OUT)
+{
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 11)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 1, i32 12)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 1, i8 0, i32 21)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 1, i8 1, i32 22)
+    mat = i22(11, 12, 21, 22);
+
+    // FXC output, for reference:
+    // mov o0.xy, l(11,12,0,0)
+    // mov o1.xy, l(21,22,0,0)
+}

tools/clang/test/CodeGenHLSL/quick-test/matrix_orientation_pragma_granularity_template_syntax.hlsl

@@ -0,0 +1,29 @@
+// RUN: %dxc /T vs_6_0 /E main > %s | FileCheck %s
+
+#pragma pack_matrix(column_major)
+
+typedef matrix<int, 2, 2
+#pragma pack_matrix(row_major)
+>
+// With FXC, we could place the #pragma pack_matrix(column_major) here
+// and still get the type be row_major. This not easy to replicate
+// in DXC because the parser looks ahead one token to see if the
+// templated type is followed by '::', which causes the execution of a #pragma
+// following the '>' one token early, but it is highly unlikely that this
+// would be a backwards compatibility issue.
+i22
+#pragma pack_matrix(column_major)
+;
+
+void main(out i22 mat : OUT)
+{
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 11)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 1, i32 12)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 1, i8 0, i32 21)
+    // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 1, i8 1, i32 22)
+    mat = i22(11, 12, 21, 22);
+
+    // FXC output, for reference:
+    // mov o0.xy, l(11,12,0,0)
+    // mov o1.xy, l(21,22,0,0)
+}

tools/clang/test/HLSL/rewriter/correct_rewrites/matrix-pack-orientation_gold.hlsl

@@ -0,0 +1,10 @@
+// Rewrite unchanged result:
+void default_noPragma(int1x1 m);
+void rowMajorAttribute_noPragma(row_major int1x1 m);
+void columnMajorAttribute_noPragma(column_major int1x1 m);
+void default_pragmaRowMajor(row_major int1x1 m);
+void rowMajorAttribute_pragmaRowMajor(row_major int1x1 m);
+void columnMajorAttribute_pragmaRowMajor(column_major int1x1 m);
+void default_pragmaColumnMajor(column_major int1x1 m);
+void rowMajorAttribute_pragmaColumnMajor(row_major int1x1 m);
+void columnMajorAttribute_pragmaColumnMajor(column_major int1x1 m);

tools/clang/test/HLSL/rewriter/matrix-pack-orientation.hlsl

@@ -0,0 +1,15 @@
+// Test that the semantics of matrix pack orientations are preserved through rewriting.
+
+void default_noPragma(int1x1 m); // The lack of pack orientation annotation should be preserved
+void rowMajorAttribute_noPragma(row_major int1x1 m);
+void columnMajorAttribute_noPragma(column_major int1x1 m);
+
+#pragma pack_matrix(row_major)
+void default_pragmaRowMajor(int1x1 m); // This should get a row_major attribute added
+void rowMajorAttribute_pragmaRowMajor(row_major int1x1 m);
+void columnMajorAttribute_pragmaRowMajor(column_major int1x1 m);
+
+#pragma pack_matrix(column_major)
+void default_pragmaColumnMajor(int1x1 m); // This should get a column_major attribute added
+void rowMajorAttribute_pragmaColumnMajor(row_major int1x1 m);
+void columnMajorAttribute_pragmaColumnMajor(column_major int1x1 m);

tools/clang/unittests/HLSL/RewriterTest.cpp

@@ -60,6 +60,7 @@ public:
   TEST_METHOD(RunIndexingOperator);
   TEST_METHOD(RunIntrinsicExamples);
   TEST_METHOD(RunMatrixAssignments);
+  TEST_METHOD(RunMatrixPackOrientation);
   TEST_METHOD(RunMatrixSyntax);
   TEST_METHOD(RunPackReg);
   TEST_METHOD(RunScalarAssignments);
@@ -332,6 +333,10 @@ TEST_F(RewriterTest, RunMatrixAssignments) {
     CheckVerifiesHLSL(L"rewriter\\matrix-assignments_noerr.hlsl", L"rewriter\\correct_rewrites\\matrix-assignments_gold.hlsl");
 }
 
+TEST_F(RewriterTest, RunMatrixPackOrientation) {
+  CheckVerifiesHLSL(L"rewriter\\matrix-pack-orientation.hlsl", L"rewriter\\correct_rewrites\\matrix-pack-orientation_gold.hlsl");
+}
+
 TEST_F(RewriterTest, RunMatrixSyntax) {
     CheckVerifiesHLSL(L"rewriter\\matrix-syntax_noerr.hlsl", L"rewriter\\correct_rewrites\\matrix-syntax_gold.hlsl");
 }