add FileCheck and fix comment

Yun-Fly · Yun-Fly · commit 3ad0e30ae516 · 2024-08-01T00:35:56.000-07:00
diff --git a/lib/gc/Transforms/IterativeTilingAndFusion.cpp b/lib/gc/Transforms/IterativeTilingAndFusion.cpp
@@ -41,9 +41,8 @@ static FailureOr<tensor::ExtractSliceOp>
 getClosestExtractSliceOfOperand(OpOperand &operand) {
   if (auto iterArg = dyn_cast<BlockArgument>(operand.get())) {
     if (auto loop =
-            dyn_cast<LoopLikeOpInterface>(iterArg.getOwner()->getParentOp())) {
+            dyn_cast<LoopLikeOpInterface>(iterArg.getOwner()->getParentOp()))
       return getClosestExtractSliceOfOperand(*loop.getTiedLoopInit(iterArg));
-    }
   }
 
   Operation *defineOp = operand.get().getDefiningOp();
@@ -69,10 +68,9 @@ getClosestInsertSliceOfResult(OpResult result) {
       sliceOp =
           dyn_cast<OffsetSizeAndStrideOpInterface>(useOfResult.getOwner());
     } else if (auto yieldOp = dyn_cast<scf::YieldOp>(useOfResult.getOwner())) {
-      if (auto loop = dyn_cast<LoopLikeOpInterface>(yieldOp->getParentOp())) {
+      if (auto loop = dyn_cast<LoopLikeOpInterface>(yieldOp->getParentOp()))
         return getClosestInsertSliceOfResult(
             loop->getResult(useOfResult.getOperandNumber()));
-      }
     }
   }
 
@@ -138,9 +136,8 @@ noTilingOnReductionFilter(RewriterBase &rewriter,
               presburger::BoundType::UB, tileSizes[resultExpr.index()], nullptr,
               true);
       if (!cstIterDomain || failed(cstTileSizes) ||
-          cstIterDomain != cstTileSizes) {
+          cstIterDomain != cstTileSizes)
         return failure();
-      }
     }
   }
   return success();
@@ -246,9 +243,8 @@ SingleCandidateInBlockFilter(RewriterBase &rewriter,
             scfX::getRealProducerOfExtractSliceOp(otherCandidate,
                                                   backwardSlice);
         if (succeeded(realProducer) &&
-            realProducer->getDefiningOp() == defOrUse.ownerOp) {
+            realProducer->getDefiningOp() == defOrUse.ownerOp)
           return failure();
-        }
       } else {
         SmallVector<OffsetSizeAndStrideOpInterface> forwardSlice;
         FailureOr<SmallVector<OpOperand *>> realConsumers =
@@ -257,9 +253,8 @@ SingleCandidateInBlockFilter(RewriterBase &rewriter,
         if (succeeded(realConsumers) &&
             llvm::any_of(*realConsumers, [&defOrUse](OpOperand *use) {
               return use->getOwner() == defOrUse.ownerOp;
-            })) {
+            }))
           return failure();
-        }
       }
     }
   }
@@ -338,9 +333,8 @@ static int TilingSizeComparer(RewriterBase &rewriter,
                          computeTileSizeProductOfCandidate(candidateA),
                      sizeProductB =
                          computeTileSizeProductOfCandidate(candidateB);
-  if (failed(sizeProductA) || failed(sizeProductB)) {
+  if (failed(sizeProductA) || failed(sizeProductB))
     return 0;
-  }
   // deal with equality
   if (*sizeProductA == *sizeProductB) {
     return 0;
@@ -401,17 +395,17 @@ tileAndFuseProducerOfOpOperand(RewriterBase &rewriter, OpOperand &operand,
   // a. Find the closest sliceOp
   FailureOr<tensor::ExtractSliceOp> closestSliceOp =
       getClosestExtractSliceOfOperand(operand);
-  if (failed(closestSliceOp)) {
+  if (failed(closestSliceOp))
     return std::nullopt;
-  }
+
   // b. Find the real producer and collect the sliceOp chain during backward
   // stage, sorted from inner to outer.
   SmallVector<tensor::ExtractSliceOp> backwardSlice;
   FailureOr<OpResult> realProducer =
       scfX::getRealProducerOfExtractSliceOp(*closestSliceOp, backwardSlice);
-  if (failed(realProducer)) {
+  if (failed(realProducer))
     return std::nullopt;
-  }
+
   // c. Check the producer of root source if is tilable.
   Operation *producer = realProducer->getDefiningOp<TilingInterface>();
   if (!producer)
@@ -451,17 +445,16 @@ tileAndFuseConsumerOfOpResult(RewriterBase &rewriter, OpResult result,
   // a. Find the closest sliceOp
   FailureOr<tensor::ExtractSliceOp> closestSliceOp =
       getClosestInsertSliceOfResult(result);
-  if (failed(closestSliceOp)) {
+  if (failed(closestSliceOp))
     return std::nullopt;
-  }
+
   // b. Find the real consumers and collect the sliceOp chain during forward
   // stage, sorted from inner to outer.
   SmallVector<OffsetSizeAndStrideOpInterface> forwardSlice;
   FailureOr<SmallVector<OpOperand *>> realConsumers =
       scfX::getRealConsumersFromInsertSliceOp(*closestSliceOp, forwardSlice);
-  if (failed(realConsumers)) {
+  if (failed(realConsumers))
     return std::nullopt;
-  }
 
   SmallVector<scf::SCFFuseConsumerOfSliceResult> fusedResultList;
   for (auto useOperand : *realConsumers) {
@@ -543,18 +536,16 @@ LogicalResult iterativelyFuseProducerAndConsumerOfTiledOp(
     // fuse producer
     for (OpOperand &operand : tiledOp->getOpOperands()) {
       if (std::optional<scf::SCFFuseProducerOfSliceResult> fuseProducerResult =
-              tileAndFuseProducerOfOpOperand(rewriter, operand, options)) {
+              tileAndFuseProducerOfOpOperand(rewriter, operand, options))
         tiledOpList.push_back(fuseProducerResult.value().tiledOps[0]);
-      }
     }
     // fuse consumer(s)
     for (OpResult result : tiledOp->getResults()) {
       if (std::optional<SmallVector<scf::SCFFuseConsumerOfSliceResult>>
               fuseConsumerResults =
                   tileAndFuseConsumerOfOpResult(rewriter, result, options)) {
-        for (auto &fuseConsumerResult : *fuseConsumerResults) {
+        for (auto &fuseConsumerResult : *fuseConsumerResults)
           tiledOpList.push_back(fuseConsumerResult.tiledOps[0]);
-        }
       }
     }
   }
@@ -573,30 +564,26 @@ LogicalResult iterativelyFuseProducerAndConsumerOfTiledOp(
 /// }
 static LogicalResult isSingleTiledOpInLoop(Operation *targetOp) {
   // 0. check tilable
-  if (!isa<TilingInterface>(targetOp)) {
+  if (!isa<TilingInterface>(targetOp))
     return failure();
-  }
   // 1. check parentOp
   auto forOp = targetOp->getParentOfType<LoopLikeOpInterface>();
-  if (!forOp) {
+  if (!forOp)
     return failure();
-  }
   // 2. check single one tiling interface in loop body
   auto walkResult = forOp->walk([&targetOp](TilingInterface op) {
     // some special op maybe already deal with in template
     if (isa<linalg::FillOp, linalg::CopyOp>(op))
       return WalkResult::skip();
     return op != targetOp ? WalkResult::interrupt() : WalkResult::advance();
   });
-  if (walkResult.wasInterrupted()) {
+  if (walkResult.wasInterrupted())
     return failure();
-  }
   // 3. check whether has either extract or insert slice op
   walkResult = forOp->walk(
       [](tensor::ExtractSliceOp) { return WalkResult::interrupt(); });
-  if (walkResult.wasInterrupted()) {
+  if (walkResult.wasInterrupted())
     return success();
-  }
   walkResult = forOp->walk(
       [](tensor::InsertSliceOp) { return WalkResult::interrupt(); });
   return success(walkResult.wasInterrupted());
@@ -690,9 +677,8 @@ static bool defaultTilingOfType(RewriterBase &rewriter, Operation *op,
       // word, all reduction dimensions should not be tiled.
       if (iterType == utils::IteratorType::parallel &&
           (en != iteratorTypes.size() - 1 ||
-           llvm::count(iteratorTypes, utils::IteratorType::reduction))) {
+           llvm::count(iteratorTypes, utils::IteratorType::reduction)))
         defaultTileSize[en] = rewriter.getIndexAttr(1);
-      }
     }
   }
   // If the tile sizes are all zero, no tiling would happen.
@@ -724,14 +710,13 @@ void iterativeTilingAndFusionUntilExhaustion(
     unTiledOps.clear();
     // Pre-order walk through funcOp
     f->walk<WalkOrder::PreOrder>([&unTiledOps](Operation *op) {
-      if (isa<LoopLikeOpInterface>(op)) {
+      if (isa<LoopLikeOpInterface>(op))
         return WalkResult::skip();
-      }
+
       if (isa<TilingInterface>(op) && !op->use_empty()) {
         auto parentLoop = op->getParentOfType<LoopLikeOpInterface>();
-        if (!parentLoop.getOperation()) {
+        if (!parentLoop.getOperation())
           unTiledOps.insert(op);
-        }
       }
       return WalkResult::advance();
     });
@@ -767,9 +752,8 @@ void iterativeTilingAndFusionUntilExhaustion(
         changed |= succeeded(iterativelyFuseProducerAndConsumerOfTiledOp(
             rewriter, tiledOp, sliceOptions));
       });
-      if (changed) {
+      if (changed)
         (void)mlir::simplifyRegions(rewriter, {f.getRegion()});
-      }
     } else {
       // Auto tiling with default tile size if no tiled op found. Follow tiling
       // priority based on OpTy: `Contraction`->`Reduction`->`Elementwise`.
@@ -803,15 +787,15 @@ static OpTileSizeMap defaultTileSizeParser(ArrayRef<std::string> strArgs) {
   for (auto str : strArgs) {
     str.erase(llvm::remove_if(str, llvm::isSpace), str.end());
     size_t pos = str.find(":");
-    if (pos == std::string::npos) {
+    if (pos == std::string::npos)
       llvm_unreachable(warning);
-    }
+
     std::string opType = str.substr(0, pos);
     std::string strTileSize = str.erase(0, pos + 1);
     if (strTileSize.size() <= 2 || strTileSize.front() != '{' ||
-        strTileSize.back() != '}') {
+        strTileSize.back() != '}')
       llvm_unreachable(warning);
-    }
+
     strTileSize = strTileSize.substr(1, strTileSize.size() - 2);
     SmallVector<int64_t> intTileSize;
     while ((pos = strTileSize.find(",")) != std::string::npos) {
diff --git a/lib/gc/Transforms/TilingUsingInterfaceX.cpp b/lib/gc/Transforms/TilingUsingInterfaceX.cpp
@@ -46,16 +46,18 @@ static std::tuple<OpResult, std::optional<OpOperand *>>
 getUntiledProducerFromSliceSource(OpOperand *source,
                                   ArrayRef<LoopLikeOpInterface> loops) {
   std::optional<OpOperand *> destinationIterArg;
-  auto loopIt = loops.rbegin();
-  while (auto iterArg = dyn_cast<BlockArgument>(source->get())) {
-    auto loop = *loopIt;
-    if (iterArg.getOwner()->getParentOp() != loop)
-      break;
-    source = loop.getTiedLoopInit(iterArg);
-    loopIt++;
+  if (!loops.empty()) {
+    auto loopIt = loops.rbegin();
+    while (auto iterArg = dyn_cast<BlockArgument>(source->get())) {
+      auto loop = *loopIt;
+      if (iterArg.getOwner()->getParentOp() != loop)
+        break;
+      source = loop.getTiedLoopInit(iterArg);
+      loopIt++;
+    }
+    if (loopIt == loops.rend())
+      destinationIterArg = source;
   }
-  if (loopIt == loops.rend())
-    destinationIterArg = source;
   return {dyn_cast<OpResult>(source->get()), destinationIterArg};
 }
 
@@ -190,8 +192,8 @@ tileAndFuseProducerOfSliceImpl(RewriterBase &rewriter,
 /// @return OpResult Producer : %0 = producer
 FailureOr<OpResult> mlir::scfX::getRealProducerOfExtractSliceOp(
     Operation *candidateSliceOp,
-    SmallVector<tensor::ExtractSliceOp> &backwardSlice, int curDepth,
-    int maxDepth) {
+    SmallVector<tensor::ExtractSliceOp> &backwardSlice, unsigned curDepth,
+    unsigned maxDepth) {
   if (!isa<tensor::ExtractSliceOp>(candidateSliceOp))
     return failure();
   // control recursive time in avoid of stack overflow
@@ -322,8 +324,8 @@ mlir::scfX::tileAndFuseProducerOfSlice(RewriterBase &rewriter,
 FailureOr<SmallVector<OpOperand *>>
 mlir::scfX::getRealConsumersFromInsertSliceOp(
     Operation *candidateSliceOp,
-    SmallVector<OffsetSizeAndStrideOpInterface> &forwardSlice, int curDepth,
-    int maxDepth) {
+    SmallVector<OffsetSizeAndStrideOpInterface> &forwardSlice,
+    unsigned curDepth, unsigned maxDepth) {
   if (!isa<tensor::InsertSliceOp, tensor::ParallelInsertSliceOp>(
           candidateSliceOp))
     return failure();
@@ -410,7 +412,8 @@ checkAssumptionForFusingConsumer(tensor::InsertSliceOp candidateSliceOp) {
 }
 
 /// Fetches the FIRST OpOperand of the tilable user (and use) of the value `val`
-/// which implements `TilingInterface` and `DestinationStyleOpInterface`.
+/// within the same block, which implements `TilingInterface` and
+/// `DestinationStyleOpInterface` and has non-empty user list.
 /// Returns failure otherwise.
 static FailureOr<OpOperand *> getConsumerFromUses(Value val,
                                                   Block *containingOpBlock) {
@@ -458,9 +461,8 @@ static LogicalResult isForOpYieldResultOfInnerLoop(LoopLikeOpInterface loop) {
 }
 
 /// Fetch the untiled consumer of a scf.for's result which is yielded by a
-/// tensor.insert_slice. This function makes the following assumptions :
-/// 1.  tensor.insert_slice has scf.yield as its only user.
-/// 2.  scf.for's corresponding result has only one use.
+/// tensor.insert_slice. This function makes the following assumptions that
+/// tensor.insert_slice has scf.yield as its only user.
 static FailureOr<OpOperand *>
 getUntiledConsumerFromSlice(tensor::InsertSliceOp candidateSliceOp) {
   if (failed(checkAssumptionForFusingConsumer(candidateSliceOp)))
@@ -564,15 +566,14 @@ static LogicalResult checkAssumptionForLoop(Operation *loopOp,
     // `ParallelInsertSlice` located inside `InParallelOp` has no same parent
     // block with any other types of operation. Thus, just redirecting to its
     // parent `InParallelOp`.
-    if (isa<tensor::ParallelInsertSliceOp>(userOp)) {
+    if (isa<tensor::ParallelInsertSliceOp>(userOp))
       userOp = userOp->getParentOfType<scf::InParallelOp>();
-    }
-    if (parentBlock != userOp->getBlock()) {
+
+    if (parentBlock != userOp->getBlock())
       return failure();
-    }
-    if (userOp->isBeforeInBlock(firstUserOfLoop)) {
+
+    if (userOp->isBeforeInBlock(firstUserOfLoop))
       firstUserOfLoop = userOp;
-    }
   }
   // Find the last define of consumer
   for (Value operand : consumerOp->getOperands()) {
@@ -582,12 +583,10 @@ static LogicalResult checkAssumptionForLoop(Operation *loopOp,
     auto defineOp = operand.getDefiningOp();
     if (defineOp == loopOp)
       continue;
-    if (!defineOp || parentBlock != defineOp->getBlock()) {
+    if (!defineOp || parentBlock != defineOp->getBlock())
       return failure();
-    }
-    if (lastDefOfConsumer->isBeforeInBlock(defineOp)) {
+    if (lastDefOfConsumer->isBeforeInBlock(defineOp))
       lastDefOfConsumer = defineOp;
-    }
   }
   if (firstUserOfLoop->isBeforeInBlock(lastDefOfConsumer)) {
     // Try to move if possible
diff --git a/lib/gc/Transforms/TilingUsingInterfaceX.h b/lib/gc/Transforms/TilingUsingInterfaceX.h
@@ -20,13 +20,13 @@ SmallVector<LoopLikeOpInterface> getOuterNestLoopsWhile(
 
 FailureOr<OpResult> getRealProducerOfExtractSliceOp(
     Operation *candidateSliceOp,
-    SmallVector<tensor::ExtractSliceOp> &backwardSlice, int curDepth = 0,
-    int maxDepth = 5);
+    SmallVector<tensor::ExtractSliceOp> &backwardSlice, unsigned curDepth = 0,
+    unsigned maxDepth = 5);
 
 FailureOr<SmallVector<OpOperand *>> getRealConsumersFromInsertSliceOp(
     Operation *candidateSliceOp,
-    SmallVector<OffsetSizeAndStrideOpInterface> &forwardSlice, int curDepth = 0,
-    int maxDepth = 5);
+    SmallVector<OffsetSizeAndStrideOpInterface> &forwardSlice, unsigned curDepth = 0,
+    unsigned maxDepth = 5);
 
 // Extension for upstream `tileAndFuseProducerOfSlice`
 std::optional<scf::SCFFuseProducerOfSliceResult>
diff --git a/test/gc/Transform/iterative-tiling-and-fusion.mlir b/test/gc/Transform/iterative-tiling-and-fusion.mlir
