diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -50,25 +50,25 @@
 
 static bool
 checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
-                             SmallPtrSetImpl<Operation *> &definedOps,
+                             SmallPtrSetImpl<Operation *> &opsWithUsers,
                              SmallPtrSetImpl<Operation *> &opsToHoist);
 static bool isOpLoopInvariant(Operation &op, Value indVar,
-                              SmallPtrSetImpl<Operation *> &definedOps,
+                              SmallPtrSetImpl<Operation *> &opsWithUsers,
                               SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
-                                 SmallPtrSetImpl<Operation *> &definedOps,
+                                 SmallPtrSetImpl<Operation *> &opsWithUsers,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
 // Returns true if the individual op is loop invariant.
 bool isOpLoopInvariant(Operation &op, Value indVar,
-                       SmallPtrSetImpl<Operation *> &definedOps,
+                       SmallPtrSetImpl<Operation *> &opsWithUsers,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
   if (isa<AffineIfOp>(op)) {
-    if (!checkInvarianceOfNestedIfOps(&op, indVar, definedOps, opsToHoist)) {
+    if (!checkInvarianceOfNestedIfOps(&op, indVar, opsWithUsers, opsToHoist)) {
       return false;
     }
   } else if (isa<AffineForOp>(op)) {
@@ -79,11 +79,8 @@
     // TODO: Support DMA ops.
     return false;
   } else if (!isa<ConstantOp>(op)) {
-    // Register op in the set of ops defined inside the loop. This set is used
-    // to prevent hoisting ops that depend on other ops defined inside the loop
-    // which are themselves not being hoisted.
-    definedOps.insert(&op);
-
+    // Register op in the set of ops that have users.
+    opsWithUsers.insert(&op);
     if (isa<AffineMapAccessInterface>(op)) {
       Value memref = isa<AffineReadOpInterface>(op)
                          ? cast<AffineReadOpInterface>(op).getMemRef()
@@ -135,7 +132,8 @@
         // If the value was defined in the loop (outside of the
         // if/else region), and that operation itself wasn't meant to
         // be hoisted, then mark this operation loop dependent.
-        if (definedOps.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
+        if (opsWithUsers.count(operandSrc) &&
+            opsToHoist.count(operandSrc) == 0) {
           return false;
         }
       }
@@ -149,12 +147,12 @@
 
 // Checks if all ops in a region (i.e. list of blocks) are loop invariant.
 bool areAllOpsInTheBlockListInvariant(
-    Region &blockList, Value indVar, SmallPtrSetImpl<Operation *> &definedOps,
+    Region &blockList, Value indVar, SmallPtrSetImpl<Operation *> &opsWithUsers,
     SmallPtrSetImpl<Operation *> &opsToHoist) {
 
   for (auto &b : blockList) {
     for (auto &op : b) {
-      if (!isOpLoopInvariant(op, indVar, definedOps, opsToHoist)) {
+      if (!isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
         return false;
       }
     }
@@ -165,17 +163,17 @@
 
 // Returns true if the affine.if op can be hoisted.
 bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
-                                  SmallPtrSetImpl<Operation *> &definedOps,
+                                  SmallPtrSetImpl<Operation *> &opsWithUsers,
                                   SmallPtrSetImpl<Operation *> &opsToHoist) {
   assert(isa<AffineIfOp>(op));
   auto ifOp = cast<AffineIfOp>(op);
 
-  if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, definedOps,
+  if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, opsWithUsers,
                                         opsToHoist)) {
     return false;
   }
 
-  if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, definedOps,
+  if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, opsWithUsers,
                                         opsToHoist)) {
     return false;
   }
@@ -187,18 +185,23 @@
   auto *loopBody = forOp.getBody();
   auto indVar = forOp.getInductionVar();
 
-  SmallPtrSet<Operation *, 8> definedOps;
   // This is the place where hoisted instructions would reside.
   OpBuilder b(forOp.getOperation());
 
   SmallPtrSet<Operation *, 8> opsToHoist;
   SmallVector<Operation *, 8> opsToMove;
+  SmallPtrSet<Operation *, 8> opsWithUsers;
 
   for (auto &op : *loopBody) {
+    // Register op in the set of ops that have users. This set is used
+    // to prevent hoisting ops that depend on these ops that are
+    // not being hoisted.
+    if (!op.use_empty())
+      opsWithUsers.insert(&op);
     // We don't hoist for loops.
     if (!isa<AffineForOp>(op)) {
       if (!isa<AffineYieldOp>(op)) {
-        if (isOpLoopInvariant(op, indVar, definedOps, opsToHoist)) {
+        if (isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
           opsToMove.push_back(&op);
         }
       }
diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -613,3 +613,68 @@
 // CHECK-NEXT:  addf
 // CHECK-NEXT:  affine.vector_store
 // CHECK-NEXT:  affine.for
+
+// -----
+
+#set = affine_set<(d0): (d0 - 10 >= 0)>
+// CHECK-LABEL:   func @affine_if_not_invariant(
+func @affine_if_not_invariant(%buffer: memref<1024xf32>) -> f32 {
+  %sum_init_0 = constant 0.0 : f32
+  %sum_init_1 = constant 1.0 : f32
+  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    %sum_next = affine.if #set(%i) -> (f32) {
+      %new_sum = addf %sum_iter, %t : f32
+      affine.yield %new_sum : f32
+    } else {
+      affine.yield %sum_iter : f32
+    }
+    %modified_sum = addf %sum_next, %sum_init_1 : f32
+    affine.yield %modified_sum : f32
+  }
+  return %res : f32
+}
+
+// CHECK:       constant 0.000000e+00 : f32
+// CHECK-NEXT:  constant 1.000000e+00 : f32
+// CHECK-NEXT:  affine.for
+// CHECK-NEXT:  affine.load
+// CHECK-NEXT:  affine.if
+// CHECK-NEXT:  addf
+// CHECK-NEXT:  affine.yield
+// CHECK-NEXT:  } else {
+// CHECK-NEXT:  affine.yield
+// CHECK-NEXT:  }
+// CHECK-NEXT:  addf
+// CHECK-NEXT:  affine.yield
+// CHECK-NEXT:  }
+
+// -----
+
+// CHECK-LABEL:   func @affine_for_not_invariant(
+func @affine_for_not_invariant(%in : memref<30x512xf32, 1>,
+                               %out : memref<30x1xf32, 1>) {
+  %sum_0 = constant 0.0 : f32
+  %cst_0 = constant 1.1 : f32
+  affine.for %j = 0 to 30 {
+    %sum = affine.for %i = 0 to 512 iter_args(%sum_iter = %sum_0) -> (f32) {
+      %t = affine.load %in[%j,%i] : memref<30x512xf32,1>
+      %sum_next = addf %sum_iter, %t : f32
+      affine.yield %sum_next : f32
+    }
+    %mod_sum = mulf %sum, %cst_0 : f32
+    affine.store %mod_sum, %out[%j, 0] : memref<30x1xf32, 1>
+  }
+  return
+}
+
+// CHECK:       constant 0.000000e+00 : f32
+// CHECK-NEXT:  constant 1.100000e+00 : f32
+// CHECK-NEXT:  affine.for
+// CHECK-NEXT:  affine.for
+// CHECK-NEXT:  affine.load
+// CHECK-NEXT:  addf
+// CHECK-NEXT:  affine.yield
+// CHECK-NEXT:  }
+// CHECK-NEXT:  mulf
+// CHECK-NEXT:  affine.store