make return value more structured.

Author: htyu

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

@@ -111,13 +111,18 @@ LogicalResult coalescePerfectlyNestedSCFForLoops(scf::ForOp op);
 void collapseParallelLoops(RewriterBase &rewriter, scf::ParallelOp loops,
                            ArrayRef<std::vector<unsigned>> combinedDimensions);
 
+struct UnrolledLoopInfo {
+  scf::ForOp mainLoopOp;
+  scf::ForOp epilogueLoopOp;
+};
+
 /// Unrolls this for operation by the specified unroll factor. Returns the
-/// unrolled main loop and the eplilog loop in sequence, if the loop is
-/// unrolled. Otherwise returns an empty vector if the loop cannot be unrolled
-/// either due to restrictions or due to invalid unroll factors. Requires
-/// positive loop bounds and step. If specified, annotates the Ops in each
-/// unrolled iteration by applying `annotateFn`.
-SmallVector<scf::ForOp> loopUnrollByFactor(
+/// unrolled main loop and the eplilog loop, if the loop is unrolled. Otherwise
+/// returns a strucutre of null fields if the loop cannot be unrolled either due
+/// to restrictions or due to invalid unroll factors. Requires positive loop
+/// bounds and step. If specified, annotates the Ops in each unrolled iteration
+/// by applying `annotateFn`.
+UnrolledLoopInfo loopUnrollByFactor(
     scf::ForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
 

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

@@ -355,9 +355,8 @@ transform::LoopUnrollOp::applyToOne(transform::TransformRewriter &rewriter,
   LogicalResult result(failure());
   if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op)) {
     auto resultLoops = loopUnrollByFactor(scfFor, getFactor());
-    result = resultLoops.empty() ? failure() : success();
-  }
-  else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
+    result = resultLoops.mainLoopOp ? success() : failure();
+  } else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
     result = loopUnrollByFactor(affineFor, getFactor());
   else
     return emitSilenceableError()

mlir/lib/Dialect/SCF/Utils/Utils.cpp

@@ -373,16 +373,15 @@ static void generateUnrolledLoop(
 }
 
 /// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
-/// eplilog loop in sequence, if the loop is unrolled. Otherwise return an empty
-/// vector.
-SmallVector<scf::ForOp> mlir::loopUnrollByFactor(
+/// eplilog loop, if the loop is unrolled. Otherwise return null.
+UnrolledLoopInfo mlir::loopUnrollByFactor(
     scf::ForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "expected positive unroll factor");
 
   // Return if the loop body is empty.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
-    return {forOp};
+    return {forOp, nullptr};
 
   // Compute tripCount = ceilDiv((upperBound - lowerBound), step) and populate
   // 'upperBoundUnrolled' and 'stepUnrolled' for static and dynamic cases.
@@ -403,8 +402,8 @@ SmallVector<scf::ForOp> mlir::loopUnrollByFactor(
     if (unrollFactor == 1) {
       if (*constTripCount == 1 &&
           failed(forOp.promoteIfSingleIteration(rewriter)))
-        return {};
-      return {forOp};
+        return {nullptr, nullptr};
+      return {forOp, nullptr};
     }
 
     int64_t tripCountEvenMultiple =
@@ -452,8 +451,7 @@ SmallVector<scf::ForOp> mlir::loopUnrollByFactor(
         boundsBuilder.create<arith::MulIOp>(loc, step, unrollFactorCst);
   }
 
-  SmallVector<scf::ForOp, 2> resultLoops;
-  resultLoops.push_back(forOp);
+  UnrolledLoopInfo resultLoops;
 
   // Create epilogue clean up loop starting at 'upperBoundUnrolled'.
   if (generateEpilogueLoop) {
@@ -473,7 +471,7 @@ SmallVector<scf::ForOp> mlir::loopUnrollByFactor(
     epilogueForOp->setOperands(epilogueForOp.getNumControlOperands(),
                                epilogueForOp.getInitArgs().size(), results);
     (void)epilogueForOp.promoteIfSingleIteration(rewriter);
-    resultLoops.push_back(epilogueForOp);
+    resultLoops.epilogueLoopOp = epilogueForOp;
   }
 
   // Create unrolled loop.
@@ -496,6 +494,7 @@ SmallVector<scf::ForOp> mlir::loopUnrollByFactor(
       annotateFn, iterArgs, yieldedValues);
   // Promote the loop body up if this has turned into a single iteration loop.
   (void)forOp.promoteIfSingleIteration(rewriter);
+  resultLoops.mainLoopOp = forOp;
   return resultLoops;
 }