Add matrix orientation attribute in return type (#1770)

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

@@ -347,6 +347,9 @@ private:
   /*TSCS*/unsigned ThreadStorageClassSpec : 2;
   unsigned SCS_extern_in_linkage_spec : 1;
 
+  // default matrix orientation specifier
+  unsigned DefaultMatrixOrientationColumnMajor : 1; // HLSL change
+
   // type-specifier
   /*TSW*/unsigned TypeSpecWidth : 2;
   /*TSC*/unsigned TypeSpecComplex : 2;
@@ -463,6 +466,12 @@ public:
     SCS_extern_in_linkage_spec = Value;
   }
 
+  // HLSL changes begin
+  void SetDefaultMatrixOrientation(bool Value) {
+    DefaultMatrixOrientationColumnMajor = Value;
+  }
+  // HLSL changes end
+
   SourceLocation getStorageClassSpecLoc() const { return StorageClassSpecLoc; }
   SourceLocation getThreadStorageClassSpecLoc() const {
     return ThreadStorageClassSpecLoc;

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

@@ -1361,6 +1361,7 @@ 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,6 +2202,21 @@ 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.

tools/clang/lib/Sema/SemaType.cpp

@@ -4302,6 +4302,23 @@ static TypeSourceInfo *GetFullTypeForDeclarator(TypeProcessingState &state,
   return S.GetTypeSourceInfoForDeclarator(D, T, TInfo);
 }
 
+// HLSL changes begin
+bool Sema::GetAttributedTypeWithMatrixPackingInfo(QualType& T, QualType *RetTy) {
+  if (getLangOpts().HLSL && hlsl::IsHLSLMatType(T)) {
+    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.
 ///
@@ -4319,6 +4336,12 @@ TypeSourceInfo *Sema::GetTypeForDeclarator(Declarator &D, Scope *S) {
   if (D.isPrototypeContext() && getLangOpts().ObjCAutoRefCount)
     inferARCWriteback(state, T);
 
+  // HLSL changes begin
+  QualType NewT;
+  if (GetAttributedTypeWithMatrixPackingInfo(T, &NewT))
+    T = NewT;
+  // HLSL changes end
+
   return GetFullTypeForDeclarator(state, T, ReturnTypeInfo);
 }
 
@@ -4511,6 +4534,16 @@ static AttributeList::Kind getAttrListKind(AttributedType::Kind kind) {
 static void fillAttributedTypeLoc(AttributedTypeLoc TL,
                                   const AttributeList *attrs,
                                   const AttributeList *DeclAttrs = nullptr) {
+
+  // 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;
+  }
+  // HLSL changes end
+
   // DeclAttrs and attrs cannot be both empty.
   assert((attrs || DeclAttrs) &&
          "no type attributes in the expected location!");

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

@@ -0,0 +1,20 @@
+// RUN: %dxc /Tvs_6_0 /Evs_main > %s | FileCheck %s
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 1, float 2.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 2, float 3.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 3, float 4.000000e+00)
+
+#pragma pack_matrix (row_major)
+
+float2x2 GetMatrix()
+{
+ float2x2 mat = {1, 2, 3, 4};
+ return mat;
+}
+
+float4 vs_main() : SV_POSITION
+{
+ float2x2 mat = GetMatrix();
+ return float4(mat[0][0], mat[0][1], mat[1][0], mat[1][1]);
+}

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

@@ -0,0 +1,28 @@
+// RUN: %dxc /Tvs_6_0 /Evs_main > %s | FileCheck %s
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 2.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 1, float 4.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 2, float 6.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 3, float 8.000000e+00)
+
+#pragma pack_matrix (row_major)
+
+float2x2 GetMatrix1()
+{
+ float2x2 mat = {1, 2, 3, 4};
+ return mat;
+}
+
+#pragma pack_matrix (column_major)
+
+float2x2 GetMatrix()
+{
+ float2x2 mat = {1, 2, 3, 4};
+ return mat + GetMatrix1();
+}
+
+float4 vs_main() : SV_POSITION
+{
+ float2x2 mat = GetMatrix();
+ return float4(mat[0][0], mat[0][1], mat[1][0], mat[1][1]);
+}

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

@@ -0,0 +1,28 @@
+// RUN: %dxc /Tvs_6_0 /Evs_main > %s | FileCheck %s
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 1, float 2.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 2, float 3.000000e+00)
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 3, float 4.000000e+00)
+
+#pragma pack_matrix (row_major)
+
+float2x2 GetMatrix(int i)
+{
+ if(i > 0)
+ {
+   float2x2 mat = {1, 2, 3, 4};
+   return mat;
+ }
+ else
+ {
+   float2x2 mat = {2, 3, 3, 4};
+   return mat;
+ }
+}
+
+float4 vs_main() : SV_POSITION
+{
+ float2x2 mat = GetMatrix(1);
+ return float4(mat[0][0], mat[0][1], mat[1][0], mat[1][1]);
+}

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

@@ -0,0 +1,27 @@
+// RUN: %dxc -ast-dump /Tvs_6_0 /Evs_main > %s | FileCheck %s
+
+// CHECK: GetMatrix 'row_major float2x2 (int)'
+
+#pragma pack_matrix (row_major)
+
+float2x2 GetMatrix(int i)
+{
+ if(i > 0)
+ {
+   #pragma pack_matrix (column_major)
+   float2x2 mat = {1, 2, 3, 4};
+   return mat;
+ }
+ else
+ {
+   #pragma pack_matrix (row_major)
+   float2x2 mat = {2, 3, 3, 4};
+   return mat;
+ }
+}
+
+float4 vs_main() : SV_POSITION
+{
+ float2x2 mat = GetMatrix(1);
+ return float4(mat[0][0], mat[0][1], mat[1][0], mat[1][1]);
+}