Fix CB Var Size reported by reflection (#2125)

Used to report next member minus this offset, which would include cb
padding in the size.  This is not correct.

Now we calculate the size based on the packing rules for the type.
- basic types are:
  (((elements * rows) - 1) * row stride) + cols * element size
  where row stride is 16, or 32 if 64-bit and cols > 2.
- for struct:
  struct size = offset of last field + size of last field
  struct size += (elements - 1) * 16-byte aligned struct size

Turned on Size check in reflection tests since it should be correct now.

lib/HLSL/DxilContainerReflection.cpp

@@ -347,6 +347,7 @@ class CShaderReflectionType : public ID3D12ShaderReflectionType
 {
 protected:
   D3D12_SHADER_TYPE_DESC              m_Desc;
+  UINT                                m_SizeInCBuffer;
   std::string                         m_Name;
   std::vector<StringRef>              m_MemberNames;
   std::vector<CShaderReflectionType*> m_MemberTypes;
@@ -382,6 +383,8 @@ public:
   bool CheckEqual(_In_ CShaderReflectionType *pOther) {
     return m_Identity == pOther->m_Identity;
   }
+
+  UINT GetCBufferSize() { return m_SizeInCBuffer; }
 };
 
 class CShaderReflectionVariable : public ID3D12ShaderReflectionVariable
@@ -747,6 +750,13 @@ HRESULT CShaderReflectionType::Initialize(
   m_Desc.Columns = 0;
   m_Desc.Elements = 0;
   m_Desc.Members = 0;
+  m_SizeInCBuffer = 0;
+
+  // Used for calculating size later
+  unsigned cbRows = 1;
+  unsigned cbCols = 1;
+  unsigned cbCompSize = 4;    // or 8 for 64-bit types.
+  unsigned cbRowStride = 16;  // or 32 if 64-bit and cols > 2.
 
   // Extract offset relative to parent.
   // Note: the `baseOffset` is used in the case where the type in
@@ -818,6 +828,7 @@ HRESULT CShaderReflectionType::Initialize(
     break;
 
   case hlsl::DXIL::ComponentType::I64:
+    cbCompSize = 8;
 #ifdef DBG
     OutputDebugStringA("DxilContainerReflection.cpp: warning: component of type 'I64' being reflected as if 'I32'\n");
 #endif
@@ -827,6 +838,7 @@ HRESULT CShaderReflectionType::Initialize(
     break;
 
   case hlsl::DXIL::ComponentType::U64:
+    cbCompSize = 8;
 #ifdef DBG
     OutputDebugStringA("DxilContainerReflection.cpp: warning: component of type 'U64' being reflected as if 'U32'\n");
 #endif
@@ -852,6 +864,7 @@ HRESULT CShaderReflectionType::Initialize(
   case hlsl::DXIL::ComponentType::F64:
   case hlsl::DXIL::ComponentType::SNormF64:
   case hlsl::DXIL::ComponentType::UNormF64:
+    cbCompSize = 8;
     componentType = D3D_SVT_DOUBLE;
     m_Name = "double";
     break;
@@ -891,6 +904,12 @@ HRESULT CShaderReflectionType::Initialize(
     m_Desc.Rows = matrixAnnotation.Rows;
     m_Desc.Columns = matrixAnnotation.Cols;
     m_Name += std::to_string(matrixAnnotation.Rows) + "x" + std::to_string(matrixAnnotation.Cols);
+
+    cbRows = m_Desc.Rows;
+    cbCols = m_Desc.Columns;
+    if (m_Desc.Class == D3D_SVC_MATRIX_COLUMNS) {
+      std::swap(cbRows, cbCols);
+    }
   }
   else if( type->isVectorTy() )
   {
@@ -904,6 +923,9 @@ HRESULT CShaderReflectionType::Initialize(
     m_Desc.Columns = type->getVectorNumElements();
 
     m_Name += std::to_string(type->getVectorNumElements());
+
+    cbRows = m_Desc.Rows;
+    cbCols = m_Desc.Columns;
   }
   else if( type->isStructTy() )
   {
@@ -943,6 +965,8 @@ HRESULT CShaderReflectionType::Initialize(
       // `struct` type from the fields (see below)
       UINT columnCounter = 0;
 
+      CShaderReflectionType *fieldReflectionType = nullptr;
+
       for(unsigned int ff = 0; ff < fieldCount; ++ff)
       {
         DxilFieldAnnotation& fieldAnnotation = structAnnotation->GetFieldAnnotation(ff);
@@ -957,7 +981,7 @@ HRESULT CShaderReflectionType::Initialize(
           continue;
         }
 
-        CShaderReflectionType *fieldReflectionType = new CShaderReflectionType();
+        fieldReflectionType = new CShaderReflectionType();
         allTypes.push_back(std::unique_ptr<CShaderReflectionType>(fieldReflectionType));
 
         fieldReflectionType->Initialize(M, fieldType, fieldAnnotation, 0, allTypes);
@@ -978,6 +1002,15 @@ HRESULT CShaderReflectionType::Initialize(
 
       m_Desc.Columns = columnCounter;
 
+      if (fieldReflectionType) {
+        // Set our size based on the last fields offset + size:
+        m_SizeInCBuffer = fieldReflectionType->m_Desc.Offset + fieldReflectionType->m_SizeInCBuffer;
+        if (m_Desc.Elements > 1) {
+          unsigned alignedSize = ((m_SizeInCBuffer + 15) & ~0xF);
+          m_SizeInCBuffer += (m_Desc.Elements - 1) * alignedSize;
+        }
+      }
+
       // Because we might have skipped fields during enumeration,
       // the `Members` count in the description might not be the same
       // as the field count of the original LLVM type.
@@ -1016,9 +1049,26 @@ HRESULT CShaderReflectionType::Initialize(
       m_Name = "dword";
       break;
     }
+
+    cbRows = 1;
+    cbCols = 1;
   }
   // TODO: are there other cases to be handled?
 
+  // Compute our cbuffer size for member reflection
+  switch (m_Desc.Class) {
+  case D3D_SVC_SCALAR:
+  case D3D_SVC_MATRIX_COLUMNS:
+  case D3D_SVC_MATRIX_ROWS:
+  case D3D_SVC_VECTOR:
+    if (m_Desc.Elements > 1)
+      cbRows = cbRows * m_Desc.Elements;
+    if (cbCompSize > 4 && cbCols > 2)
+      cbRowStride = 32;
+    m_SizeInCBuffer = cbRowStride * (cbRows - 1) + cbCompSize * cbCols;
+    break;
+  }
+
   m_Desc.Name = m_Name.c_str();
 
   return S_OK;
@@ -1050,7 +1100,6 @@ void CShaderReflectionConstantBuffer::Initialize(
     return;
 
   m_Desc.Variables = ST->getNumContainedTypes();
-  unsigned lastIndex = ST->getNumContainedTypes() - 1;
 
   for (unsigned i = 0; i < ST->getNumContainedTypes(); ++i) {
     DxilFieldAnnotation &fieldAnnotation = annotation->GetFieldAnnotation(i);
@@ -1068,14 +1117,7 @@ void CShaderReflectionConstantBuffer::Initialize(
 
     VarDesc.Name = fieldAnnotation.GetFieldName().c_str();
     VarDesc.StartOffset = fieldAnnotation.GetCBufferOffset();
-    if (i < lastIndex) {
-      DxilFieldAnnotation &nextFieldAnnotation =
-          annotation->GetFieldAnnotation(i + 1);
-      VarDesc.Size = nextFieldAnnotation.GetCBufferOffset() - fieldAnnotation.GetCBufferOffset();
-    }
-    else {
-      VarDesc.Size = CB.GetSize() - fieldAnnotation.GetCBufferOffset();
-    }
+    VarDesc.Size = pVarType->GetCBufferSize();
     Var.Initialize(this, &VarDesc, pVarType, pDefaultValue);
     m_Variables.push_back(Var);
   }

tools/clang/test/CodeGenHLSL/batch/misc/d3dreflect/cb_sizes.hlsl

@@ -0,0 +1,238 @@
+// RUN: %dxc -E main -T vs_6_0 %s | %D3DReflect %s | FileCheck %s
+
+// Verify CB variable sizes align with expectations.
+// This also tests some matrix, struct, and array cases that may
+// have not been covered sufficiently elsewhere.
+
+#if 0
+// CHECK:  Constant Buffers:
+// CHECK-NEXT:    ID3D12ShaderReflectionConstantBuffer:
+// CHECK-NEXT:      D3D12_SHADER_BUFFER_DESC: Name: CB
+// CHECK-NEXT:        Type: D3D_CT_CBUFFER
+// CHECK-NEXT:        Size: 2176
+// CHECK-NEXT:        uFlags: 0
+// CHECK-NEXT:        Num Variables: 8
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: A
+// CHECK-NEXT:            Size: 8
+// CHECK-NEXT:            StartOffset: 0
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: B
+// CHECK-NEXT:            Size: 16
+// CHECK-NEXT:            StartOffset: 16
+// CHECK-NEXT:            uFlags: 0x2
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: D
+// CHECK-NEXT:            Size: 8
+// CHECK-NEXT:            StartOffset: 32
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: double
+// CHECK-NEXT:              Class: D3D_SVC_SCALAR
+// CHECK-NEXT:              Type: D3D_SVT_DOUBLE
+// CHECK-NEXT:              Elements: 0
+// CHECK-NEXT:              Rows: 1
+// CHECK-NEXT:              Columns: 1
+// CHECK-NEXT:              Members: 0
+// CHECK-NEXT:              Offset: 0
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: f3x2
+// CHECK-NEXT:            Size: 28
+// CHECK-NEXT:            StartOffset: 48
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:              Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:              Type: D3D_SVT_FLOAT
+// CHECK-NEXT:              Elements: 0
+// CHECK-NEXT:              Rows: 3
+// CHECK-NEXT:              Columns: 2
+// CHECK-NEXT:              Members: 0
+// CHECK-NEXT:              Offset: 0
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: f3x2_row
+// CHECK-NEXT:            Size: 40
+// CHECK-NEXT:            StartOffset: 80
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:              Class: D3D_SVC_MATRIX_ROWS
+// CHECK-NEXT:              Type: D3D_SVT_FLOAT
+// CHECK-NEXT:              Elements: 0
+// CHECK-NEXT:              Rows: 3
+// CHECK-NEXT:              Columns: 2
+// CHECK-NEXT:              Members: 0
+// CHECK-NEXT:              Offset: 0
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: d3x4
+// CHECK-NEXT:            Size: 120
+// CHECK-NEXT:            StartOffset: 128
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: double3x4
+// CHECK-NEXT:              Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:              Type: D3D_SVT_DOUBLE
+// CHECK-NEXT:              Elements: 0
+// CHECK-NEXT:              Rows: 3
+// CHECK-NEXT:              Columns: 4
+// CHECK-NEXT:              Members: 0
+// CHECK-NEXT:              Offset: 0
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: s1
+// CHECK-NEXT:            Size: 312
+// CHECK-NEXT:            StartOffset: 256
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: S1
+// CHECK-NEXT:              Class: D3D_SVC_STRUCT
+// CHECK-NEXT:              Type: D3D_SVT_VOID
+// CHECK-NEXT:              Elements: 0
+// CHECK-NEXT:              Rows: 1
+// CHECK-NEXT:              Columns: 39
+// CHECK-NEXT:              Members: 5
+// CHECK-NEXT:              Offset: 0
+// CHECK-NEXT:            {
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: int
+// CHECK-NEXT:                  Class: D3D_SVC_SCALAR
+// CHECK-NEXT:                  Type: D3D_SVT_INT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 1
+// CHECK-NEXT:                  Columns: 1
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 0
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_FLOAT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 2
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 16
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: double3x4
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_DOUBLE
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 4
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 48
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: int2x1
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_INT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 2
+// CHECK-NEXT:                  Columns: 1
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 168
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_ROWS
+// CHECK-NEXT:                  Type: D3D_SVT_FLOAT
+// CHECK-NEXT:                  Elements: 3
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 2
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 176
+
+// CHECK:        ID3D12ShaderReflectionVariable:
+// CHECK-NEXT:          D3D12_SHADER_VARIABLE_DESC: Name: s1_arr
+// CHECK-NEXT:            Size: 1592
+// CHECK-NEXT:            StartOffset: 576
+// CHECK-NEXT:            uFlags: 0
+// CHECK-NEXT:            DefaultValue: <nullptr>
+// CHECK-NEXT:          ID3D12ShaderReflectionType:
+// CHECK-NEXT:            D3D12_SHADER_TYPE_DESC: Name: S1
+// CHECK-NEXT:              Class: D3D_SVC_STRUCT
+// CHECK-NEXT:              Type: D3D_SVT_VOID
+// CHECK-NEXT:              Elements: 5
+// CHECK-NEXT:              Rows: 1
+// CHECK-NEXT:              Columns: 39
+// CHECK-NEXT:              Members: 5
+// CHECK-NEXT:              Offset: 0
+// CHECK-NEXT:            {
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: int
+// CHECK-NEXT:                  Class: D3D_SVC_SCALAR
+// CHECK-NEXT:                  Type: D3D_SVT_INT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 1
+// CHECK-NEXT:                  Columns: 1
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 0
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_FLOAT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 2
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 16
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: double3x4
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_DOUBLE
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 4
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 48
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: int2x1
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_COLUMNS
+// CHECK-NEXT:                  Type: D3D_SVT_INT
+// CHECK-NEXT:                  Elements: 0
+// CHECK-NEXT:                  Rows: 2
+// CHECK-NEXT:                  Columns: 1
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 168
+// CHECK-NEXT:              ID3D12ShaderReflectionType:
+// CHECK-NEXT:                D3D12_SHADER_TYPE_DESC: Name: float3x2
+// CHECK-NEXT:                  Class: D3D_SVC_MATRIX_ROWS
+// CHECK-NEXT:                  Type: D3D_SVT_FLOAT
+// CHECK-NEXT:                  Elements: 3
+// CHECK-NEXT:                  Rows: 3
+// CHECK-NEXT:                  Columns: 2
+// CHECK-NEXT:                  Members: 0
+// CHECK-NEXT:                  Offset: 176
+
+#endif
+
+struct S1 {
+  int i;
+  float3x2 f3x2;
+  double3x4 d3x4;
+  int2x1 i2x1;
+  row_major float3x2 f3x2_row[3];
+};
+
+cbuffer CB {
+  float2 A;
+  float4 B;
+  double D;
+  float3x2 f3x2;
+  row_major float3x2 f3x2_row;
+  double3x4 d3x4;
+  S1 s1;
+  S1 s1_arr[5];
+}
+
+float4 main() : OUT {
+  return B;
+}

tools/clang/unittests/HLSL/DxilContainerTest.cpp

@@ -286,8 +286,8 @@ public:
           D3D12_SHADER_VARIABLE_DESC baseConst = variableMap[testConst.Name];
           VERIFY_ARE_EQUAL(testConst.uFlags, baseConst.uFlags);
           VERIFY_ARE_EQUAL(testConst.StartOffset, baseConst.StartOffset);
-          // TODO: enalbe size cmp.
-          //VERIFY_ARE_EQUAL(testConst.Size, baseConst.Size);
+
+          VERIFY_ARE_EQUAL(testConst.Size, baseConst.Size);
 
           ID3D12ShaderReflectionType* pTestType = pTestConst->GetType();
           VERIFY_IS_NOT_NULL(pTestType);