Remove stacksave and stackstore. (#914)

include/dxc/HLSL/DxilGenerationPass.h

@@ -54,6 +54,7 @@ FunctionPass *createDxilExpandTrigIntrinsicsPass();
 ModulePass *createDxilLoadMetadataPass();
 ModulePass *createDxilDeadFunctionEliminationPass();
 ModulePass *createHLDeadFunctionEliminationPass();
+ModulePass *createHLPreprocessPass();
 ModulePass *createDxilPrecisePropagatePass();
 FunctionPass *createDxilPreserveAllOutputsPass();
 FunctionPass *createDxilLegalizeResourceUsePass();
@@ -83,6 +84,7 @@ void initializeDxilExpandTrigIntrinsicsPass(llvm::PassRegistry&);
 void initializeDxilLoadMetadataPass(llvm::PassRegistry&);
 void initializeDxilDeadFunctionEliminationPass(llvm::PassRegistry&);
 void initializeHLDeadFunctionEliminationPass(llvm::PassRegistry&);
+void initializeHLPreprocessPass(llvm::PassRegistry&);
 void initializeDxilPrecisePropagatePassPass(llvm::PassRegistry&);
 void initializeDxilPreserveAllOutputsPass(llvm::PassRegistry&);
 void initializeDxilLegalizeResourceUsePassPass(llvm::PassRegistry&);

lib/HLSL/CMakeLists.txt

@@ -45,6 +45,7 @@ add_llvm_library(LLVMHLSL
   HLOperations.cpp
   HLOperationLower.cpp
   HLOperationLowerExtension.cpp
+  HLPreprocess.cpp
   HLResource.cpp
   HLSignatureLower.cpp
   PauseResumePasses.cpp

lib/HLSL/DxcOptimizer.cpp

@@ -117,6 +117,7 @@ HRESULT SetupRegistryPassForHLSL() {
     initializeHLEmitMetadataPass(Registry);
     initializeHLEnsureMetadataPass(Registry);
     initializeHLMatrixLowerPassPass(Registry);
+    initializeHLPreprocessPass(Registry);
     initializeHoistConstantArrayPass(Registry);
     initializeIPSCCPPass(Registry);
     initializeIndVarSimplifyPass(Registry);

lib/HLSL/HLPreprocess.cpp

@@ -0,0 +1,105 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// HLPreprocess.cpp                                                          //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Preprocess HLModule after inline.                                         //
+//                                                                           //
+///////////////////////////////////////////////////////////////////////////////
+
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Intrinsics.h"
+
+#include "dxc/HLSL/DxilGenerationPass.h"
+
+using namespace llvm;
+
+///////////////////////////////////////////////////////////////////////////////
+// HLPreprocess.
+// Inliner will create stacksave stackstore if there are allocas inside block
+// other than entry block. HLPreprocess will remove stacksave and stackstore and
+// put all allocas inside entry block.
+//
+namespace {
+
+class HLPreprocess : public ModulePass {
+public:
+  static char ID; // Pass identification, replacement for typeid
+  explicit HLPreprocess() : ModulePass(ID) {}
+
+  const char *getPassName() const override {
+    return "Preprocess HLModule after inline";
+  }
+
+  bool runOnModule(Module &M) override {
+    bool bUpdated = false;
+    // Remove stacksave and stackstore.
+    // Get the two intrinsics we care about.
+    Function *StackSave = Intrinsic::getDeclaration(&M, Intrinsic::stacksave);
+    Function *StackRestore =
+        Intrinsic::getDeclaration(&M, Intrinsic::stackrestore);
+    // If has user, remove user first.
+    if (!StackSave->user_empty() || !StackRestore->user_empty()) {
+      for (auto it = StackRestore->user_begin();
+           it != StackRestore->user_end();) {
+        Instruction *I = cast<Instruction>(*(it++));
+        I->eraseFromParent();
+      }
+
+      for (auto it = StackSave->user_begin(); it != StackSave->user_end(); ) {
+        Instruction *I = cast<Instruction>(*(it++));
+        I->eraseFromParent();
+      }
+      bUpdated = true;
+    }
+
+    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);
+      }
+    }
+
+    return bUpdated;
+  }
+
+private:
+  void MoveAllocasToEntryBlock(Function *F);
+};
+
+char HLPreprocess::ID = 0;
+
+// Make sure all allocas are in entry block.
+void HLPreprocess::MoveAllocasToEntryBlock(Function *F) {
+  if (F->getBasicBlockList().size() < 2)
+    return;
+  BasicBlock &Entry = F->getEntryBlock();
+  IRBuilder<> Builder(Entry.getFirstInsertionPt());
+
+  for (auto bb = F->begin(); bb != F->end(); bb++) {
+    BasicBlock *BB = bb;
+    if (BB == &Entry)
+      continue;
+    for (auto it = BB->begin(); it != BB->end();) {
+      Instruction *I = (it++);
+      if (isa<AllocaInst>(I)) {
+        I->removeFromParent();
+        Builder.Insert(I);
+      }
+    }
+  }
+}
+
+} // namespace
+
+ModulePass *llvm::createHLPreprocessPass() { return new HLPreprocess(); }
+
+INITIALIZE_PASS(HLPreprocess, "hl-preprocess",
+                "Preprocess HLModule after inline", false, false)

lib/Transforms/IPO/PassManagerBuilder.cpp

@@ -210,6 +210,8 @@ static void addHLSLPasses(bool HLSLHighLevel, bool NoOpt, hlsl::HLSLExtensionsCo
     return;
   }
 
+  MPM.add(createHLPreprocessPass());
+
   if (!NoOpt) {
     MPM.add(createHLDeadFunctionEliminationPass());
   }

utils/hct/hctdb.py

@@ -1330,6 +1330,7 @@ class db_dxil(object):
         add_pass('hlsl-dxilload', 'DxilLoadMetadata', 'HLSL DXIL Metadata Load', [])
         add_pass('dxil-dfe', 'DxilDeadFunctionElimination', 'Remove all unused function except entry from DxilModule', [])
         add_pass('hl-dfe', 'HLDeadFunctionElimination', 'Remove all unused function except entry from HLModule', [])
+        add_pass('hl-preprocess', 'HLPreprocess', 'Preprocess HLModule after inline', [])
         add_pass('hlsl-dxil-expand-trig', 'DxilExpandTrigIntrinsics', 'DXIL expand trig intrinsics', [])
         add_pass('hlsl-hca', 'HoistConstantArray', 'HLSL constant array hoisting', [])
         add_pass('hlsl-dxil-preserve-all-outputs', 'DxilPreserveAllOutputs', 'DXIL write to all outputs in signature', [])