diff --git a/hpvm/lib/Transforms/Utils/HPVMLoopTransforms.cpp b/hpvm/lib/Transforms/Utils/HPVMLoopTransforms.cpp
index 083fe374b230bc482afd6cdf9b47e8afe7d863f1..3c33d4062d49030aaa70b12174c28c1b26c0412c 100644
--- a/hpvm/lib/Transforms/Utils/HPVMLoopTransforms.cpp
+++ b/hpvm/lib/Transforms/Utils/HPVMLoopTransforms.cpp
@@ -358,95 +358,35 @@ void updateNZLoop(Loop *L, LoopInfo &LI, ScalarEvolution &SE, Function *F,
 
   std::vector<IntrinsicInst *> IItoRemove;
 
-  if (!L->isInvalid()) {
-    // The loop has not been fully unrolled, find and update the NZLoop
-    // intrinsic
-    PHINode *IndVar = L->getInductionVariable(SE);
-    bool found = false;
-    for (auto *User : IndVar->users()) {
-      if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
-        if (II->getIntrinsicID() != Intrinsic::hpvm_nz_loop) {
-          continue;
-        }
-        if (found) {
-          // This is the same copy of NZLoop, most likely a result of inner-more
-          // loops getting unrolled. Delete it!
-          IItoRemove.push_back(II);
-          continue;
-        }
-        found = true;
-        auto TripCount =
-            dyn_cast<ConstantInt>(II->getOperand(1))->getZExtValue();
-        auto NewTripCount = TripCount / UF;
-        LLVM_DEBUG(errs() << "Updating TripCount in Intrinsic: " << TripCount
-                     << " -> " << NewTripCount << "\n");
-        ConstantInt *NTC =
-            ConstantInt::get(Type::getInt32Ty(F->getContext()), NewTripCount);
-        II->setOperand(1, NTC);
-        LLVM_DEBUG(errs() << "Updated Intrinsic: " << *NTC << "\n");
+  // The loop has not been fully unrolled, find and update the NZLoop
+  // intrinsic
+  PHINode *IndVar = L->getInductionVariable(SE);
+  bool found = false;
+  for (auto *User : IndVar->users()) {
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(User)) {
+      if (II->getIntrinsicID() != Intrinsic::hpvm_nz_loop) {
+        continue;
       }
+      if (found) {
+        // This is the same copy of NZLoop, most likely a result of inner-more
+        // loops getting unrolled. Delete it!
+        IItoRemove.push_back(II);
+        continue;
+      }
+      found = true;
+      auto TripCount =
+          dyn_cast<ConstantInt>(II->getOperand(1))->getZExtValue();
+      auto NewTripCount = TripCount / UF;
+      LLVM_DEBUG(errs() << "Updating TripCount in Intrinsic: " << TripCount
+                    << " -> " << NewTripCount << "\n");
+      ConstantInt *NTC =
+          ConstantInt::get(Type::getInt32Ty(F->getContext()), NewTripCount);
+      II->setOperand(1, NTC);
+      LLVM_DEBUG(errs() << "Updated Intrinsic: " << *NTC << "\n");
     }
-    assert(found && "Expecting to have found an NZLoop for this loop!");
-  }
-  // for (auto ib = inst_begin(F); ib != inst_end(F); ++ib) {
-  //   Instruction *I = &*ib;
-  //   if (BuildDFG::isHPVMIntrinsic(I)) {
-  //     // Find the NZLoop intrinsic
-  //     IntrinsicInst *II = dyn_cast<IntrinsicInst>(I);
-  //     if (II->getIntrinsicID() != Intrinsic::hpvm_nz_loop) {
-  //       continue;
-  //     }
-  //     assert(II->getIntrinsicID() == Intrinsic::hpvm_nz_loop &&
-  //            "Unhandled intrinsic!");
-  //     LLVM_DEBUG(errs() << "Found Intrinsic: " << *I << "\n");
-  //     // Check if the intrinsic is in a header BB, if not, make sure it is
-  //     part
-  //     // of L and remove it. this is the result of unrolling.
-  //     auto *IIBB = II->getParent();
-  //     if (!LI.isLoopHeader(IIBB)) {
-  //       LLVM_DEBUG(errs() << "The intrinsic is not in a header, and also not part
-  //       "
-  //                       "of this loop!\n");
-  //       LLVM_DEBUG(errs() << "Removing Intrisic: " << *II << "\n");
-  //       IItoRemove.push_back(II);
-  //     } else {
-  //       // If this intrinsic is in L's header, first, check if its operand is
-  //       // the phi node, if it is, we need to update the value for the trip
-  //       // count. Otherwise, it has to be removed.
-  //       if (!L->isInvalid() && (L->getHeader() == IIBB)) {
-  //         PHINode *IIOp = dyn_cast<PHINode>(II->getOperand(0));
-  //         bool NeedsRemoval = false;
-  //         if (IIOp == nullptr) {
-  //           if (Instruction *IIOpInst =
-  //                   dyn_cast<Instruction>(II->getOperand(0))) {
-  //             if (!(isa<PHINode>(IIOpInst->getOperand(0)) ||
-  //                   isa<PHINode>(IIOpInst->getOperand(1)))) {
-  //               NeedsRemoval = true;
-  //             }
-  //           } else {
-  //             NeedsRemoval = true;
-  //           }
-  //         }
-  //         if (NeedsRemoval) {
-  //           LLVM_DEBUG(errs() << "Removing Intrisici that is result of unrolling:
-  //           "
-  //                        << *II << "\n");
-  //           IItoRemove.push_back(II);
-  //         } else {
-  //           auto TripCount =
-  //               dyn_cast<ConstantInt>(II->getOperand(1))->getZExtValue();
-  //           auto NewTripCount = TripCount / UF;
-  //           LLVM_DEBUG(errs() << "Updating TripCount in Intrinsic: " << TripCount
-  //                        << " -> " << NewTripCount << "\n");
-  //           ConstantInt *NTC = ConstantInt::get(
-  //               Type::getInt32Ty(F->getContext()), NewTripCount);
-  //           II->setOperand(1, NTC);
-  //           LLVM_DEBUG(errs() << "Updated Intrinsic: " << *NTC << "\n");
-  //         }
-  //       }
-  //     }
-  //   }
-  // }
+  }
+  assert(found && "Expecting to have found an NZLoop for this loop!");
+
   for (std::vector<IntrinsicInst *>::reverse_iterator ri = IItoRemove.rbegin(),
                                                       re = IItoRemove.rend();
        ri != re; ++ri)
@@ -1265,16 +1205,20 @@ void hpvm::unrollLoopByCount(Loop *L, LoopInfo &LI, DominatorTree &DT,
     if (_PreserveLCSSA)
       formLCSSA(*L, DT, &LI, &SE);
 
-    hpvmUnrollLoop(L, ULO, &LI, &SE, &DT, &AC, &ORE, &TTI, _TripCount, _TripMultiple,
+    LoopUnrollResult Result = hpvmUnrollLoop(L, ULO, &LI, &SE, &DT, &AC, &ORE, &TTI, _TripCount, _TripMultiple,
       _PreserveLCSSA);
 
-    // Without this, SE performs no extra analysis and can't determine the
-    // recurrence for the induction variable, needed in updateNZLoop
-    if (!L->isInvalid())
+
+    // These only needs to happen when the loop is not fully unrolled.
+    if (Result == LoopUnrollResult::PartiallyUnrolled) {
+      // Without this, SE performs no extra analysis and can't determine the
+      // recurrence for the induction variable, needed in updateNZLoop
       SE.forgetLoop(L);
 
-    if (_AllowRuntime && (_TripCount > 0))
-      updateNZLoop(L, LI, SE, F, _Count);
+      if (_AllowRuntime && (_TripCount > 0))
+        updateNZLoop(L, LI, SE, F, _Count);
+    
+    }
   }
 }