Fix invalid IR from scalarrepl-param-hlsl in ReplaceConstantWithInst (#6556)

ReplaceConstantWithInst(C, V) replaces uses of C in the current function
with V. If such a use C is an instruction I, the it replaces uses of C
in I with V. However, this function did not make sure to only perform
this replacement if V dominates I. As a result, it may end up replacing
uses of C in instructions before the definition of V.

The fix is to lazily compute the dominator tree in
ReplaceConstantWithInst so that we can guard the replacement with that
dominance check.

Author: amaiorano

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -3271,15 +3271,34 @@ bool SROA_Helper::DoScalarReplacement(GlobalVariable *GV,
   return true;
 }
 
-static void ReplaceConstantWithInst(Constant *C, Value *V,
+// Replaces uses of constant C in the current function
+// with V, when those uses are dominated by V.
+// Returns true if it was completely replaced.
+static bool ReplaceConstantWithInst(Constant *C, Value *V,
                                     IRBuilder<> &Builder) {
+  bool bReplacedAll = true;
   Function *F = Builder.GetInsertBlock()->getParent();
+  Instruction *VInst = dyn_cast<Instruction>(V);
+  // Lazily calculate dominance
+  DominatorTree DT;
+  bool Calculated = false;
+  auto Dominates = [&](llvm::Instruction *Def, llvm::Instruction *User) {
+    if (!Calculated) {
+      DT.recalculate(*F);
+      Calculated = true;
+    }
+    return DT.dominates(Def, User);
+  };
+
   for (auto it = C->user_begin(); it != C->user_end();) {
     User *U = *(it++);
     if (Instruction *I = dyn_cast<Instruction>(U)) {
       if (I->getParent()->getParent() != F)
         continue;
-      I->replaceUsesOfWith(C, V);
+      if (VInst && Dominates(VInst, I))
+        I->replaceUsesOfWith(C, V);
+      else
+        bReplacedAll = false;
     } else {
       // Skip unused ConstantExpr.
       if (U->user_empty())
@@ -3288,10 +3307,12 @@ static void ReplaceConstantWithInst(Constant *C, Value *V,
       Instruction *Inst = CE->getAsInstruction();
       Builder.Insert(Inst);
       Inst->replaceUsesOfWith(C, V);
-      ReplaceConstantWithInst(CE, Inst, Builder);
+      if (!ReplaceConstantWithInst(CE, Inst, Builder))
+        bReplacedAll = false;
     }
   }
   C->removeDeadConstantUsers();
+  return bReplacedAll;
 }
 
 static void ReplaceUnboundedArrayUses(Value *V, Value *Src) {
@@ -3531,15 +3552,17 @@ static bool ReplaceMemcpy(Value *V, Value *Src, MemCpyInst *MC,
       } else {
         // Replace Constant with a non-Constant.
         IRBuilder<> Builder(MC);
-        ReplaceConstantWithInst(C, Src, Builder);
+        if (!ReplaceConstantWithInst(C, Src, Builder))
+          return false;
       }
     } else {
       // Try convert special pattern for cbuffer which copy array of float4 to
       // array of float.
       if (!tryToReplaceCBVec4ArrayToScalarArray(V, TyV, Src, TySrc, MC, DL)) {
         IRBuilder<> Builder(MC);
         Src = Builder.CreateBitCast(Src, V->getType());
-        ReplaceConstantWithInst(C, Src, Builder);
+        if (!ReplaceConstantWithInst(C, Src, Builder))
+          return false;
       }
     }
   } else {
@@ -5449,9 +5472,9 @@ void SROA_Parameter_HLSL::flattenArgument(
         if (Ty->isPointerTy())
           Ty = Ty->getPointerElementType();
         unsigned size = DL.getTypeAllocSize(Ty);
-#if 0  // HLSL Change
+#if 0 // HLSL Change
         DIExpression *DDIExp = DIB.createBitPieceExpression(debugOffset, size);
-#else  // HLSL Change
+#else // HLSL Change
         Type *argTy = Arg->getType();
         if (argTy->isPointerTy())
           argTy = argTy->getPointerElementType();

tools/clang/test/DXC/Passes/ScalarReplHLSL/scalarrepl-param-hlsl-const-to-local-and-back.hlsl

@@ -0,0 +1,45 @@
+// RUN: not %dxc -T ps_6_2 %s 2>&1 | FileCheck %s
+
+// Validate that copying from static array to local, then back to static
+// array does not crash the compiler. This was resulting in an invalid
+// ReplaceConstantWithInst from ScalarReplAggregatesHLSL, which would
+// result in referenced deleted instruction in a later pass.
+
+// CHECK: error: Loop must have break.
+
+static int arr1[10] = (int[10])0;
+static int arr2[10] = (int[10])0;
+static float result = 0;
+ByteAddressBuffer buff : register(t0);
+
+void foo() {
+  int i = 0;
+  if (buff.Load(0u)) {
+    return;
+  }
+  arr2[i] = arr1[i];
+  result = float(arr1[0]);
+}
+
+struct tint_symbol {
+  float4 value : SV_Target0;
+};
+
+float main_inner() {
+  foo();
+  bool cond = false;
+  while (true) {
+    if (cond) { break; }
+  }
+  int arr1_copy[10] = arr1; // constant to local
+  arr1 = arr1_copy; // local to constant
+  foo();
+  return result;
+}
+
+tint_symbol main() {
+  float inner_result = main_inner();
+  tint_symbol wrapper_result = (tint_symbol)0;
+  wrapper_result.value.x = inner_result;
+  return wrapper_result;
+}