[mlir][linalg] Hoist padding simplifications (NFC).

[gysit]

Summary

Remove unused members and store the indexing and packing loops in SmallVector.

mlir/include/mlir/Dialect/Linalg/Transforms/HoistPadding.h

 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 class Value;
 
 namespace linalg {
 class PadTensorOp;
 
-/// Mechanically hoist padding operations on tensors by `nLoops` into a new,
+/// Mechanically hoist padding operations on tensors by `numLoops` into a new,
 /// generally larger tensor. This achieves packing of multiple padding ops into
 /// a larger tensor. On success, `padTensorOp` is replaced by the cloned version
 /// in the packing loop so the caller can continue reasoning about the padding
 /// operation.
 ///
 /// Example in pseudo-mlir:
 /// =======================
 ///

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

 ///   4. The backward slice from the pad op to the scf::ForOp to hoist above is
 ///      empty.
 ///   5. The source tensor of pad op is not defined by an extract slice op.
 ///   6. The source tensor of the extract slice op is not defined outside of
 ///      the outermost enclosing scf::ForOp.
 ///   7. There is no enclosing scf::ForOp that indexes the padded data.
 /// Other cases succeed and will trigger hoisting of the pad op.
 struct HoistingAnalysis {
-  HoistingAnalysis(PadTensorOp padTensorOp, int nLevels);
+  HoistingAnalysis(PadTensorOp padTensorOp, int numLoops);
 
   bool isValid() { return valid; }
 
   /// Footprint of the packedTensor, computed from the packingLoops and
   /// `backwardSlice`.
   FailureOr<SmallVector<Value>> getPackedTensorSizes(ImplicitLocOpBuilder &b);
 
-  /// The padTensorOp that needs to be hoisted.
-  PadTensorOp padTensorOp;
-
-  /// The maximum number of immediately enclosing scf::ForOp to hoist over.
-  int nLevels;
-
   /// The outermost loop, determined by `nLevels` above which `padTensorOp` will
   /// be hoisted.
   scf::ForOp outermostEnclosingForOp;
 
   /// Backward slice rooted at `padTensorOp` and nested under
   /// `outermostEnclosingForOp`.
   SetVector<Operation *> backwardSlice;
 
   /// The scf::ForOp immediately enclosing `padTensorOp` such that:
   ///  1. they are nested under `outermostEnclosingForOp` (inclusive)
   ///  2. whose induction variable is used, directly or indirectly, in the
   ///     computation of `padTensorOp`.
   /// The span of these loops determines the footprint of the packed tensor.
-  /// SmallSetVector<scf::ForOp> packingLoops;
-  SetVector<scf::ForOp, SmallVector<scf::ForOp>, DenseSet<Operation *>>
-      packingLoops;
+  SmallVector<scf::ForOp> packingLoops;
 
 private:
   /// Returns the loops in `backwardSlice` used to index the padded data. The
   /// method starts from `padTensorOp` and `sliceOp`, follows the use-def
   /// chains of their index operands, and stores any enclosing loop whose
   /// induction variable is part of the walked index computation.
   ///
   /// Example:
   /// ```
   /// %source = linalg.fill(%cst, %arg0)
   /// scf.for %i
   ///   scf.for %j
   ///     scf.for %k // not used to index %source!
   ///       %ubi = affine.min #map(%i)
   ///       %ubj = affine.min #map(%j)
   ///       %slice = tensor.extract_slice %source [%i, %j] [%ubi, %ubj]
   ///       %padded_slice = linalg.pad_tensor %slice
   /// ```
   /// getIndexingLoops(%padded_slice, %slice) returns [scf.for %i, scf.for %j]
-  SetVector<Operation *> getIndexingLoops(PadTensorOp padTensorOp,
+  SmallVector<scf::ForOp> getIndexingLoops(PadTensorOp padTensorOp,
                                            tensor::ExtractSliceOp sliceOp);
 
   /// Encodes whether the analysis is valid and hoisting can proceed.
   bool valid;
 };
 
 /// Return true if all uses of `padTensorOp` are an input tensor of some
 /// LinalgOp.
 static bool isOnlyUsedAsInputOfLinalgOp(PadTensorOp padTensorOp) {
LLVM_DEBUG(
         DBGS() << "loops: ";
         outermostEnclosingForOp.getInductionVar().printAsOperand(dbgs(), state);
         dbgs() << "\n");
     reverseEnclosingLoops.push_back(outermostEnclosingForOp);
     nextEnclosingOp = outermostEnclosingForOp->getParentOp();
   }
 }
 
-HoistingAnalysis::HoistingAnalysis(PadTensorOp padTensorOp, int nLevels)
-    : padTensorOp(padTensorOp), nLevels(nLevels), valid(false) {
-  AsmState state(padTensorOp->getParentOfType<mlir::FuncOp>());
-  (void)state;
+HoistingAnalysis::HoistingAnalysis(PadTensorOp padTensorOp, int numLoops) {
+  valid = false;
 
   // Bail on any use that isn't an input of a Linalg op.
   // Hoisting of inplace updates happens after vectorization.
   if (!isOnlyUsedAsInputOfLinalgOp(padTensorOp))
     return;
 
   // Get at most nLevels of immediately enclosing loops.
   SmallVector<scf::ForOp> reverseEnclosingLoops;
-  getAtMostNEnclosingLoops(padTensorOp, nLevels, reverseEnclosingLoops);
+  getAtMostNEnclosingLoops(padTensorOp, numLoops, reverseEnclosingLoops);
   if (reverseEnclosingLoops.empty()) {
     LLVM_DEBUG(DBGS() << "No immediately enclosing loop -> skip\n");
     return;
   }
 
   outermostEnclosingForOp = reverseEnclosingLoops.back();
 
   // Get all the ops in the backwards slice starting from `padTensorOp` and that
if (!sliceOp) {
     return;
   }
   if (!outermostEnclosingForOp.isDefinedOutsideOfLoop(sliceOp.source())) {
     LLVM_DEBUG(DBGS() << "Source not defined outside of loops -> skip\n");
     return;
   }
 
   // Search the loops found in `backwardSlice` used to index the padded data.
-  SetVector<Operation *> indexingLoops = getIndexingLoops(padTensorOp, sliceOp);
+  SmallVector<scf::ForOp> indexingLoops =
+      getIndexingLoops(padTensorOp, sliceOp);
 
   // Add only the loops part of `indexingLoops` to the packing loops. All other
   // loops are not used to index the padded data and consequently access the
   // same data in every loop iteration. Adding them to the packing loops would
   // increase the cache footprint of the packed data by storing the same data
   // multiple times.
-  for (scf::ForOp forOp : llvm::reverse(reverseEnclosingLoops)) {
-    if (indexingLoops.contains(forOp))
-      packingLoops.insert(forOp);
-  }
-  assert(indexingLoops.size() == packingLoops.size() &&
+  for (scf::ForOp forOp : llvm::reverse(reverseEnclosingLoops))
+    if (!indexingLoops.empty() && indexingLoops.back() == forOp)
+      packingLoops.push_back(indexingLoops.pop_back_val());
+  assert(indexingLoops.empty() &&
          "expect the all indexing loops are enclosing loops");
+
   if (packingLoops.empty()) {
     LLVM_DEBUG(DBGS() << "Cannot find a packing loop -> skip\n");
     return;
   }
 
   // The analysis is valid and hoisting can occur.
   valid = true;
 }
 
 /// Add all index operands of `operation` to `indexEdges`. An index operand is
 /// an operand of type index.
 static void addIndexOperandsToIndexEdges(Operation *operation,
                                          SetVector<Value> &indexEdges) {
   for (Value operand : operation->getOperands())
     if (operand.getType().isIndex())
       indexEdges.insert(operand);
 }
 
-SetVector<Operation *>
+SmallVector<scf::ForOp>
 HoistingAnalysis::getIndexingLoops(PadTensorOp padTensorOp,
                                    tensor::ExtractSliceOp sliceOp) {
   // Set of all values used for index computation.
   SetVector<Value> indexEdges;
 
   // Starting from `padTensorOp` and `sliceOp` walk the use-def edges of index
   // type in `backwardSlice`. Add the index operands of an operation to
   // `indexEdges` if one of its results is an index edge found so far and store
   // all loops part of the index computation to `indexingLoops`.
   //
   // Example:
   // ```
   // %source = linalg.fill(%cst, %arg0)
   // scf.for %i
   //   scf.for %j
   //     scf.for %k // not used to index %source!
   //       %ubi = affine.min #map(%i)
   //       %ubj = affine.min #map(%j)
   //       %slice = tensor.extract_slice %source [%i, %j] [%ubi, %ubj]
   //       %padded_slice = linalg.pad_tensor %slice
   // ```
   // After iterating `backwardSlice` we obtain:
   // indexEdges = [%i, %j, %ubi, %ubj]
   // indexingLoops = [scf.for %i, scf.for %j]
-  SetVector<Operation *> indexingLoops;
+  SmallVector<scf::ForOp> indexingLoops;
   for (Operation *op : llvm::reverse(backwardSlice)) {
     // Add the index operands of `padTensorOp` and `sliceOp` to start the
     // exploration of the index computation.
     if (op == padTensorOp || op == sliceOp) {
       addIndexOperandsToIndexEdges(op, indexEdges);
       continue;
     }
     // Add the index operands of the loop if its induction variable is
     // used for index computation. Additionally, insert the loop into
     // `indexingLoops`
     if (auto forOp = dyn_cast<scf::ForOp>(op)) {
       if (indexEdges.contains(forOp.getInductionVar())) {
         addIndexOperandsToIndexEdges(op, indexEdges);
-        indexingLoops.insert(forOp);
+        indexingLoops.push_back(forOp);
         continue;
       }
     }
     // Add the index operands of all other operations if at least one result is
     // used for index computation.
     if (llvm::any_of(op->getResults(),
                      [&](Value result) { return indexEdges.contains(result); }))
       addIndexOperandsToIndexEdges(op, indexEdges);
auto constraints = FlatAffineValueConstraints::getHyperrectangular(
           packingLoops, [](scf::ForOp op) { return op.getInductionVar(); })),
       llvm::to_vector<8>(llvm::map_range(
           packingLoops, [](scf::ForOp op) { return op.lowerBound(); })),
       llvm::to_vector<8>(llvm::map_range(
           packingLoops, [](scf::ForOp op) { return op.upperBound(); })));
 
   // Iteratively try to fold the upper bounds into the constraints set.
   if (failed(foldUpperBoundsIntoConstraintsSet(
-          constraints, outermostEnclosingForOp, packingLoops.getArrayRef())))
+          constraints, outermostEnclosingForOp, packingLoops)))
     return failure();
 
   int nPackedLoops = packingLoops.size();
   SmallVector<AffineMap> lbs(nPackedLoops), ubs(nPackedLoops);
   // Compute the bounds of the first positions, assuming the others are fixed.
   constraints.getSliceBounds(/*pos=*/0, /*num=*/nPackedLoops,
                              outermostEnclosingForOp->getContext(), &lbs, &ubs);
 
if (hasNoEffects &&
         (op->getNumRegions() == 0 || isa<linalg::PadTensorOp>(op))) {
       b.clone(*op, bvm);
       continue;
     }
     // TODO: support more cases as they appear.
     auto forOp = dyn_cast<scf::ForOp>(op);
     assert(forOp && "Expected scf::ForOp when hoisting pad ops");
     // Unused loop, just skip it.
-    if (!analysis.packingLoops.contains(forOp))
+    if (!llvm::is_contained(analysis.packingLoops, forOp))
       continue;
 
     auto clonedForOp =
         b.create<scf::ForOp>(loc, bvm.lookupOrDefault(forOp.lowerBound()),
                              bvm.lookupOrDefault(forOp.upperBound()),
                              bvm.lookupOrDefault(forOp.step()), packedTensor);
     // Map the induction var, region args and results to the `clonedForOp`.
     bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());