Fix hang/write-back drop for inout structs bitcast

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(BitCastInst *BCI) {
+  // Unused bitcast may be leftover from temporary memcpy
+  if (BCI->use_empty())
+    return true;
+
+  Type *DstTy = BCI->getType();
+  Type *SrcTy = BCI->getOperand(0)->getType();
+
+  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(BC))
+        return MarkUnsafe(Info, User);
       isSafeForScalarRepl(BC, Offset, Info);
     } else if (GetElementPtrInst *GEPI = dyn_cast<GetElementPtrInst>(User)) {
       uint64_t GEPOffset = Offset;
@@ -2687,7 +2722,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 +4433,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 +4500,6 @@ class SROA_Parameter_HLSL : public ModulePass {
     while (!WorkList.empty()) {
       Function *F = WorkList.front();
       WorkList.pop_front();
-      RewriteBitcastWithIdenticalStructs(F);
       createFlattenedFunction(F);
     }
 
@@ -4609,28 +4644,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 +4674,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() {

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/similar_layout_structs/cast_as_func_inout_param-strictudt.hlsl

@@ -0,0 +1,40 @@
+// RUN: %dxc -T ps_6_0 -HV 2021 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2021 -DSTATIC_GLOBAL %s | FileCheck %s -check-prefix=STATIC_GLOBAL
+
+// Validate that passing a struct to an inout parameter of a differently laid
+// out struct does not hang the compiler.
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+
+// STATIC_GLOBAL: error: Unsupported cast between struct types with different layouts.
+
+struct S {
+  float2 v;
+};
+
+struct T {
+  float x, y;
+};
+
+void f(inout S s) {
+  s.v.x = 1;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f((S)g);
+  return g.x;
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f((S)t);
+  return t.x;
+}
+
+#endif

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

@@ -0,0 +1,39 @@
+// RUN: %dxc -T ps_6_0 -HV 2018 %s | FileCheck %s
+// RUN: %dxc -T ps_6_0 -HV 2018 -DSTATIC_GLOBAL %s | FileCheck %s -check-prefix=STATIC_GLOBAL
+
+// Validate that passing a struct to an inout parameter of a differently laid
+// out struct does not hang the compiler.
+
+// CHECK: call void @dx.op.storeOutput.f32(i32 5, i32 0, i32 0, i8 0, float 1.000000e+00)
+// STATIC_GLOBAL: error: Unsupported cast between struct types with different layouts.
+
+struct S {
+  float2 v;
+};
+
+struct T {
+  float x, y;
+};
+
+void f(inout S s) {
+  s.v.x = 1;
+}
+
+#ifdef STATIC_GLOBAL
+
+static T g;
+
+float4 main() : SV_Target {
+  f(g);
+  return g.x;
+}
+
+#else
+
+float4 main() : SV_Target {
+  T t = (T)0;
+  f(t);
+  return t.x;
+}
+
+#endif