Fix hang/write-back drop for inout structs bitcast

docs/ReleaseNotes.md

@@ -33,6 +33,10 @@ line upon naming the release. Refer to previous for appropriate section names.
 - Fix a crash generating DXIL from sources containing a dynamic resource heap
   access that was discarded. Identified during development of SPIR-V support for
   [descriptor heaps](https://github.com/microsoft/DirectXShaderCompiler/pull/8517#discussion_r3752113078).
+- Fixed a dropped `inout` write-back when an argument is passed to an `inout`
+  parameter of a struct type with an identical layout.
+- Fixed a compiler hang when an argument is passed to an `inout` parameter of a
+  struct type with a different layout.
 
 #### HLSL Language
 

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -1533,6 +1533,39 @@ static bool isUDTIntrinsicArg(CallInst *CI, unsigned OpIdx) {
   return false;
 }
 
+/// isSafeBitCastForScalarRepl - Check if a bitcast can be handled for scalar
+/// replacement.  A struct to struct pointer cast is only rewritable when the
+/// destination is reachable through the leading elements of the source, or
+/// when both structs have an identical layout.
+static bool isSafeBitCastForScalarRepl(BitCastOperator *BC) {
+  // Unused bitcast may be leftover from temporary memcpy
+  if (BC->use_empty())
+    return true;
+
+  Type *DstTy = BC->getDestTy();
+  Type *SrcTy = BC->getSrcTy();
+
+  if (!DstTy->isPointerTy() || !SrcTy->isPointerTy())
+    return true;
+
+  StructType *DstST = dyn_cast<StructType>(DstTy->getPointerElementType());
+  StructType *SrcST = dyn_cast<StructType>(SrcTy->getPointerElementType());
+
+  // A non-struct destination is rewritten per llvm.lifetime.* intrinsic user
+  // A non-struct source never reaches the struct to struct rewrite
+  if (!DstST || !SrcST)
+    return true;
+
+  for (StructType *ST = SrcST; ST && ST->getNumElements();) {
+    Type *EltTy = ST->getElementType(0);
+    if (EltTy == DstST)
+      return true;
+    ST = dyn_cast<StructType>(EltTy);
+  }
+
+  return SrcST->isLayoutIdentical(DstST);
+}
+
 /// isSafeForScalarRepl - Check if instruction I is a safe use with regard to
 /// performing scalar replacement of alloca AI.  The results are flagged in
 /// the Info parameter.  Offset indicates the position within AI that is
@@ -1548,6 +1581,8 @@ void isSafeForScalarRepl(Instruction *I, uint64_t Offset, AllocaInfo &Info) {
     Instruction *User = cast<Instruction>(U.getUser());
 
     if (BitCastInst *BC = dyn_cast<BitCastInst>(User)) {
+      if (!isSafeBitCastForScalarRepl(cast<BitCastOperator>(BC)))
+        return MarkUnsafe(Info, User);
       isSafeForScalarRepl(BC, Offset, Info);
     } else if (GetElementPtrInst *GEPI = dyn_cast<GetElementPtrInst>(User)) {
       uint64_t GEPOffset = Offset;
@@ -1669,6 +1704,24 @@ bool isSafeAllocaToScalarRepl(AllocaInst *AI) {
 
   return true;
 }
+
+/// isSafeGlobalToScalarRepl - Check if a global can be broken down into
+/// elements.  Recursively walks the pointer producing users, which for a global
+/// may be constant expressions rather than instructions, and returns false when
+/// any of them cannot be rewritten.
+bool isSafeGlobalToScalarRepl(Value *V) {
+  for (User *U : V->users()) {
+    if (BitCastOperator *BC = dyn_cast<BitCastOperator>(U)) {
+      if (!isSafeBitCastForScalarRepl(BC))
+        return false;
+    } else if (!isa<GEPOperator>(U)) {
+      continue;
+    }
+    if (!isSafeGlobalToScalarRepl(U))
+      return false;
+  }
+  return true;
+}
 } // namespace
 
 namespace {
@@ -2668,7 +2721,7 @@ void SROA_Helper::RewriteBitCast(BitCastInst *BCI) {
 
   bool bTypeMatch = false;
   unsigned level = 0;
-  while (SrcST) {
+  while (SrcST && SrcST->getNumElements()) {
     level++;
     Type *EltTy = SrcST->getElementType(0);
     if (EltTy == DstST) {
@@ -2687,7 +2740,10 @@ void SROA_Helper::RewriteBitCast(BitCastInst *BCI) {
       BCI->eraseFromParent();
       return;
     }
-    assert(0 && "Type mismatch.");
+    dxilutil::EmitErrorOnInstruction(
+        BCI, "Unsupported cast between struct types with different layouts.");
+    BCI->replaceAllUsesWith(UndefValue::get(BCI->getType()));
+    BCI->eraseFromParent();
     return;
   }
 
@@ -4395,8 +4451,6 @@ class SROA_Parameter_HLSL : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   explicit SROA_Parameter_HLSL() : ModulePass(ID) {}
   StringRef getPassName() const override { return "SROA Parameter HLSL"; }
-  static void RewriteBitcastWithIdenticalStructs(Function *F);
-  static void RewriteBitcastWithIdenticalStructs(BitCastInst *BCI);
   static bool DeleteSimpleStoreOnlyAlloca(AllocaInst *AI);
   static bool IsSimpleStoreOnlyAlloca(AllocaInst *AI);
 
@@ -4464,7 +4518,6 @@ class SROA_Parameter_HLSL : public ModulePass {
     while (!WorkList.empty()) {
       Function *F = WorkList.front();
       WorkList.pop_front();
-      RewriteBitcastWithIdenticalStructs(F);
       createFlattenedFunction(F);
     }
 
@@ -4609,28 +4662,6 @@ INITIALIZE_PASS(SROA_Parameter_HLSL, "scalarrepl-param-hlsl",
                 "Scalar Replacement of Aggregates HLSL (parameters)", false,
                 false)
 
-void SROA_Parameter_HLSL::RewriteBitcastWithIdenticalStructs(Function *F) {
-  if (F->isDeclaration())
-    return;
-  // Gather list of bitcast involving src and dest structs with identical layout
-  std::vector<BitCastInst *> worklist;
-  for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I) {
-    if (BitCastInst *BCI = dyn_cast<BitCastInst>(&*I)) {
-      Type *DstTy = BCI->getDestTy();
-      Type *SrcTy = BCI->getSrcTy();
-      if (ArePointersToStructsOfIdenticalLayouts(DstTy, SrcTy))
-        worklist.push_back(BCI);
-    }
-  }
-
-  // Replace bitcast involving src and dest structs with identical layout
-  while (!worklist.empty()) {
-    BitCastInst *BCI = worklist.back();
-    worklist.pop_back();
-    RewriteBitcastWithIdenticalStructs(BCI);
-  }
-}
-
 bool SROA_Parameter_HLSL::IsSimpleStoreOnlyAlloca(AllocaInst *AI) {
   if (!AI->getAllocatedType()->isSingleValueType())
     return false;
@@ -4661,23 +4692,6 @@ bool SROA_Parameter_HLSL::DeleteSimpleStoreOnlyAlloca(AllocaInst *AI) {
   return true;
 }
 
-void SROA_Parameter_HLSL::RewriteBitcastWithIdenticalStructs(BitCastInst *BCI) {
-  StructType *srcStTy =
-      cast<StructType>(BCI->getSrcTy()->getPointerElementType());
-  StructType *destStTy =
-      cast<StructType>(BCI->getDestTy()->getPointerElementType());
-  Value *srcPtr = BCI->getOperand(0);
-  IRBuilder<> AllocaBuilder(
-      dxilutil::FindAllocaInsertionPt(BCI->getParent()->getParent()));
-  AllocaInst *destPtr = AllocaBuilder.CreateAlloca(destStTy);
-  IRBuilder<> InstBuilder(BCI);
-  std::vector<unsigned> idxlist = {0};
-  CopyElementsOfStructsWithIdenticalLayout(InstBuilder, destPtr, srcPtr,
-                                           srcStTy, idxlist);
-  BCI->replaceAllUsesWith(destPtr);
-  BCI->eraseFromParent();
-}
-
 /// DeleteDeadInstructions - Erase instructions on the DeadInstrs list,
 /// recursively including all their operands that become trivially dead.
 void SROA_Parameter_HLSL::DeleteDeadInstructions() {
@@ -6645,7 +6659,8 @@ class LowerStaticGlobalIntoAlloca : public ModulePass {
       } else {
         EltTy = dxilutil::GetArrayEltTy(EltTy);
         // Lower static [array of] resources
-        if (dxilutil::IsHLSLObjectType(EltTy) || EltTy == handleTy) {
+        if (dxilutil::IsHLSLObjectType(EltTy) || EltTy == handleTy ||
+            !isSafeGlobalToScalarRepl(&GV)) {
           staticGVs.emplace_back(&GV);
         }
       }

tools/clang/test/HLSLFileCheck/hlsl/types/cast/identical_layout_structs/cast_as_func_inout_param-strictudt.hlsl

@@ -0,0 +1,42 @@
+// RUN: %dxc -T ps_6_0 -HV 2021 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2021 -DSTATIC_GLOBAL %s | FileCheck %s
+
+// Validate that the copy-out of an inout argument survives a cast between two
+// structs of identical layout.
+
+// 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)
+
+struct S {
+  float a;
+  int b;
+};
+
+struct T {
+  float c;
+  int d;
+};
+
+void f(inout S s) {
+  s.a = 1;
+  s.b = 2;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f((S)g);
+  return float4(g.c, g.d, 0, 0);
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f((S)t);
+  return float4(t.c, t.d, 0, 0);
+}
+
+#endif

tools/clang/test/HLSLFileCheck/hlsl/types/cast/identical_layout_structs/cast_as_func_inout_param.hlsl

@@ -0,0 +1,42 @@
+// RUN: %dxc -T ps_6_0 -HV 2018 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2018 -DSTATIC_GLOBAL %s | FileCheck %s
+
+// Validate that the copy-out of an inout argument survives a cast between two
+// structs of identical layout.
+
+// 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)
+
+struct S {
+  float a;
+  int b;
+};
+
+struct T {
+  float c;
+  int d;
+};
+
+void f(inout S s) {
+  s.a = 1;
+  s.b = 2;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f(g);
+  return float4(g.c, g.d, 0, 0);
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f(t);
+  return float4(t.c, t.d, 0, 0);
+}
+
+#endif

tools/clang/test/HLSLFileCheck/hlsl/types/cast/identical_layout_structs/cast_as_func_inout_param_empty_base-strictudt.hlsl

@@ -0,0 +1,46 @@
+// RUN: %dxc -T ps_6_0 -HV 2021 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2021 -DSTATIC_GLOBAL %s | FileCheck %s
+
+// Validate that a cast between two structs of identical layout whose leading
+// members bottom out in an empty struct does not crash the compiler.
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+
+struct Empty {};
+
+struct Inner {
+  Empty e;
+};
+
+struct S {
+  Inner i;
+  float a;
+};
+
+struct T {
+  Inner i;
+  float a;
+};
+
+void f(inout S s) {
+  s.a = 1;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f((S)g);
+  return g.a;
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f((S)t);
+  return t.a;
+}
+
+#endif

tools/clang/test/HLSLFileCheck/hlsl/types/cast/identical_layout_structs/cast_as_func_inout_param_empty_base.hlsl

@@ -0,0 +1,46 @@
+// RUN: %dxc -T ps_6_0 -HV 2018 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2018 -DSTATIC_GLOBAL %s | FileCheck %s
+
+// Validate that a cast between two structs of identical layout whose leading
+// members bottom out in an empty struct does not crash the compiler.
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+
+struct Empty {};
+
+struct Inner {
+  Empty e;
+};
+
+struct S {
+  Inner i;
+  float a;
+};
+
+struct T {
+  Inner i;
+  float a;
+};
+
+void f(inout S s) {
+  s.a = 1;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f(g);
+  return g.a;
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f(t);
+  return t.a;
+}
+
+#endif

tools/clang/test/HLSLFileCheck/hlsl/types/cast/similar_layout_structs/cast_as_func_inout_param-lib.hlsl

@@ -0,0 +1,27 @@
+// RUN: %dxc -T lib_6_3 %s | FileCheck %s
+
+// Validate that passing a static global to an inout parameter of a differently
+// laid out struct is diagnosed rather than hanging the compiler, when the
+// global is used from an exported library function.
+
+// CHECK: error: Unsupported cast between struct types with different layouts.
+
+struct S {
+  float2 v;
+};
+
+struct T {
+  float x, y;
+};
+
+static T g;
+
+void f(inout S s) {
+  s.v.x = 1;
+}
+
+export float4 fn(float y) {
+  g.y = y;
+  f((S)g);
+  return float4(g.x, g.y, 0, 0);
+}