Fix issues with resource array aliasing support (#3810)

SROA: Skip arrays of object/matrix as well.

Fix issues with ReplaceUseOfZeroInit*

  - Fix ReplaceUseOfZeroInitPostDom bailing on use of memcpy dest.
    When memcpy not in entry block, any use of dest after memcpy would abort
    replacement, as if it was unsafe to replace dest with src.  But uses
    of dest dominated by memcpy *are* safe to replace.
  - Fix ReplaceUseOfZeroInit* misuse of post-dom.
    Post-dom was used to detect whether it was safe to replace uses of dest
    before the memcpy with zero when dest was zeroinitialized.  But post-dom
    is not the right way to tell if this is safe.  It is unsafe if any uses
    *could* follow the memcpy.  So the new test is to gather a set of blocks
    that could be reachable from the successors of the memcpy block (where
    memcpy ends the block because we split it there).  If so, it is not safe
    to replace the use with zeroinitializer, and because it was not dominated
    by the memcpy either, it's not safe to replace dest with src here either,
    so memcpy replacement must abort.
  - Re-merge blocks after splitting in ReplaceUseOfZeroInitBeforeDef.
    Restore CFG after ReplaceUseOfZeroInitBeforeDef, rather than leaving it
    in more of a mess to be cleaned up much later.
    This has the side benefit of preserving more trivial entry-block
    replacement opportunities.

DFE should apply to internal functions in lib target as well.

  - Also, iterating over RemoveUnusedFunctions will allow removal of more
    functions only called by other internal functions being removed.

Add LowerStaticGlobalIntoAlloca before SROA for LowerMemcpy opportunities

  - Existing location after SROA must also be there for non-resource aggregates
  - Don't preserve zeroinitializer when moving static GV to alloca for objects
  - Minor opt in LowerStaticGlobalIntoAlloca: skip fn decls more cheaply

Author: tex3d

lib/DXIL/DxilUtil.cpp

@@ -121,7 +121,8 @@ bool RemoveUnusedFunctions(Module &M, Function *EntryFunc,
   for (auto &F : M.functions()) {
     if (&F == EntryFunc || &F == PatchConstantFunc)
       continue;
-    if (F.isDeclaration() || !IsLib) {
+    if (F.isDeclaration() || !IsLib ||
+        F.hasInternalLinkage()) {
       if (F.user_empty())
         deadList.emplace_back(&F);
     }

lib/HLSL/HLDeadFunctionElimination.cpp

@@ -35,8 +35,11 @@ class HLDeadFunctionElimination : public ModulePass {
       Function *EntryFunc = HLM.GetEntryFunction();
       Function *PatchConstantFunc = HLM.GetPatchConstantFunction();
 
-      return dxilutil::RemoveUnusedFunctions(M, EntryFunc, PatchConstantFunc,
-                                             IsLib);
+      bool bChanged = false;
+      while (dxilutil::RemoveUnusedFunctions(M, EntryFunc, PatchConstantFunc,
+                                             IsLib))
+        bChanged = true;
+      return bChanged;
     }
 
     return false;

lib/Transforms/IPO/PassManagerBuilder.cpp

@@ -223,6 +223,10 @@ static void addHLSLPasses(bool HLSLHighLevel, unsigned OptLevel, bool OnlyWarnOn
     MPM.add(createHLDeadFunctionEliminationPass());
   }
 
+  // Do this before scalarrepl-param-hlsl for opportunities to move things
+  // like resource arrays to alloca, allowing more likely memcpy replacement.
+  MPM.add(createLowerStaticGlobalIntoAlloca());
+
   // Expand buffer store intrinsics before we SROA
   MPM.add(createHLExpandStoreIntrinsicsPass());
 

lib/Transforms/Scalar/SROA.cpp

@@ -627,6 +627,17 @@ static Value *foldPHINodeOrSelectInst(Instruction &I) {
   return foldSelectInst(cast<SelectInst>(I));
 }
 
+// HLSL Change - Detect HLSL Object or Matrix [array] type
+// These types should be SROA'd elsewhere as necessary.
+bool SkipHLSLType(Type *Ty, bool SkipHLSLMat) {
+  if (Ty->isPointerTy())
+    Ty = Ty->getPointerElementType();
+  while (Ty->isArrayTy())
+    Ty = Ty->getArrayElementType();
+  return (SkipHLSLMat && hlsl::HLMatrixType::isa(Ty)) ||
+         hlsl::dxilutil::IsHLSLObjectType(Ty);
+}
+
 /// \brief Builder for the alloca slices.
 ///
 /// This class builds a set of alloca slices by recursively visiting the uses
@@ -697,21 +708,10 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
     if (BC.use_empty())
       return markAsDead(BC);
     // HLSL Change Begin - not sroa matrix type.
-    if (PointerType *PT = dyn_cast<PointerType>(BC.getType())) {
-      Type *EltTy = PT->getElementType();
-      if ((SkipHLSLMat && hlsl::HLMatrixType::isa(EltTy)) ||
-          hlsl::dxilutil::IsHLSLObjectType(EltTy)) {
-        AS.PointerEscapingInstr = &BC;
-        return;
-      }
-      if (PointerType *SrcPT = dyn_cast<PointerType>(BC.getSrcTy())) {
-        Type *SrcEltTy = SrcPT->getElementType();
-        if ((SkipHLSLMat && hlsl::HLMatrixType::isa(SrcEltTy)) ||
-            hlsl::dxilutil::IsHLSLObjectType(SrcEltTy)) {
-          AS.PointerEscapingInstr = &BC;
-          return;
-        }
-      }
+    if (SkipHLSLType(BC.getType(), SkipHLSLMat) ||
+        SkipHLSLType(BC.getSrcTy(), SkipHLSLMat)) {
+      AS.PointerEscapingInstr = &BC;
+      return;
     }
     // HLSL Change End.
     return Base::visitBitCastInst(BC);
@@ -775,8 +775,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
 
   void visitLoadInst(LoadInst &LI) {
     // HLSL Change Begin - not sroa matrix type.
-    if ((SkipHLSLMat && hlsl::HLMatrixType::isa(LI.getType())) ||
-        hlsl::dxilutil::IsHLSLObjectType(LI.getType()))
+    if (SkipHLSLType(LI.getType(), SkipHLSLMat))
       return PI.setEscapedAndAborted(&LI);
     // HLSL Change End.
     assert((!LI.isSimple() || LI.getType()->isSingleValueType()) &&
@@ -796,8 +795,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
     if (ValOp == *U)
       return PI.setEscapedAndAborted(&SI);
     // HLSL Change Begin - not sroa matrix type.
-    if ((SkipHLSLMat && hlsl::HLMatrixType::isa(ValOp->getType())) ||
-        hlsl::dxilutil::IsHLSLObjectType(ValOp->getType()))
+    if (SkipHLSLType(ValOp->getType(), SkipHLSLMat))
       return PI.setEscapedAndAborted(&SI);
     // HLSL Change End.
 
@@ -3366,8 +3364,7 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     if (!LI.isSimple() || LI.getType()->isSingleValueType())
       return false;
     // HLSL Change Begin - not sroa matrix type.
-    if ((SkipHLSLMat && hlsl::HLMatrixType::isa(LI.getType())) ||
-        hlsl::dxilutil::IsHLSLObjectType(LI.getType()))
+    if (SkipHLSLType(LI.getType(), SkipHLSLMat))
       return false;
     // HLSL Change End.
 
@@ -3405,8 +3402,7 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     if (V->getType()->isSingleValueType())
       return false;
     // HLSL Change Begin - not sroa matrix type.
-    if ((SkipHLSLMat && hlsl::HLMatrixType::isa(V->getType())) ||
-        hlsl::dxilutil::IsHLSLObjectType(V->getType()))
+    if (SkipHLSLType(V->getType(), SkipHLSLMat))
       return false;
     // HLSL Change End.
     // We have an aggregate being stored, split it apart.
@@ -3419,17 +3415,9 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
 
   bool visitBitCastInst(BitCastInst &BC) {
     // HLSL Change Begin - not sroa matrix type.
-    if (PointerType *PT = dyn_cast<PointerType>(BC.getType())) {
-      Type *EltTy = PT->getElementType();
-      if ((SkipHLSLMat && hlsl::HLMatrixType::isa(EltTy)) ||
-          hlsl::dxilutil::IsHLSLObjectType(EltTy))
-        return false;
-      if (PointerType *SrcPT = dyn_cast<PointerType>(BC.getSrcTy())) {
-        Type *SrcEltTy = SrcPT->getElementType();
-        if ((SkipHLSLMat && hlsl::HLMatrixType::isa(SrcEltTy)) ||
-            hlsl::dxilutil::IsHLSLObjectType(SrcEltTy))
-          return false;
-      }
+    if (SkipHLSLType(BC.getType(), SkipHLSLMat) ||
+        SkipHLSLType(BC.getSrcTy(), SkipHLSLMat)) {
+      return false;
     }
     // HLSL Change End.
     enqueueUsers(BC);
@@ -4420,8 +4408,7 @@ bool SROA::runOnAlloca(AllocaInst &AI) {
       hlsl::dxilutil::IsHLSLObjectType(
           AI.getAllocatedType()) || // HLSL Change - not sroa resource type.
       // HLSL Change Begin - not sroa matrix type.
-      (SkipHLSLMat &&
-       hlsl::HLMatrixType::isa(AI.getAllocatedType())) ||
+      SkipHLSLType(AI.getAllocatedType(), SkipHLSLMat) ||
       // HLSL Change End.
       DL.getTypeAllocSize(AI.getAllocatedType()) == 0)
     return false;

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -49,6 +49,7 @@
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include "llvm/Transforms/Utils/SSAUpdater.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "dxc/HLSL/HLOperations.h"
 #include "dxc/DXIL/DxilConstants.h"
 #include "dxc/HLSL/HLModule.h"
@@ -3443,7 +3444,7 @@ static bool ReplaceMemcpy(Value *V, Value *Src, MemCpyInst *MC,
   if (V != Src && V->hasOneUse() && Src->hasOneUse())
     return false;
 
-  // If the memcpy doesn't dominate all its users,
+  // If the source of the memcpy (Src) doesn't dominate all users of dest (V),
   // full replacement isn't possible without complicated PHI insertion
   // This will likely replace with ld/st which will be replaced in mem2reg
   if (Instruction *SrcI = dyn_cast<Instruction>(Src))
@@ -3587,32 +3588,40 @@ static bool ReplaceUseOfZeroInitEntry(Instruction *I, Value *V) {
   return true;
 }
 
-static bool ReplaceUseOfZeroInitPostDom(Instruction *I, Value *V,
-                                    PostDominatorTree &PDT) {
+// If a V user is dominated by memcpy (I),
+//    skip it - memcpy dest can simply alias to src for this user.
+// If the V user may follow the memcpy (I),
+//    return false - memcpy dest not safe to replace with src.
+// Otherwise,
+//    replace use with zeroinitializer.
+static bool ReplaceUseOfZeroInit(Instruction *I, Value *V,
+                                 DominatorTree &DT,
+                                 SmallPtrSet<BasicBlock*, 8> &Reachable) {
   BasicBlock *BB = I->getParent();
   Function *F = I->getParent()->getParent();
   for (auto U = V->user_begin(); U != V->user_end(); ) {
     Instruction *UI = dyn_cast<Instruction>(*(U++));
-    if (!UI)
+    if (!UI || UI == I)
       continue;
     if (UI->getParent()->getParent() != F)
       continue;
 
-    if (!PDT.dominates(BB, UI->getParent()))
+    // Skip properly dominated users
+    if (DT.properlyDominates(BB, UI->getParent()))
+      continue;
+
+    // If user is found in memcpy successor list
+    // then the user is not safe to replace with zeroinitializer.
+    if (Reachable.count(UI->getParent()))
       return false;
 
+    // Remaining cases are where I:
+    // - is at the end of the same block
+    // - does not precede UI on any path
     if (isa<GetElementPtrInst>(UI) || isa<BitCastInst>(UI)) {
-      if (!ReplaceUseOfZeroInitPostDom(I, UI, PDT))
-        return false;
-      else
+      if (ReplaceUseOfZeroInit(I, UI, DT, Reachable))
         continue;
-    }
-
-    if (BB != UI->getParent() || UI == I)
-      continue;
-    // I is the last inst in the block after split.
-    // Any inst in current block is before I.
-    if (LoadInst *LI = dyn_cast<LoadInst>(UI)) {
+    } else if (LoadInst *LI = dyn_cast<LoadInst>(UI)) {
       LI->replaceAllUsesWith(ConstantAggregateZero::get(LI->getType()));
       LI->eraseFromParent();
       continue;
@@ -3621,23 +3630,42 @@ static bool ReplaceUseOfZeroInitPostDom(Instruction *I, Value *V,
   }
   return true;
 }
+
+// Recursively collect all successors of BB and BB's successors.
+// BB will not be in set unless it's reachable through its successors.
+static void CollectReachableBBs(BasicBlock *BB, SmallPtrSet<BasicBlock*, 8> &Reachable) {
+  for (auto S : successors(BB)) {
+    if (Reachable.insert(S).second)
+      CollectReachableBBs(S, Reachable);
+  }
+}
+
 // When zero initialized GV has only one define, all uses before the def should
 // use zero.
 static bool ReplaceUseOfZeroInitBeforeDef(Instruction *I, GlobalVariable *GV) {
   BasicBlock *BB = I->getParent();
   Function *F = I->getParent()->getParent();
   // Make sure I is the last inst for BB.
+  BasicBlock *NewBB = nullptr;
   if (I != BB->getTerminator())
-    BB->splitBasicBlock(I->getNextNode());
+    NewBB = BB->splitBasicBlock(I->getNextNode());
 
+  bool bSuccess = false;
   if (&F->getEntryBlock() == I->getParent()) {
-    return ReplaceUseOfZeroInitEntry(I, GV);
+    bSuccess = ReplaceUseOfZeroInitEntry(I, GV);
   } else {
-    // Post dominator tree.
-    PostDominatorTree PDT;
-    PDT.runOnFunction(*F);
-    return ReplaceUseOfZeroInitPostDom(I, GV, PDT);
+    DominatorTree DT;
+    DT.recalculate(*F);
+    SmallPtrSet<BasicBlock*, 8> Reachable;
+    CollectReachableBBs(BB, Reachable);
+    bSuccess = ReplaceUseOfZeroInit(I, GV, DT, Reachable);
   }
+
+  // Re-merge basic block to keep things simpler
+  if (NewBB)
+    llvm::MergeBlockIntoPredecessor(NewBB);
+
+  return bSuccess;
 }
 
 // Use `DT` to trace all users and make sure `I`'s BB dominates them all
@@ -5918,7 +5946,7 @@ class LowerStaticGlobalIntoAlloca : public ModulePass {
         }
       } else {
         for (Function &F : M) {
-          if (!HLM.IsEntry(&F)) {
+          if (F.isDeclaration() || !HLM.IsEntry(&F)) {
             continue;
           }
           entryAndInitFunctionSet.insert(&F);
@@ -5936,7 +5964,7 @@ class LowerStaticGlobalIntoAlloca : public ModulePass {
         }
       } else {
         for (Function &F : M) {
-          if (!DM.IsEntry(&F))
+          if (F.isDeclaration() || !DM.IsEntry(&F))
             continue;
           entryAndInitFunctionSet.insert(&F);
         }
@@ -6171,14 +6199,17 @@ bool LowerStaticGlobalIntoAlloca::lowerStaticGlobalIntoAlloca(
     GlobalVariable *GV, const DataLayout &DL, DxilTypeSystem &typeSys,
     SetVector<Function *> &entryAndInitFunctionSet) {
   GV->removeDeadConstantUsers();
+  bool bIsObjectTy = dxilutil::IsHLSLObjectType(
+      dxilutil::StripArrayTypes(GV->getType()->getElementType()));
   // Create alloca for each entry.
   DenseMap<Function *, AllocaInst *> allocaMap;
   for (Function *F : entryAndInitFunctionSet) {
     IRBuilder<> Builder(dxilutil::FindAllocaInsertionPt(F));
     AllocaInst *AI = Builder.CreateAlloca(GV->getType()->getElementType());
     allocaMap[F] = AI;
     // Store initializer is exist.
-    if (GV->hasInitializer() && !isa<UndefValue>(GV->getInitializer())) {
+    if (GV->hasInitializer() && !isa<UndefValue>(GV->getInitializer()) &&
+        !bIsObjectTy) { // Do not zerio-initialize object allocas
       Builder.CreateStore(GV->getInitializer(), GV);
     }
   }

tools/clang/test/HLSLFileCheck/passes/llvm/sroa/SROA_res-array.hlsl

@@ -0,0 +1,38 @@
+// RUN: %dxc -T lib_6_5 %s  | FileCheck %s
+
+// Make sure we don't use alloca, meaning we successfully aliased the static
+// and local resource arrays to the original global array.
+
+// CHECK: define void @main()
+// CHECK-NOT: alloca
+
+Texture2D<float4> buf[16];
+static Texture2D<float4> s_buf[16];
+
+void Init(Texture2D<float4> buf[16]) {
+  s_buf = buf;
+}
+
+float4 Use(Texture2D<float4> buf[16], uint i) {
+  return buf[i][int2(3,4)];
+}
+
+float4 Use(uint i) {
+  return Use(s_buf, i);
+}
+
+[shader("pixel")]
+float4 main(uint i:I) : SV_Target {
+  Texture2D<float4> lbuf[16];
+  lbuf = buf;
+  Init(lbuf);
+  return Use(i);
+}
+
+[shader("pixel")]
+float4 other_main(uint i:I) : SV_Target {
+  Texture2D<float4> lbuf[16];
+  lbuf = buf;
+  Init(lbuf);
+  return Use(i);
+}