Correct reflection stripping detection (#2696)

When stripping reflection data, structures are searched for any member
or submember scalars that are smaller than 32 bits. When one is found
and min precision is being used, the structure is identified as not
being able to be stripped.

In the case of a shader that has a resource that contains a struct which
contains an array which contains a struct, the struct was not being
processed as a struct, but rather as a scalar, which would cause it to
return a size of zero, which was found to be less than 32 bits. This
isn't right if the innermost struct contained a 32 bit integer. Either
way, this resulted in inconsistent results for the parent struct and the
innermost struct. As a result, the annotations were stripped from the
innermost, but left on the outermost. Because the parent struct had
annotations, validation was attempted, but ultimately hit an assert or
caused validation failures because of the innermost struct being
stripped.

By reorganizing the function that searches for 16 bit scalars to look
for a struct only after processing the arrays, the struct is properly
traversed and all annotations are stripped, avoiding the failures.

Rather than determining all aspects of keeping or stripping reflection
annotation based on the individual struct and its contents, all structs
contained in a struct are collected so that they can be excluded from
stripping should their parent struct be found needing of maintaining its
annotations.

Thanks to @tex3d for suggesting (and mostly implementing) this addition

lib/DXIL/DxilModule.cpp

@@ -31,6 +31,7 @@
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include <unordered_set>
 
 using namespace llvm;
@@ -1571,23 +1572,29 @@ StripResourcesReflection(std::vector<std::unique_ptr<TResource>> &vec) {
   return bChanged;
 }
 
-static bool ResourceTypeRequiresTranslation(const StructType* Ty) {
+// Return true if any members or components of struct <Ty> contain
+// scalars of less than 32 bits or are matrices, in which case translation is required
+typedef llvm::SmallSetVector<const StructType*, 4> SmallStructSetVector;
+static bool ResourceTypeRequiresTranslation(const StructType * Ty, SmallStructSetVector & containedStructs) {
   if (Ty->getName().startswith("class.matrix."))
     return true;
+  bool bResult = false;
+  containedStructs.insert(Ty);
   for (auto eTy : Ty->elements()) {
-    if (StructType *structTy = dyn_cast<StructType>(eTy)) {
-      if (ResourceTypeRequiresTranslation(structTy))
-        return true;
-    }
+    // Skip past all levels of sequential types to test their elements
     SequentialType *seqTy;
     while ((seqTy = dyn_cast<SequentialType>(eTy))) {
       eTy = seqTy->getElementType();
     }
-    if (eTy->getScalarSizeInBits() < 32) {
-      return true;
+    // Recursively call this function again to process internal structs
+    if (StructType *structTy = dyn_cast<StructType>(eTy)) {
+      if (ResourceTypeRequiresTranslation(structTy, containedStructs))
+        bResult = true;
+    } else if (eTy->getScalarSizeInBits() < 32) { // test scalar sizes
+      bResult = true;
     }
   }
-  return false;
+  return bResult;
 }
 
 bool DxilModule::StripReflection() {
@@ -1614,11 +1621,19 @@ bool DxilModule::StripReflection() {
   {
     // We must preserve struct annotations for resources containing min-precision types,
     // since they have not yet been converted for legacy layout.
-    SmallVector<const StructType*, 4> structsToRemove;
+    // Keep all structs contained in any we must keep.
+    SmallStructSetVector structsToKeep;
+    SmallStructSetVector structsToRemove;
     for (auto &item : m_pTypeSystem->GetStructAnnotationMap()) {
-      if (!ResourceTypeRequiresTranslation(item.first))
-        structsToRemove.emplace_back(item.first);
+      SmallStructSetVector containedStructs;
+      if (!ResourceTypeRequiresTranslation(item.first, containedStructs))
+        structsToRemove.insert(item.first);
+      else
+        structsToKeep.insert(containedStructs.begin(), containedStructs.end());
     }
+
+    for (auto Ty : structsToKeep)
+      structsToRemove.remove(Ty);
     for (auto Ty : structsToRemove) {
       m_pTypeSystem->GetStructAnnotationMap().erase(Ty);
     }

tools/clang/test/HLSLFileCheck/d3dreflect/cbuffer_alignment_16.hlsl

@@ -0,0 +1,51 @@
+// RUN: %dxc -T lib_6_3 %s | FileCheck %s
+
+// CHECK: ; Buffer Definitions:
+// CHECK: ;
+// CHECK: ; cbuffer cbuf
+// CHECK: ; {
+// CHECK: ;
+// CHECK: ;   struct cbuf
+// CHECK: ;   {
+// CHECK: ;
+// CHECK: ;       struct struct.Agg
+// CHECK: ;       {
+// CHECK: ;
+// CHECK: ;           struct struct.Value
+// CHECK: ;           {
+// CHECK: ;
+// CHECK: ;               uint v;                               ; Offset:    0
+// CHECK: ;
+// CHECK: ;           } value[1];;                              ; Offset:    0
+// CHECK: ;
+// CHECK: ;           min16ui fodder;                           ; Offset:    4
+// CHECK: ;
+// CHECK: ;       } aggie;                                      ; Offset:    0
+// CHECK: ;
+// CHECK: ;
+// CHECK: ;   } cbuf;                                           ; Offset:    0 Size:     8
+// CHECK: ;
+// CHECK: ; }
+
+
+// Test Cbuffer validation likely to cause mistaken overlaps
+struct Value { uint v; }; // This may be stripped because it has nothing below 32 bits
+
+struct Agg {
+
+  Value value[1];
+  min16uint fodder; // This will cause this struct to be preserved
+};
+
+cbuffer cbuf : register(b1)
+{
+  Agg aggie;
+}
+
+RWBuffer<int> Out : register(u0);
+
+[shader("raygeneration")]
+void main()
+{
+  Out[0] = aggie.value[0].v;
+}

tools/clang/test/HLSLFileCheck/d3dreflect/cbuffer_alignment_matrix.hlsl

@@ -0,0 +1,52 @@
+// RUN: %dxc -T lib_6_3 %s | FileCheck %s
+
+// CHECK: ; Buffer Definitions:
+// CHECK: ;
+// CHECK: ; cbuffer cbuf
+// CHECK: ; {
+// CHECK: ;
+// CHECK: ;   struct cbuf
+// CHECK: ;   {
+// CHECK: ;
+// CHECK: ;       struct struct.Agg
+// CHECK: ;       {
+// CHECK: ;
+// CHECK: ;           column_major uint2x2 mat;                 ; Offset:    0
+// CHECK: ;           struct struct.Value
+// CHECK: ;           {
+// CHECK: ;
+// CHECK: ;               uint v;                               ; Offset:    32
+// CHECK: ;
+// CHECK: ;           } value[1];;                              ; Offset:    32
+// CHECK: ;
+// CHECK: ;           uint fodder;                              ; Offset:    36
+// CHECK: ;
+// CHECK: ;       } aggie;                                      ; Offset:    0
+// CHECK: ;
+// CHECK: ;
+// CHECK: ;   } cbuf;                                           ; Offset:    0 Size:     40
+// CHECK: ;
+// CHECK: ; }
+
+
+// Test Cbuffer validation likely to cause mistaken overlaps
+struct Value { uint v; }; // This will be stripped because it has nothing below 32 bits
+
+struct Agg {
+  uint2x2 mat; // This will cause the struct to be 
+  Value value[1];
+  uint fodder;// This will give the struct something to overlap
+};
+
+cbuffer cbuf : register(b1)
+{
+  Agg aggie;
+}
+
+RWBuffer<int> Out : register(u0);
+
+[shader("raygeneration")]
+void main()
+{
+  Out[0] = aggie.value[0].v;
+}