[SM6.10] Make DxilStructAnnotation aware of target types (#8282)

Update `DxilStructAnnotation` to be aware of target types. Otherwise the annotation would declare a struct with a single LinAlg Matrix type as empty, leading to invalid codegen.

Declaration names match the style of the surrounding code. 

Fixes #8279

Author: hekota

include/dxc/DXIL/DxilTypeSystem.h

@@ -148,10 +148,12 @@ class DxilStructAnnotation {
   void SetCBufferSize(unsigned size);
   void MarkEmptyStruct();
   bool IsEmptyStruct();
-  // Since resources don't take real space, IsEmptyBesidesResources
-  // determines if the structure is empty or contains only resources.
-  bool IsEmptyBesidesResources();
+  // Since resources and target types don't take real space,
+  // IsEmptyBesidesResourcesAndTargetTypes() determines if the structure is
+  // empty or contains only resources or target types.
+  bool IsEmptyBesidesResourcesAndTargetTypes() const;
   bool ContainsResources() const;
+  bool ContainsTargetTypes() const;
 
   // For template args, GetNumTemplateArgs() will return 0 if not a template
   unsigned GetNumTemplateArgs() const;
@@ -174,6 +176,15 @@ class DxilStructAnnotation {
     False,
     Only
   } m_ResourcesContained = HasResources::False;
+
+  void SetContainsTargetTypes();
+  // HasTargetTypes::Only will be set on MarkEmptyStruct() when
+  // HasTargetTypes::True
+  enum class HasTargetTypes {
+    True,
+    False,
+    Only
+  } m_TargetTypesContained = HasTargetTypes::False;
 };
 
 /// Use this class to represent type annotation for DXR payload field.
@@ -341,7 +352,10 @@ class DxilTypeSystem {
   void SetMinPrecision(bool bMinPrecision);
 
   // Determines whether type is a resource or contains a resource
-  bool IsResourceContained(llvm::Type *Ty);
+  bool IsResourceContained(llvm::Type *Ty) const;
+
+  // Determines whether type is a target type or contains a target type
+  bool IsTargetTypeContained(llvm::Type *Ty) const;
 
 private:
   llvm::Module *m_pModule;

lib/DXIL/DxilTypeSystem.cpp

@@ -236,14 +236,17 @@ void DxilStructAnnotation::SetCBufferSize(unsigned size) {
 void DxilStructAnnotation::MarkEmptyStruct() {
   if (m_ResourcesContained == HasResources::True)
     m_ResourcesContained = HasResources::Only;
+  else if (m_TargetTypesContained == HasTargetTypes::True)
+    m_TargetTypesContained = HasTargetTypes::Only;
   else
     m_FieldAnnotations.clear();
 }
 bool DxilStructAnnotation::IsEmptyStruct() {
   return m_FieldAnnotations.empty();
 }
-bool DxilStructAnnotation::IsEmptyBesidesResources() {
+bool DxilStructAnnotation::IsEmptyBesidesResourcesAndTargetTypes() const {
   return m_ResourcesContained == HasResources::Only ||
+         m_TargetTypesContained == HasTargetTypes::Only ||
          m_FieldAnnotations.empty();
 }
 
@@ -256,6 +259,14 @@ bool DxilStructAnnotation::ContainsResources() const {
   return m_ResourcesContained != HasResources::False;
 }
 
+void DxilStructAnnotation::SetContainsTargetTypes() {
+  if (m_TargetTypesContained == HasTargetTypes::False)
+    m_TargetTypesContained = HasTargetTypes::True;
+}
+bool DxilStructAnnotation::ContainsTargetTypes() const {
+  return m_TargetTypesContained != HasTargetTypes::False;
+}
+
 // For template args, GetNumTemplateArgs() will return 0 if not a template
 unsigned DxilStructAnnotation::GetNumTemplateArgs() const {
   return (unsigned)m_TemplateAnnotations.size();
@@ -393,6 +404,13 @@ void DxilTypeSystem::FinishStructAnnotation(DxilStructAnnotation &SA) {
       SA.SetContainsResources();
   }
 
+  // Update target type containment
+  for (unsigned i = 0; i < SA.GetNumFields() && !SA.ContainsTargetTypes();
+       i++) {
+    if (IsTargetTypeContained(ST->getElementType(i)))
+      SA.SetContainsTargetTypes();
+  }
+
   // Mark if empty
   if (SA.GetCBufferSize() == 0)
     SA.MarkEmptyStruct();
@@ -854,20 +872,36 @@ void DxilTypeSystem::SetMinPrecision(bool bMinPrecision) {
   m_LowPrecisionMode = mode;
 }
 
-bool DxilTypeSystem::IsResourceContained(llvm::Type *Ty) {
+bool DxilTypeSystem::IsResourceContained(llvm::Type *Ty) const {
   // strip pointer/array
   if (Ty->isPointerTy())
     Ty = Ty->getPointerElementType();
   if (Ty->isArrayTy())
     Ty = Ty->getArrayElementType();
 
-  if (auto ST = dyn_cast<StructType>(Ty)) {
-    if (dxilutil::IsHLSLResourceType(Ty)) {
+  if (auto *ST = dyn_cast<StructType>(Ty)) {
+    if (dxilutil::IsHLSLResourceType(Ty))
       return true;
-    } else if (auto SA = GetStructAnnotation(ST)) {
+    else if (auto *SA = GetStructAnnotation(ST))
       if (SA->ContainsResources())
         return true;
-    }
+  }
+  return false;
+}
+
+bool DxilTypeSystem::IsTargetTypeContained(llvm::Type *Ty) const {
+  // strip pointer/array
+  if (Ty->isPointerTy())
+    Ty = Ty->getPointerElementType();
+  if (Ty->isArrayTy())
+    Ty = Ty->getArrayElementType();
+
+  if (auto *ST = dyn_cast<StructType>(Ty)) {
+    if (dxilutil::IsHLSLKnownTargetType(Ty))
+      return true;
+    else if (auto *SA = GetStructAnnotation(ST))
+      if (SA->ContainsTargetTypes())
+        return true;
   }
   return false;
 }

lib/HLSL/DxilContainerReflection.cpp

@@ -1278,7 +1278,8 @@ HRESULT CShaderReflectionType::Initialize(
 
       // There is no annotation for empty structs
       unsigned int fieldCount = 0;
-      if (structAnnotation && !structAnnotation->IsEmptyBesidesResources())
+      if (structAnnotation &&
+          !structAnnotation->IsEmptyBesidesResourcesAndTargetTypes())
         fieldCount = type->getStructNumElements();
 
       // The DXBC reflection info computes `Columns` for a

tools/clang/test/CodeGenDXIL/hlsl/linalg/matrix-target-type-in-struct.hlsl

@@ -0,0 +1,22 @@
+// REQUIRES: dxil-1-10
+// RUN: %dxc -T cs_6_10 -enable-16bit-types %s | FileCheck %s
+
+using MyHandleT = __builtin_LinAlgMatrix [[__LinAlgMatrix_Attributes(9, 4, 4, 0, 1)]];
+
+class MyMatrix {
+  MyHandleT handle;
+
+  static MyMatrix Splat(float Val) {
+    MyMatrix Result;
+    __builtin_LinAlg_FillMatrix(Result.handle, Val);
+    return Result;
+  }
+};
+
+[numthreads(4, 4, 4)]
+void main() {
+  MyMatrix MatA = MyMatrix::Splat(1.0f);
+}
+
+// CHECK: call %dx.types.LinAlgMatrixC9M4N4U0S1 @dx.op.linAlgFillMatrix.mC9M4N4U0S1.f32(i32 -2147483636, float 1.000000e+00)  ; LinAlgFillMatrix(value)
+// CHECK-NOT: @llvm.trap()