diff --git a/llvm/lib/Target/DirectX/DXILLegalizePass.cpp b/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
index f9a494ce63dd3..395311e430fbb 100644
--- a/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
+++ b/llvm/lib/Target/DirectX/DXILLegalizePass.cpp
@@ -5,12 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===---------------------------------------------------------------------===//
-//===---------------------------------------------------------------------===//
-///
-/// \file This file contains a pass to remove i8 truncations and i64 extract
-/// and insert elements.
-///
-//===----------------------------------------------------------------------===//
+
 #include "DXILLegalizePass.h"
 #include "DirectX.h"
 #include "llvm/IR/Function.h"
@@ -20,37 +15,24 @@
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <functional>
-#include <map>
-#include <stack>
-#include <vector>
 
 #define DEBUG_TYPE "dxil-legalize"
 
 using namespace llvm;
-namespace {
 
 static void fixI8TruncUseChain(Instruction &I,
-                               std::stack<Instruction *> &ToRemove,
-                               std::map<Value *, Value *> &ReplacedValues) {
-
-  auto *Cmp = dyn_cast<CmpInst>(&I);
+                               SmallVectorImpl<Instruction *> &ToRemove,
+                               DenseMap<Value *, Value *> &ReplacedValues) {
 
-  if (auto *Trunc = dyn_cast<TruncInst>(&I)) {
-    if (Trunc->getDestTy()->isIntegerTy(8)) {
-      ReplacedValues[Trunc] = Trunc->getOperand(0);
-      ToRemove.push(Trunc);
-    }
-  } else if (I.getType()->isIntegerTy(8) ||
-             (Cmp && Cmp->getOperand(0)->getType()->isIntegerTy(8))) {
-    IRBuilder<> Builder(&I);
-
-    std::vector<Value *> NewOperands;
+  auto ProcessOperands = [&](SmallVector<Value *> &NewOperands) {
     Type *InstrType = IntegerType::get(I.getContext(), 32);
+
     for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {
       Value *Op = I.getOperand(OpIdx);
       if (ReplacedValues.count(Op))
         InstrType = ReplacedValues[Op]->getType();
     }
+
     for (unsigned OpIdx = 0; OpIdx < I.getNumOperands(); ++OpIdx) {
       Value *Op = I.getOperand(OpIdx);
       if (ReplacedValues.count(Op))
@@ -61,6 +43,8 @@ static void fixI8TruncUseChain(Instruction &I,
         // Note: options here are sext or sextOrTrunc.
         // Since i8 isn't supported, we assume new values
         // will always have a higher bitness.
+        assert(NewBitWidth > Value.getBitWidth() &&
+               "Replacement's BitWidth should be larger than Current.");
         APInt NewValue = Value.sext(NewBitWidth);
         NewOperands.push_back(ConstantInt::get(InstrType, NewValue));
       } else {
@@ -68,31 +52,50 @@ static void fixI8TruncUseChain(Instruction &I,
         NewOperands.push_back(Op);
       }
     }
-
-    Value *NewInst = nullptr;
-    if (auto *BO = dyn_cast<BinaryOperator>(&I)) {
-      NewInst =
-          Builder.CreateBinOp(BO->getOpcode(), NewOperands[0], NewOperands[1]);
-
-      if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(&I)) {
-        if (OBO->hasNoSignedWrap())
-          cast<BinaryOperator>(NewInst)->setHasNoSignedWrap();
-        if (OBO->hasNoUnsignedWrap())
-          cast<BinaryOperator>(NewInst)->setHasNoUnsignedWrap();
-      }
-    } else if (Cmp) {
-      NewInst = Builder.CreateCmp(Cmp->getPredicate(), NewOperands[0],
-                                  NewOperands[1]);
-      Cmp->replaceAllUsesWith(NewInst);
+  };
+  IRBuilder<> Builder(&I);
+  if (auto *Trunc = dyn_cast<TruncInst>(&I)) {
+    if (Trunc->getDestTy()->isIntegerTy(8)) {
+      ReplacedValues[Trunc] = Trunc->getOperand(0);
+      ToRemove.push_back(Trunc);
+      return;
     }
+  }
 
-    if (NewInst) {
-      ReplacedValues[&I] = NewInst;
-      ToRemove.push(&I);
+  if (auto *BO = dyn_cast<BinaryOperator>(&I)) {
+    if (!I.getType()->isIntegerTy(8))
+      return;
+    SmallVector<Value *> NewOperands;
+    ProcessOperands(NewOperands);
+    Value *NewInst =
+        Builder.CreateBinOp(BO->getOpcode(), NewOperands[0], NewOperands[1]);
+    if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(&I)) {
+      if (OBO->hasNoSignedWrap())
+        cast<BinaryOperator>(NewInst)->setHasNoSignedWrap();
+      if (OBO->hasNoUnsignedWrap())
+        cast<BinaryOperator>(NewInst)->setHasNoUnsignedWrap();
     }
-  } else if (auto *Cast = dyn_cast<CastInst>(&I)) {
+    ReplacedValues[BO] = NewInst;
+    ToRemove.push_back(BO);
+    return;
+  }
+
+  if (auto *Cmp = dyn_cast<CmpInst>(&I)) {
+    if (!Cmp->getOperand(0)->getType()->isIntegerTy(8))
+      return;
+    SmallVector<Value *> NewOperands;
+    ProcessOperands(NewOperands);
+    Value *NewInst =
+        Builder.CreateCmp(Cmp->getPredicate(), NewOperands[0], NewOperands[1]);
+    Cmp->replaceAllUsesWith(NewInst);
+    ReplacedValues[Cmp] = NewInst;
+    ToRemove.push_back(Cmp);
+    return;
+  }
+
+  if (auto *Cast = dyn_cast<CastInst>(&I)) {
     if (Cast->getSrcTy()->isIntegerTy(8)) {
-      ToRemove.push(Cast);
+      ToRemove.push_back(Cast);
       Cast->replaceAllUsesWith(ReplacedValues[Cast->getOperand(0)]);
     }
   }
@@ -100,8 +103,8 @@ static void fixI8TruncUseChain(Instruction &I,
 
 static void
 downcastI64toI32InsertExtractElements(Instruction &I,
-                                      std::stack<Instruction *> &ToRemove,
-                                      std::map<Value *, Value *> &) {
+                                      SmallVectorImpl<Instruction *> &ToRemove,
+                                      DenseMap<Value *, Value *> &) {
 
   if (auto *Extract = dyn_cast<ExtractElementInst>(&I)) {
     Value *Idx = Extract->getIndexOperand();
@@ -115,7 +118,7 @@ downcastI64toI32InsertExtractElements(Instruction &I,
           Extract->getVectorOperand(), Idx32, Extract->getName());
 
       Extract->replaceAllUsesWith(NewExtract);
-      ToRemove.push(Extract);
+      ToRemove.push_back(Extract);
     }
   }
 
@@ -132,38 +135,35 @@ downcastI64toI32InsertExtractElements(Instruction &I,
           Insert->getName());
 
       Insert->replaceAllUsesWith(Insert32Index);
-      ToRemove.push(Insert);
+      ToRemove.push_back(Insert);
     }
   }
 }
 
+namespace {
 class DXILLegalizationPipeline {
 
 public:
   DXILLegalizationPipeline() { initializeLegalizationPipeline(); }
 
   bool runLegalizationPipeline(Function &F) {
-    std::stack<Instruction *> ToRemove;
-    std::map<Value *, Value *> ReplacedValues;
+    SmallVector<Instruction *> ToRemove;
+    DenseMap<Value *, Value *> ReplacedValues;
     for (auto &I : instructions(F)) {
-      for (auto &LegalizationFn : LegalizationPipeline) {
+      for (auto &LegalizationFn : LegalizationPipeline)
         LegalizationFn(I, ToRemove, ReplacedValues);
-      }
     }
-    bool MadeChanges = !ToRemove.empty();
 
-    while (!ToRemove.empty()) {
-      Instruction *I = ToRemove.top();
-      I->eraseFromParent();
-      ToRemove.pop();
-    }
+    for (auto *Inst : reverse(ToRemove))
+      Inst->eraseFromParent();
 
-    return MadeChanges;
+    return !ToRemove.empty();
   }
 
 private:
-  std::vector<std::function<void(Instruction &, std::stack<Instruction *> &,
-                                 std::map<Value *, Value *> &)>>
+  SmallVector<
+      std::function<void(Instruction &, SmallVectorImpl<Instruction *> &,
+                         DenseMap<Value *, Value *> &)>>
       LegalizationPipeline;
 
   void initializeLegalizationPipeline() {