[SM6.10] Make DxilStructAnnotation aware of target types

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();
   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.
@@ -343,6 +354,9 @@ class DxilTypeSystem {
   // Determines whether type is a resource or contains a resource
   bool IsResourceContained(llvm::Type *Ty);
 
+  // Determines whether type is a target type or contains a target type
+  bool IsTargetTypeContained(llvm::Type *Ty);
+
 private:
   llvm::Module *m_pModule;
   StructAnnotationMap m_StructAnnotations;

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() {
   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();
@@ -872,6 +890,24 @@ bool DxilTypeSystem::IsResourceContained(llvm::Type *Ty) {
   return false;
 }
 
+bool DxilTypeSystem::IsTargetTypeContained(llvm::Type *Ty) {
+  // 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;
+}
+
 DxilStructTypeIterator::DxilStructTypeIterator(
     llvm::StructType *sTy, DxilStructAnnotation *sAnnotation, unsigned idx)
     : STy(sTy), SAnnotation(sAnnotation), index(idx) {

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 -I %hlsl_headers -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:  %1 = call %dx.types.LinAlgMatrixC9M4N4U0S1 @dx.op.linAlgFillMatrix.mC9M4N4U0S1.f32(i32 -2147483636, float 1.000000e+00)  ; LinAlgFillMatrix(value)
+// CHECK-NOT: @llvm.trap()