[MLIR][NFC] Fix/improve style in affine-licm pass

[bondhugula]

Summary

The code here appears to be out of line with proper style and
guidelines. Fix this.

[Groverkss]

LGTM with a minor suggestion.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

100–102

Braces can be removed here.

[bondhugula]

Address comment.

[bondhugula]

Merged at https://reviews.llvm.org/rGaa7aac9a3fa0

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

 struct LoopInvariantCodeMotion
     : public impl::AffineLoopInvariantCodeMotionBase<LoopInvariantCodeMotion> {
   void runOnOperation() override;
   void runOnAffineForOp(AffineForOp forOp);
 };
 } // namespace
 
 static bool
-checkInvarianceOfNestedIfOps(Operation *op, Value indVar, ValueRange iterArgs,
+checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar, ValueRange iterArgs,
                              SmallPtrSetImpl<Operation *> &opsWithUsers,
                              SmallPtrSetImpl<Operation *> &opsToHoist);
 static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
                               SmallPtrSetImpl<Operation *> &opsWithUsers,
                               SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
                                  ValueRange iterArgs,
                                  SmallPtrSetImpl<Operation *> &opsWithUsers,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
 // Returns true if the individual op is loop invariant.
 bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
                        SmallPtrSetImpl<Operation *> &opsWithUsers,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
-  if (isa<AffineIfOp>(op)) {
-    if (!checkInvarianceOfNestedIfOps(&op, indVar, iterArgs, opsWithUsers,
-                                      opsToHoist)) {
+  if (auto ifOp = dyn_cast<AffineIfOp>(op)) {
+    if (!checkInvarianceOfNestedIfOps(ifOp, indVar, iterArgs, opsWithUsers,
+                                      opsToHoist))
       return false;
-    }
   } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
     if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs,
-                                          opsWithUsers, opsToHoist)) {
+                                          opsWithUsers, opsToHoist))
       return false;
-    }
   } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
     // TODO: Support DMA ops.
     return false;
   } else if (!matchPattern(&op, m_Constant())) {
     // 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()
+      auto read = dyn_cast<AffineReadOpInterface>(op);
+      Value memref = read ? read.getMemRef()
                           : cast<AffineWriteOpInterface>(op).getMemRef();
       for (auto *user : memref.getUsers()) {
         // If this memref has a user that is a DMA, give up because these
         // operations write to this memref.
-        if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
+        if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op))
           return false;
-        }
         // If the memref used by the load/store is used in a store elsewhere in
         // the loop nest, we do not hoist. Similarly, if the memref used in a
         // load is also being stored too, we do not hoist the load.
         if (isa<AffineWriteOpInterface>(user) ||
             (isa<AffineReadOpInterface>(user) &&
              isa<AffineWriteOpInterface>(op))) {
           if (&op != user) {
             SmallVector<AffineForOp, 8> userIVs;
             getLoopIVs(*user, &userIVs);
             // Check that userIVs don't contain the for loop around the op.
-            if (llvm::is_contained(userIVs, getForInductionVarOwner(indVar))) {
+            if (llvm::is_contained(userIVs, getForInductionVarOwner(indVar)))
               return false;
-            }
-          }
-        }
-      }
-    }
 
     if (op.getNumOperands() == 0 && !isa<AffineYieldOp>(op)) {
-      LLVM_DEBUG(llvm::dbgs() << "\nNon-constant op with 0 operands\n");
+      LLVM_DEBUG(llvm::dbgs() << "Non-constant op with 0 operands\n");
       return false;
     }
   }
 
   // Check operands.
   for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
     auto *operandSrc = op.getOperand(i).getDefiningOp();
 
     LLVM_DEBUG(
-        op.getOperand(i).print(llvm::dbgs() << "\nIterating on operand\n"));
+        op.getOperand(i).print(llvm::dbgs() << "Iterating on operand\n"));
 
     // If the loop IV is the operand, this op isn't loop invariant.
     if (indVar == op.getOperand(i)) {
-      LLVM_DEBUG(llvm::dbgs() << "\nLoop IV is the operand\n");
+      LLVM_DEBUG(llvm::dbgs() << "Loop IV is the operand\n");
       return false;
     }
 
     // If the one of the iter_args is the operand, this op isn't loop invariant.
     if (llvm::is_contained(iterArgs, op.getOperand(i))) {
-      LLVM_DEBUG(llvm::dbgs() << "\nOne of the iter_args is the operand\n");
+      LLVM_DEBUG(llvm::dbgs() << "One of the iter_args is the operand\n");
       return false;
     }
 
-    if (operandSrc != nullptr) {
-      LLVM_DEBUG(llvm::dbgs() << *operandSrc << "\nIterating on operand src\n");
+    if (operandSrc) {
+      LLVM_DEBUG(llvm::dbgs() << *operandSrc << "Iterating on operand src\n");
 
-      // 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 (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
+      // 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 (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0)
         return false;
-      }
-    }
-  }
 
   // If no operand was loop variant, mark this op for motion.
   opsToHoist.insert(&op);
   return true;
 }
 
 // Checks if all ops in a region (i.e. list of blocks) are loop invariant.
 bool areAllOpsInTheBlockListInvariant(
     Region &blockList, Value indVar, ValueRange iterArgs,
     SmallPtrSetImpl<Operation *> &opsWithUsers,
     SmallPtrSetImpl<Operation *> &opsToHoist) {
 
   for (auto &b : blockList) {
     for (auto &op : b) {
-      if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) {
+      if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist))
         return false;
-      }
-    }
-  }
 
   return true;
 }
 
 // Returns true if the affine.if op can be hoisted.
-bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
+bool checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar,
                                   ValueRange iterArgs,
                                   SmallPtrSetImpl<Operation *> &opsWithUsers,
                                   SmallPtrSetImpl<Operation *> &opsToHoist) {
-  assert(isa<AffineIfOp>(op));
-  auto ifOp = cast<AffineIfOp>(op);
-
   if (!areAllOpsInTheBlockListInvariant(ifOp.getThenRegion(), indVar, iterArgs,
-                                        opsWithUsers, opsToHoist)) {
+                                        opsWithUsers, opsToHoist))
     return false;
-  }
 
   if (!areAllOpsInTheBlockListInvariant(ifOp.getElseRegion(), indVar, iterArgs,
-                                        opsWithUsers, opsToHoist)) {
+                                        opsWithUsers, opsToHoist))
     return false;
-  }
 
   return true;
 }
 
 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   auto *loopBody = forOp.getBody();
   auto indVar = forOp.getInductionVar();
   ValueRange iterArgs = forOp.getRegionIterArgs();