Always move allocas to entry block (#1465)

lib/HLSL/HLPreprocess.cpp

@@ -59,27 +59,30 @@ public:
 
     StackSave->eraseFromParent();
     StackRestore->eraseFromParent();
-    // Has stacksave/store mean alloca not in entry block.
-    if (bUpdated) {
-      // Make sure all allocas are in entry block.
-      for (Function &F : M.functions()) {
-        MoveAllocasToEntryBlock(&F);
-      }
+
+    // If stacksave/store is present, it means alloca not in the
+    // entry block. However, there could be other cases where allocas
+    // could be present in the non-entry blocks.
+    // Therefore, always go through all non-entry blocks and
+    // make sure all allocas are moved to the entry block.
+    for (Function &F : M.functions()) {
+      bUpdated |= MoveAllocasToEntryBlock(&F);
     }
 
     return bUpdated;
   }
 
 private:
-  void MoveAllocasToEntryBlock(Function *F);
+  bool MoveAllocasToEntryBlock(Function *F);
 };
 
 char HLPreprocess::ID = 0;
 
 // Make sure all allocas are in entry block.
-void HLPreprocess::MoveAllocasToEntryBlock(Function *F) {
+bool HLPreprocess::MoveAllocasToEntryBlock(Function *F) {
+  bool changed = false;
   if (F->getBasicBlockList().size() < 2)
-    return;
+    return changed;
   BasicBlock &Entry = F->getEntryBlock();
   IRBuilder<> Builder(Entry.getFirstInsertionPt());
 
@@ -92,9 +95,11 @@ void HLPreprocess::MoveAllocasToEntryBlock(Function *F) {
       if (isa<AllocaInst>(I)) {
         I->removeFromParent();
         Builder.Insert(I);
+        changed = true;
       }
     }
   }
+  return changed;
 }
 
 } // namespace

tools/clang/test/CodeGenHLSL/quick-test/alloca-non-entry-basic-block.hlsl

@@ -0,0 +1,17 @@
+// RUN: %dxc /Tps_6_0 /Emain > %s | FileCheck %s
+// CHECK: define void @main()
+// CHECK: entry
+
+float2 foo() {
+  return float2(1.0f, -1.0f);
+}
+
+[RootSignature("")]
+float main() : SV_Target {
+  for (int c = 0; c < 2; ++c) {
+    if (foo()[c] >= 1.0f) {
+      return foo()[c];
+    }
+  }
+  return 1.0f;
+}