[mlir][linalg] Improve hoist padding buffer size computation.

[gysit]

Summary

Adapt the Fourier Motzkin elimination to take into account affine computations happening outside of the cloned loop nest.

Depends On D112713

[gysit]

Cleanup.

[nicolasvasilache]

Very nice!

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

341

Similar refactoring, comment, checkm re side-effects should be done before landing please.

[gysit]

Address comments by improving comments and making the index slice computation more concise.

[gysit]

Rebase.

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

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 *>
 HoistingAnalysis::getIndexingLoops(PadTensorOp padTensorOp,
                                    tensor::ExtractSliceOp sliceOp) {
   // Set of all values used for index computation.
   SetVector<Value> indexEdges;
 
-  // Helper function that adds all index operands of an operation to
-  // `indexEdges`. An operand is an index operand if it is of index type.
-  auto addIndexOperandsToIndexEdges = [&](Operation *op) {
-    for (Value operand : op->getOperands())
-      if (operand.getType().isIndex())
-        indexEdges.insert(operand);
-  };
-
   // 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;
   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);
+      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);
+        addIndexOperandsToIndexEdges(op, indexEdges);
         indexingLoops.insert(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);
+      addIndexOperandsToIndexEdges(op, indexEdges);
   }
   return indexingLoops;
 }
 
 static bool isDefinedOutsideOrConstant(scf::ForOp outer, Value v) {
   return outer.isDefinedOutsideOfLoop(v) || v.getDefiningOp<ConstantOp>();
 }
 
 /// For each loop in `loops`, determine the ops involved in the construction of
 /// its upper bound---up to the outerLimit loop--- and fold them as new
 /// inequalities in the constraint set.
 /// This is achieved by computing the backwardSlice of the loop's upper bound
 /// and iteratively folding each op in reverse topological order to guarantee
 /// use-def ordering.
 /// As operations are folded in, their result is projected out of the
 /// constraints set.
 /// The following operations are supported:
 ///   - scf::ForOp are simply skipped.
 ///   - AffineApplyOp are composed to replace the result by an equality.
 ///   - AffineMinOp are composed by adding each entry as an upper bound.
+/// Additionally, the following terminal operations are handled:
+///   - DimOp and ConstantOp are skipped.
 /// If any other operation is met, return failure.
 // TODO: extend on a per-need basis.
 static LogicalResult
 foldUpperBoundsIntoConstraintsSet(FlatAffineValueConstraints &constraints,
                                   scf::ForOp outerLimit,
                                   ArrayRef<scf::ForOp> loops) {
   SetVector<Value> toProjectOut;
   for (scf::ForOp loop : loops) {
     auto ub = loop.upperBound();
-    if (isDefinedOutsideOrConstant(outerLimit, ub))
-      continue;
 
-    // Compute a backward slice up to, but not including, `outerLimit`.
-    SetVector<Operation *> backwardSlice;
-    getBackwardSlice(ub, &backwardSlice, [&](Operation *op) {
-      return outerLimit->isProperAncestor(op);
+    // Set of all values used for index computation.
+    SetVector<Value> indexEdges;
+    indexEdges.insert(ub);
+
+    // Compute the backward slice `indexSlice` containing the index computation
+    // performed to obtain the upper bound `ub`. Starting from `ub` add the
+    // index operands of an operation to `indexEdges` if one of its results is
+    // an index edge. Otherwise, stop the slice computation. For a loop, check
+    // if its induction variable is an index edge.
+    //
+    // Example:
+    // ```
+    // %c0 = arith.constant 0
+    // scf.for %i = %c0 to ...
+    //   scf.for %j = %c0 to ...
+    //     %ub = affine.min #map(%i)
+    //     scf.for %k = %c0 to %ub
+    // ```
+    // After computing the backward slice we obtain:
+    // indexEdges = [%ub, %i, %c0]
+    // indexSlice = [arith.constant 0, scf.for %i, affine.min #map(%i)]
+    SetVector<Operation *> indexSlice;
+    getBackwardSlice(ub, &indexSlice, [&](Operation *op) {
+      // Continue only along the index operands of the ForOp.
+      if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+        // Consider only loops part of the enclosing loops.
+        if (!outerLimit->isAncestor(op))
+          return false;
+        if (!indexEdges.contains(forOp.getInductionVar()))
+          return false;
+        addIndexOperandsToIndexEdges(op, indexEdges);
+        return true;
+      }
+      // All supported index operations have one result.
+      assert(op->getNumResults() == 1 &&
+             "expect operations to have one result");
+      if (!indexEdges.contains(op->getResult(0)))
+        return false;
+      addIndexOperandsToIndexEdges(op, indexEdges);
+      return true;
     });
-    backwardSlice.insert(ub.getDefiningOp());
+    indexSlice.insert(ub.getDefiningOp());
 
     // Iterate over all ops in the slice and compose them in the constraints.
-    for (Operation *op : llvm::reverse(backwardSlice)) {
-      if (!isa<scf::ForOp, AffineApplyOp, AffineMinOp>(op))
-        return failure();
-      if (isa<scf::ForOp>(op))
+    for (Operation *op : llvm::reverse(indexSlice)) {
+      // All ForOps have previously been added to the constraints and ConstantOp
+      // and DimOp are terminals of the index computation.
+      if (isa<scf::ForOp, arith::ConstantOp, tensor::DimOp>(op))
         continue;
-      // Ensure there is a
+      // Check all index computation operations are supported.
+      if (!isa<AffineApplyOp, AffineMinOp>(op))
+        return failure();
+      // Ensure there is an id.
       auto ensureIdFailed = [&](Value v) {
         if (constraints.containsId(v)) {
           unsigned pos;
           constraints.findId(v, &pos);
           return pos >= constraints.getNumDimIds();
         }
         constraints.appendDimId(v);
         return false;
       };
 
       // Ensure all ids exist and add results for later projection.
       if (llvm::any_of(op->getResults(), ensureIdFailed) ||
           llvm::any_of(op->getOperands(), ensureIdFailed))
         return failure();
 
       // All supported ops have 1 result.
       // TODO: extend when needed.
+      assert(op->getNumResults() == 1 &&
+             "expect operations to have one result");
       toProjectOut.insert(op->getResult(0));
 
       // Compose supported ops.
       if (auto affineApplyOp = dyn_cast<AffineApplyOp>(op)) {
         AffineValueMap avm(affineApplyOp.getAffineMap(),
                            affineApplyOp.getOperands(),
                            affineApplyOp.getResult());
         if (failed(constraints.composeMap(&avm)))

mlir/test/Dialect/Linalg/hoist-padding.mlir

func @dot(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<f32>)
   %c0 = arith.constant 0 : index
   %c4 = arith.constant 4 : index
   %c2 = arith.constant 2 : index
   %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
 
   //      CHECK: scf.for %[[I:[0-9a-z]+]] =
   //
   //      CHECK:   %[[MR8:.*]] = affine.min #[[$MIN_REST8]](%[[I]])
-  //      CHECK:   %[[D0:.*]] = affine.apply #[[$DIV4]](%[[MR8]])
   // Init tensor and pack.
-  //      CHECK:   %[[INIT_PACKED_A:.*]] = linalg.init_tensor [%[[D0]], 2, 2] : tensor<?x2x2xf32>
-  //      CHECK:   %[[CAST_INIT_PACKED_A:.*]] = tensor.cast %[[INIT_PACKED_A]] : tensor<?x2x2xf32> to tensor<?x?x2xf32>
+  //      CHECK:   %[[INIT_PACKED_A:.*]] = linalg.init_tensor [2, 2, 2] : tensor<2x2x2xf32>
+  //      CHECK:   %[[CAST_INIT_PACKED_A:.*]] = tensor.cast %[[INIT_PACKED_A]] : tensor<2x2x2xf32> to tensor<?x?x2xf32>
   //      CHECK:   %[[PACKED_A:.*]] = scf.for %[[II:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[CAST_INIT_PACKED_A]]) -> (tensor<?x?x2xf32>) {
   //      CHECK:     scf.for %[[III:[0-9a-z]+]] =
   //      CHECK:       tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
   //
-  //      CHECK:   %[[D0_2:.*]] = affine.apply #[[$DIV4]](%[[MR8]])
   // Init tensor and pack.
-  //      CHECK:   %[[INIT_PACKED_B:.*]] = linalg.init_tensor [%[[D0_2]], 2, 2] : tensor<?x2x2xf32>
-  //      CHECK:   %[[CAST_INIT_PACKED_B:.*]] = tensor.cast %[[INIT_PACKED_B]] : tensor<?x2x2xf32> to tensor<?x?x2xf32>
+  //      CHECK:   %[[INIT_PACKED_B:.*]] = linalg.init_tensor [2, 2, 2] : tensor<2x2x2xf32>
+  //      CHECK:   %[[CAST_INIT_PACKED_B:.*]] = tensor.cast %[[INIT_PACKED_B]] : tensor<2x2x2xf32> to tensor<?x?x2xf32>
   //      CHECK:   %[[PACKED_B:.*]] = scf.for %[[II_2:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[CAST_INIT_PACKED_B]]) -> (tensor<?x?x2xf32>) {
   //      CHECK:     scf.for %[[III_2:[0-9a-z]+]] =
   //      CHECK:       tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
   // Compute.
   //      CHECK:   scf.for %[[II_3:[0-9a-z]+]] =
   //      CHECK:     scf.for %[[III_3:[0-9a-z]+]] = {{.*}} iter_args(%[[C:.*]] = %{{.*}}) -> (tensor<f32>) {
   //      CHECK:       %[[IDX0:.*]] = affine.apply #[[$DIV4]](%[[II_3]])
   //      CHECK:       %[[IDX1:.*]] = affine.apply #[[$DIV2]](%[[III_3]])

mlir/test/Dialect/Linalg/pad-and-hoist.mlir

%0 = scf.for %arg3 = %c0 to %c24 step %c15 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
 
     //    CHECK-DOUBLE:  scf.for %[[IV1:[0-9a-zA-Z]*]] =
     %1 = scf.for %arg5 = %c0 to %c25 step %c16 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
       %2 = affine.min #map0(%arg3)
       %3 = affine.min #map1(%arg5)
       %4 = tensor.extract_slice %arg6[%arg3, %arg5] [%2, %3] [1, 1] : tensor<24x25xf32> to tensor<?x?xf32>
 
       // Packing the second input operand.
-      //    CHECK-DOUBLE:  = linalg.init_tensor [%{{.*}}, 12, 6]
+      //    CHECK-DOUBLE:  = linalg.init_tensor [3, 12, 6]
       //    CHECK-DOUBLE:  %[[PT1:.*]] = scf.for %[[PIV1:[0-9a-z]+]] =
       //      CHECK-DOUBLE:   %[[PIDX1:.*]] = affine.apply #[[DIV6]](%[[PIV1]])
       //      CHECK-DOUBLE:   %[[T3:.*]] = tensor.extract_slice %[[ARG1]]
       //      CHECK-DOUBLE:   %[[T4:.*]] = linalg.pad_tensor %[[T3]] nofold
       //      CHECK-DOUBLE:   %[[T5:.*]] = tensor.insert_slice %[[T4:.*]] into %{{.*}}[%[[PIDX1]], 0, 0]
       //      CHECK-DOUBLE:   scf.yield %[[T5:.*]]
 
       //    CHECK-DOUBLE:  scf.for %[[IV2:[0-9a-zA-Z]*]] =