[mlir][linalg] Improve the padding packing loop computation.

[gysit]

Summary

The revision updates the packing loop search in hoist padding. Instead of considering all loops in the backward slice, we now compute a separate backward slice containing the index computations only. This modification ensures we do not add packing loops that are not used to index the packed buffer due to spurious dependencies. One instance where such spurious dependencies can appear is the extract slice operation introduced between the tile loops of a double tiling.

Depends On D112412

[gysit]

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

172

The assertion triggers for the double tiling test case since an affine min computation in the body of the outermost loop actually is in front of the loop. This is OK since they are not dependent meaning both orders are valid topological sorts.

Additionally, I think it may be possible that the outermost loop is not in the backward slice (no data dependency) and we may still want to hoist out of it.

[gysit]

Rename IndexEdge to IndexEdges.

[gysit]

Improve tests and comments.

[nicolasvasilache]

Let's proceed with this modulo refactoring and at least trying to implement a check re side-effect that would work correctly in the scf::ForOp slice computation case.

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

170–209

seems this has reached a natural limit and would benefit from its own function ?

171

A loop

177

Could you rephrase since and with the exception of ?

178

can we have some comment on what this represents ?
If this is to better filter out the fact that the backward slice is too coarse and we want finer control over which values used in scf::For should propagate of not please explain do in the broader context of proper documentation about the behavior of the slice algorithm that you are using.

185

A bit more contained and precise explanation + an example in pseudoIR would be useful here.

The backward slice computation algorithm takes a filter function that ... 
This has the effect of ..
When applied to ... turns all enclosing loops into packing loops regardless of whether they are actually used

etc

Example:
for i
  for j
    for k
      ..

191

are these actual assertions that can never happen or failure conditions where we should abort the transformation?

204

The way you are using this seems to hint at an interface for performing slice analysis that would be more precise.
The trickiness seems related to ops that contain other ops with side effects.

Is the restriction you are making here valid in the presence of say load/store on memref within the scf::ForOp ?
Maybe that is part of what needs to be computed here: is there a nested op with side effects (which in practice we know doesn't happen in the case we are interested in but we need to be conservative).

214

sounds even more like an interface for slicing

[gysit]

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

204

@nicolasvasilache
After reading this I think we should consider switching to an explicit model where we encode the ops that can appear in the backward slice and how we handle them. For example, if the induction variable of a for op is part of the backward slice we only add the bounds and the step to the backward slice but not the iteration argument which is usually completely unrelated.

For the ForOp the fact that it has a body and side effects is not the only problem. It is more the fact that it has a complex semantic and we are not interested in the iteration argument part. For example, an hypothetical cast_tensor_and_affine_apply op would also be problematic and would require special treatment depending on which of its results is found in the backward slice. I.e. we either continue the backward slice computation with the cast or the apply specific operands.

Loads and stores on memrefs are a different beast. They would usually not appear in the backward slice but can in a later point in time indeed make hoisting invalid. We can probably handle this by adding a check what operations may be part of the hoisted / cloned loop nest?

[gysit]

Address comments and use a more fail safe strategy.

[gysit]

Rebase and minor fixes.

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

getBackwardSlice(
           return false;
         }
         return true;
       });
 
   if (analysisFailure || backwardSlice.empty())
     return;
 
-  // Backward slice is a topologically sorted list of ops starting at
-  // `outermostEnclosingForOp`.
-  assert(outermostEnclosingForOp == backwardSlice.front());
-
-  // Filter out the loops whose induction variable is not used to compute the
-  // padded result. As a first approximation, just look for IVs that have no use
-  // in the backwardSlice.
-  // These are the dimensions of reuse that we can exploit to reduce the amount
-  // of copy / memory.
+  // Compute the backward slice used to index the padded operand and filter out
+  // any enclosing loop not used by the indexing. All iterations of a loop not
+  // part of the index computation can share the same packed data. As a result,
+  // detecting these loops is performance critical as it increases cache reuse
+  // and minizes the size of the packing. We approximate the index computation
+  // by keeping track of all values holding index computations or data consumed
+  // by the padding. We add only index and data edges if an operation result
+  // is found in the index or data edges collected so far. Note that this
+  // algorithm overapproximates the index computation since it is not aware of
+  // the op semantics.
+  SetVector<Value> indexEdge;
+  indexEdge.insert(padTensorOp.getResult());
+  SetVector<Operation *> indexSlice;
+  getBackwardSlice(padTensorOp.getOperation(), &indexSlice, [&](Operation *op) {
+    if (!domInfo.dominates(outermostEnclosingForOp, op))
+      return false;
+    // Add ForOps if the induction variable is part of the index computation. At
+    // the moment, ForOps are the only operations with a body containing index
+    // computation.
+    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+      // Adding the block arguments mapped to the iteration arguments is not
+      // needed since we do not support hoisting for output tensors. Verify
+      // neither the block arguments nor the results are part of the index
+      // computation.
+      assert(llvm::none_of(forOp.getIterOpOperands(),
+                           [&](OpOperand &opOperand) {
+                             return indexEdge.contains(
+                                 forOp.getRegionIterArgForOpOperand(opOperand));
+                           }) &&
+             "expect the loop iter args are not part of the index computation");
+      assert(llvm::none_of(
+                 forOp.getResults(),
+                 [&](Value value) { return indexEdge.contains(value); }) &&
+             "expect the loop results are not part of the index computation");
+      if (!indexEdge.contains(forOp.getInductionVar()))
+        return false;
+      indexEdge.insert(forOp.lowerBound());
+      indexEdge.insert(forOp.upperBound());
+      indexEdge.insert(forOp.step());
+      return true;
+    }
+    // All other operations are part of the index computation if the result is
+    // an index edge collected so far.
+    if (llvm::none_of(op->getResults(),
+                      [&](Value value) { return indexEdge.contains(value); }))
+      return false;
+    for (Value value : op->getOperands())
+      indexEdge.insert(value);
+    return true;
+  });
+
+  // Add only the loops part of the index slice to the packing loops.
   for (scf::ForOp forOp : llvm::reverse(reverseEnclosingLoops)) {
-    for (Operation *user : forOp.getInductionVar().getUsers()) {
-      if (backwardSlice.contains(user)) {
-        packingLoops.insert(forOp);
-        break;
-      }
-    }
+    if (indexSlice.contains(forOp))
   }
+  if (packingLoops.empty())
+    valid = false;
 
   // The analysis is valid and hoisting can occur.
   valid = true;
 }
 
 static bool isDefinedOutsideOrConstant(scf::ForOp outer, Value v) {
   return outer.isDefinedOutsideOfLoop(v) || v.getDefiningOp<ConstantOp>();
 }

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

 // RUN: mlir-opt %s -test-linalg-transform-patterns="test-pad-pattern pack-paddings=1,1,0 hoist-paddings=2,1,0" -cse -canonicalize -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-transform-patterns="test-pad-pattern pack-paddings=1,1,0 hoist-paddings=4,3,0" -cse -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-DOUBLE
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-pad-pattern pack-paddings=1,1,0 hoist-paddings=3,2,0" -cse -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-DOUBLE
 
 // CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (5, -d0 + 24)>
 // CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0) -> (8, -d0 + 12)>
 // CHECK-DAG: #[[DIV6:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 6)>
 #map0 = affine_map<(d0) -> (5, -d0 + 24)>
 #map1 = affine_map<(d0) -> (8, -d0 + 12)>
 #map2 = affine_map<(d0) -> (7, -d0 + 25)>
 
func @double_tiling(%arg0: tensor<24x12xf32>,
   %c15 = arith.constant 15 : index
   %c16 = arith.constant 16 : index
   %c24 = arith.constant 24 : index
   %c25 = arith.constant 25 : index
   %c0 = arith.constant 0 : index
   %c5 = arith.constant 5 : index
   %c6 = arith.constant 6 : index
 
-  // Packing the first input operand.
-  //    CHECK-DOUBLE:  = linalg.init_tensor
-  //    CHECK-DOUBLE:  = linalg.pad_tensor {{.*}} nofold
-
   //    CHECK-DOUBLE:  scf.for %[[IV0:[0-9a-zA-Z]*]] =
   %0 = scf.for %arg3 = %c0 to %c24 step %c15 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
 
-    // Packing the second input operand.
-    //    CHECK-DOUBLE:  = linalg.init_tensor
+    // Packing the first input operand.
+    //    CHECK-DOUBLE:  = linalg.init_tensor [3, 5, 12]
     //    CHECK-DOUBLE:  = linalg.pad_tensor {{.*}} nofold
 
     //    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.pad_tensor {{.*}} nofold
+
       //    CHECK-DOUBLE:  scf.for %[[IV2:[0-9a-zA-Z]*]] =
       %5 = scf.for %arg7 = %c0 to %2 step %c5 iter_args(%arg8 = %4) -> (tensor<?x?xf32>) {
 
         //    CHECK-DOUBLE:  scf.for %[[IV3:[0-9a-zA-Z]*]] =
         %7 = scf.for %arg9 = %c0 to %3 step %c6 iter_args(%arg10 = %arg8) -> (tensor<?x?xf32>) {
           %8 = affine.min #map2(%arg7, %2)
           %9 = affine.apply #map3(%arg7, %arg3)
           %10 = tensor.extract_slice %arg0[%9, 0] [%8, 12] [1, 1] : tensor<24x12xf32> to tensor<?x12xf32>