[MLIR] NFC. affine data copy generate utility return value cleanup

[bondhugula]

Summary

Clean up return value on affineDataCopyGenerate utility. Return the
actual success/failure status instead of the "number of bytes" which
isn't being used in the codebase in any way. The success/failure status
wasn't being sent out earlier.

[bondhugula]

Remove duplicated comments.

[dcaballe]

LGTM. Just a minor comment.

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

98

Not sure I understand the LogicalResult -> void in runOnBlock. If runOnBlock is calling affineDataCopyGenerate which is now returning a LogicalResult, wouldn't it make sense that runOnBlock continues to return LogicalResult and and propagates the return value of affineDataCopyGenerate when appropriate?

[bondhugula]

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

98

runOnBlock makes multiple calls to affineDataCopyGenerate and although we could make it return success if everything succeeded, the overall failure or success status of it isn't that meaningful or used in any way at the top level. runOnFunction itself calls runOnBlock multiple times and the IR is still valid and useful if even on partial failure. So I decided to not propagate "aggregate status" up further.

mlir/include/mlir/Transforms/LoopUtils.h

struct AffineCopyOptions {
   unsigned tagMemorySpace;
   // Capacity of the fast memory space in bytes.
   uint64_t fastMemCapacityBytes;
 };
 
 /// Performs explicit copying for the contiguous sequence of operations in the
 /// block iterator range [`begin', `end'), where `end' can't be past the
 /// terminator of the block (since additional operations are potentially
-/// inserted right before `end`. Returns the total size of fast memory space
-/// buffers used. `copyOptions` provides various parameters, and the output
-/// argument `copyNests` is the set of all copy nests inserted, each represented
-/// by its root affine.for. Since we generate alloc's and dealloc's for all fast
-/// buffers (before and after the range of operations resp. or at a hoisted
-/// position), all of the fast memory capacity is assumed to be available for
-/// processing this block range. When 'filterMemRef' is specified, copies are
-/// only generated for the provided MemRef.
-uint64_t affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
+/// inserted right before `end`. `copyOptions` provides various parameters, and
+/// the output argument `copyNests` is the set of all copy nests inserted, each
+/// represented by its root affine.for. Since we generate alloc's and dealloc's
+/// for all fast buffers (before and after the range of operations resp. or at a
+/// hoisted position), all of the fast memory capacity is assumed to be
+/// available for processing this block range. When 'filterMemRef' is specified,
+/// copies are only generated for the provided MemRef. Returns success if the
+/// explicit copying succeeded for all memrefs on which affine load/stores were
+/// encountered.
+LogicalResult affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
                                      const AffineCopyOptions &copyOptions,
                                      Optional<Value> filterMemRef,
                                      DenseSet<Operation *> &copyNests);
 
 /// A convenience version of affineDataCopyGenerate for all ops in the body of
 /// an AffineForOp.
-uint64_t affineDataCopyGenerate(AffineForOp forOp,
+LogicalResult affineDataCopyGenerate(AffineForOp forOp,
                                      const AffineCopyOptions &copyOptions,
                                      Optional<Value> filterMemRef,
                                      DenseSet<Operation *> &copyNests);
 
 /// Result for calling generateCopyForMemRegion.
 struct CopyGenerateResult {
   // Number of bytes used by alloc.
   uint64_t sizeInBytes;
 
   // The newly created buffer allocation.
   Operation *alloc;

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

explicit AffineDataCopyGeneration(unsigned slowMemorySpace,
     this->slowMemorySpace = slowMemorySpace;
     this->fastMemorySpace = fastMemorySpace;
     this->tagMemorySpace = tagMemorySpace;
     this->minDmaTransferSize = minDmaTransferSize;
     this->fastMemoryCapacity = fastMemCapacityBytes / 1024;
   }
 
   void runOnFunction() override;
-  LogicalResult runOnBlock(Block *block, DenseSet<Operation *> &copyNests);
+  void runOnBlock(Block *block, DenseSet<Operation *> &copyNests);
 
   // Constant zero index to avoid too many duplicates.
   Value zeroIndex = nullptr;
 };
 
 } // namespace
 
 /// Generates copies for memref's living in 'slowMemorySpace' into newly created

 mlir::createAffineDataCopyGenerationPass() {
   return std::make_unique<AffineDataCopyGeneration>();
 }
 
 /// Generate copies for this block. The block is partitioned into separate
 /// ranges: each range is either a sequence of one or more operations starting
 /// and ending with an affine load or store op, or just an affine.forop (which
 /// could have other affine for op's nested within).
-LogicalResult
-AffineDataCopyGeneration::runOnBlock(Block *block,
+void AffineDataCopyGeneration::runOnBlock(Block *block,
                                           DenseSet<Operation *> &copyNests) {
   if (block->empty())
-    return success();
+    return;
 
   uint64_t fastMemCapacityBytes =
       fastMemoryCapacity != std::numeric_limits<uint64_t>::max()
           ? fastMemoryCapacity * 1024
           : fastMemoryCapacity;
   AffineCopyOptions copyOptions = {generateDma, slowMemorySpace,
                                    fastMemorySpace, tagMemorySpace,
                                    fastMemCapacityBytes};
 
-  // Every affine.forop in the block starts and ends a block range for copying;
+  // Every affine.for op in the block starts and ends a block range for copying;
   // in addition, a contiguous sequence of operations starting with a
   // load/store op but not including any copy nests themselves is also
   // identified as a copy block range. Straightline code (a contiguous chunk of
   // operations excluding AffineForOp's) are always assumed to not exhaust
   // memory. As a result, this approach is conservative in some cases at the
   // moment; we do a check later and report an error with location info.
   // TODO: An 'affine.if' operation is being treated similar to an
   // operation. 'affine.if''s could have 'affine.for's in them;
   // treat them separately.
 
   // Get to the first load, store, or for op (that is not a copy nest itself).
   auto curBegin =
       std::find_if(block->begin(), block->end(), [&](Operation &op) {
         return isa<AffineLoadOp, AffineStoreOp, AffineForOp>(op) &&
                copyNests.count(&op) == 0;
       });
 
   // Create [begin, end) ranges.
   auto it = curBegin;
   while (it != block->end()) {
     AffineForOp forOp;
     // If you hit a non-copy for loop, we will split there.
     if ((forOp = dyn_cast<AffineForOp>(&*it)) && copyNests.count(forOp) == 0) {
       // Perform the copying up unti this 'for' op first.
-      affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/it, copyOptions,
+      (void)affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/it, copyOptions,
                                    /*filterMemRef=*/llvm::None, copyNests);
 
       // Returns true if the footprint is known to exceed capacity.
       auto exceedsCapacity = [&](AffineForOp forOp) {
         Optional<int64_t> footprint =
             getMemoryFootprintBytes(forOp,
                                     /*memorySpace=*/0);
         return (footprint.hasValue() &&
                 static_cast<uint64_t>(footprint.getValue()) >
                     fastMemCapacityBytes);
       };
 
       // If the memory footprint of the 'affine.for' loop is higher than fast
       // memory capacity (when provided), we recurse to copy at an inner level
       // until we find a depth at which footprint fits in fast mem capacity. If
       // the footprint can't be calculated, we assume for now it fits. Recurse
       // inside if footprint for 'forOp' exceeds capacity, or when
       // skipNonUnitStrideLoops is set and the step size is not one.
       bool recurseInner = skipNonUnitStrideLoops ? forOp.getStep() != 1
                                                  : exceedsCapacity(forOp);
       if (recurseInner) {
         // We'll recurse and do the copies at an inner level for 'forInst'.
         // Recurse onto the body of this loop.
-        (void)runOnBlock(forOp.getBody(), copyNests);
+        runOnBlock(forOp.getBody(), copyNests);
       } else {
         // We have enough capacity, i.e., copies will be computed for the
         // portion of the block until 'it', and for 'it', which is 'forOp'. Note
         // that for the latter, the copies are placed just before this loop (for
         // incoming copies) and right after (for outgoing ones).
 
         // Inner loop copies have their own scope - we don't thus update
         // consumed capacity. The footprint check above guarantees this inner
         // loop's footprint fits.
-        affineDataCopyGenerate(/*begin=*/it, /*end=*/std::next(it), copyOptions,
+        (void)affineDataCopyGenerate(/*begin=*/it, /*end=*/std::next(it),
+                                     copyOptions,
                                      /*filterMemRef=*/llvm::None, copyNests);
       }
       // Get to the next load or store op after 'forOp'.
       curBegin = std::find_if(std::next(it), block->end(), [&](Operation &op) {
         return isa<AffineLoadOp, AffineStoreOp, AffineForOp>(op) &&
                copyNests.count(&op) == 0;
       });
       it = curBegin;
     } else {
       assert(copyNests.count(&*it) == 0 &&
              "all copy nests generated should have been skipped above");
       // We simply include this op in the current range and continue for more.
       ++it;
     }
   }
 
   // Generate the copy for the final block range.
   if (curBegin != block->end()) {
     // Can't be a terminator because it would have been skipped above.
     assert(!curBegin->hasTrait<OpTrait::IsTerminator>() &&
            "can't be a terminator");
     // Exclude the affine.yield - hence, the std::prev.
-    affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/std::prev(block->end()),
-                           copyOptions, /*filterMemRef=*/llvm::None, copyNests);
+    (void)affineDataCopyGenerate(/*begin=*/curBegin,
+                                 /*end=*/std::prev(block->end()), copyOptions,
+                                 /*filterMemRef=*/llvm::None, copyNests);
   }
-
-  return success();
 }
 
 void AffineDataCopyGeneration::runOnFunction() {
   FuncOp f = getFunction();
   OpBuilder topBuilder(f.getBody());
   zeroIndex = topBuilder.create<arith::ConstantIndexOp>(f.getLoc(), 0);
 
   // Nests that are copy-in's or copy-out's; the root AffineForOps of those
   // nests are stored herein.
   DenseSet<Operation *> copyNests;
 
   // Clear recorded copy nests.
   copyNests.clear();
 
   for (auto &block : f)
-    (void)runOnBlock(&block, copyNests);
+    runOnBlock(&block, copyNests);
 
   // Promote any single iteration loops in the copy nests and collect
   // load/stores to simplify.
   SmallVector<Operation *, 4> copyOps;
   for (Operation *nest : copyNests)
     // With a post order walk, the erasure of loops does not affect
     // continuation of the walk or the collection of load/store ops.
     nest->walk([&](Operation *op) {

mlir/lib/Transforms/Utils/LoopUtils.cpp

for (unsigned d = 0; d < rank; d++) {
     auto dimSize = memRefType.getDimSize(d);
     assert(dimSize > 0 && "filtered dynamic shapes above");
     regionCst->addBound(FlatAffineConstraints::LB, d, 0);
     regionCst->addBound(FlatAffineConstraints::UB, d, dimSize - 1);
   }
   return true;
 }
 
-/// Performs explicit copying for the contiguous sequence of operations in the
-/// block iterator range [`begin', `end'), where `end' can't be past the
-/// terminator of the block (since additional operations are potentially
-/// inserted right before `end`. Returns the total size of fast memory space
-/// buffers used. `copyOptions` provides various parameters, and the output
-/// argument `copyNests` is the set of all copy nests inserted, each represented
-/// by its root affine.for. Since we generate alloc's and dealloc's for all fast
-/// buffers (before and after the range of operations resp. or at a hoisted
-/// position), all of the fast memory capacity is assumed to be available for
-/// processing this block range. When 'filterMemRef' is specified, copies are
-/// only generated for the provided MemRef.
-uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
+LogicalResult mlir::affineDataCopyGenerate(Block::iterator begin,
                                            Block::iterator end,
                                            const AffineCopyOptions &copyOptions,
                                            Optional<Value> filterMemRef,
                                            DenseSet<Operation *> &copyNests) {
   if (begin == end)
-    return 0;
+    return success();
 
   assert(begin->getBlock() == std::prev(end)->getBlock() &&
          "Inconsistent block begin/end args");
   assert(end != end->getBlock()->end() && "end can't be the block terminator");
 
   Block *block = begin->getBlock();
 
   // Copies will be generated for this depth, i.e., symbolic in all loops
block->walk(begin, end, [&](Operation *opInst) {
     if (region->isWrite() && !existsInWrite) {
       writeRegions[region->memref] = std::move(region);
     } else if (!region->isWrite() && !existsInRead) {
       readRegions[region->memref] = std::move(region);
     }
   });
 
   if (error) {
-    begin->emitError(
-        "copy generation failed for one or more memref's in this block\n");
-    return 0;
+    LLVM_DEBUG(begin->emitError(
+        "copy generation failed for one or more memref's in this block\n"));
+    return failure();
   }
 
   uint64_t totalCopyBuffersSizeInBytes = 0;
   bool ret = true;
   auto processRegions =
       [&](const SmallMapVector<Value, std::unique_ptr<MemRefRegion>, 4>
               &regions) {
         for (const auto &regionEntry : regions) {
auto processRegions =
           }
           ret = ret & succeeded(iRet);
         }
       };
   processRegions(readRegions);
   processRegions(writeRegions);
 
   if (!ret) {
-    begin->emitError(
-        "copy generation failed for one or more memref's in this block\n");
-    return totalCopyBuffersSizeInBytes;
+    LLVM_DEBUG(begin->emitError(
+        "copy generation failed for one or more memref's in this block\n"));
+    return failure();
   }
 
   // For a range of operations, a note will be emitted at the caller.
   AffineForOp forOp;
   uint64_t sizeInKib = llvm::divideCeil(totalCopyBuffersSizeInBytes, 1024);
   if (llvm::DebugFlag && (forOp = dyn_cast<AffineForOp>(&*begin))) {
-    forOp.emitRemark()
+    LLVM_DEBUG(forOp.emitRemark()
                << sizeInKib
-        << " KiB of copy buffers in fast memory space for this block\n";
+               << " KiB of copy buffers in fast memory space for this block\n");
   }
 
   if (totalCopyBuffersSizeInBytes > copyOptions.fastMemCapacityBytes) {
     StringRef str = "Total size of all copy buffers' for this block "
                     "exceeds fast memory capacity\n";
     block->getParentOp()->emitWarning(str);
   }
 
-  return totalCopyBuffersSizeInBytes;
+  return success();
 }
 
 // A convenience version of affineDataCopyGenerate for all ops in the body of
 // an AffineForOp.
-uint64_t mlir::affineDataCopyGenerate(AffineForOp forOp,
+LogicalResult mlir::affineDataCopyGenerate(AffineForOp forOp,
                                            const AffineCopyOptions &copyOptions,
                                            Optional<Value> filterMemRef,
                                            DenseSet<Operation *> &copyNests) {
   return affineDataCopyGenerate(forOp.getBody()->begin(),
                                 std::prev(forOp.getBody()->end()), copyOptions,
                                 filterMemRef, copyNests);
 }
 
 LogicalResult mlir::generateCopyForMemRegion(
     const MemRefRegion &memrefRegion, Operation *analyzedOp,
     const AffineCopyOptions &copyOptions, CopyGenerateResult &result) {

mlir/test/Dialect/Affine/dma-generate.mlir

func @dma_unknown_size(%arg0: memref<?x?xf32>) {
   %M = memref.dim %arg0, %c0 : memref<? x ? x f32>
   %N = memref.dim %arg0, %c0 : memref<? x ? x f32>
   affine.for %i = 0 to %M {
     affine.for %j = 0 to %N {
       // If this loop nest isn't tiled, the access requires a non-constant DMA
       // size -- not yet implemented.
       // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
       affine.load %arg0[%i, %j] : memref<? x ? x f32>
-      // expected-error@-6 {{copy generation failed for one or more memref's in this block}}
     }
   }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @dma_memref_3d
 func @dma_memref_3d(%arg0: memref<1024x1024x1024xf32>) {
   affine.for %i = 0 to 1024 {
     affine.for %j = 0 to 1024 {
       affine.for %k = 0 to 1024 {
         %idx = affine.apply affine_map<(d0) -> (d0 mod 128)>(%i)
         %idy = affine.apply affine_map<(d0) -> (d0 mod 128)>(%j)
         %idz = affine.apply affine_map<(d0) -> (d0 mod 128)>(%k)
         // DMA with nested striding (or emulating with loop around strided DMA)
         // not yet implemented.
         // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<1024x1024x1024xf32>
         %v = affine.load %arg0[%idx, %idy, %idz] : memref<1024 x 1024 x 1024 x f32>
-        // expected-error@-10 {{copy generation failed for one or more memref's in this block}}
       }
     }
   }
   return
 }
 
 // -----
 

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

void TestAffineDataCopy::runOnFunction() {
 
   AffineCopyOptions copyOptions = {/*generateDma=*/false,
                                    /*slowMemorySpace=*/0,
                                    /*fastMemorySpace=*/0,
                                    /*tagMemorySpace=*/0,
                                    /*fastMemCapacityBytes=*/32 * 1024 * 1024UL};
   DenseSet<Operation *> copyNests;
   if (clMemRefFilter) {
-    affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
+    if (failed(affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(),
+                                      copyNests)))
+      return;
   } else if (clTestGenerateCopyForMemRegion) {
     CopyGenerateResult result;
     MemRefRegion region(loopNest.getLoc());
-    (void)region.compute(load, /*loopDepth=*/0);
-    (void)generateCopyForMemRegion(region, loopNest, copyOptions, result);
+    if (failed(region.compute(load, /*loopDepth=*/0)))
+      return;
+    if (failed(generateCopyForMemRegion(region, loopNest, copyOptions, result)))
+      return;
   }
 
   // Promote any single iteration loops in the copy nests and simplify
   // load/stores.
   SmallVector<Operation *, 4> copyOps;
   for (Operation *nest : copyNests) {
     // With a post order walk, the erasure of loops does not affect
     // continuation of the walk or the collection of load/store ops.