[mlir][linalg] Perform checks early in hoist padding.

[gysit]

Summary

Instead of checking for unexpected operations (any operation with a region except for scf::For and padTensorOp or operations with a memory effect) while cloning the packing loop nest perform the checks early. Update dropNonIndexDependencies to check for unexpected operations. Additionally, check all of these operations have index type operands only.

Depends On D114428

[gysit]

Add additional test with a loop that prevents hoisting.

[gysit]

Remove temporary check.

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

 using namespace mlir::linalg;
 
 /// Analysis class to support PadTensorOp hoisting across multiple enclosing
 /// loops. The failure conditions are:
 ///   1. Pad op has a use that is not an input of a LinalgOp.
 ///   2. Pad op does not have a constant padding value.
 ///   3. There is no immediately enclosing scf::ForOp.
 ///   4. The backward slice from the pad op to the scf::ForOp to hoist above
-///      contains an unknown op with a region.
+///      contains an unknown op with non index type operands, a region, or a
+///      memory effect.
 ///   5. The backward slice from the pad op to the scf::ForOp to hoist above is
 ///      empty.
 ///   6. The source tensor of pad op is not defined by an extract slice op.
 ///   7. The source tensor of the extract slice op is not defined outside of
 ///      the outermost enclosing scf::ForOp.
 ///   8. There is no enclosing scf::ForOp that indexes the padded data.
 /// Other cases succeed and will trigger hoisting of the pad op.
 struct HoistingAnalysis {
struct HoistingAnalysis {
   SmallVector<scf::ForOp> packingLoops;
 
 private:
   /// Drop any non-index dependencies of `padTensorOp` and `sliceOp` from
   /// `backwardSlice`. The method follows the use-def chains of the index
   /// operands consumed by `padTensorOp` and `sliceOp` and drops the operations
   /// not part of this index computation. Afterwards, the filtered
   /// `backwardSlice` contains only the loops whose induction variable is used,
-  /// directly or indirectly, to index the padded tensor.
+  /// directly or indirectly, to index the padded tensor. The method returns
+  /// failure if the filtered backward slice contains an unexpected operation.
   ///
   /// Example:
   /// ```
   /// %source = linalg.fill(%cst, %arg0)
   /// scf.for %i
   ///   %unrelated = linalg.fill(%cst, %arg1)    // not used to index %source!
   ///   scf.for %j (%arg2 = %unrelated)
   ///     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
   /// ```
   /// dropNonIndexDependencies(%padded_slice, %slice)
   /// removes [scf.for %k, linalg.fill(%cst, %arg1)] from backwardSlice.
-  void dropNonIndexDependencies(PadTensorOp padTensorOp,
+  LogicalResult dropNonIndexDependencies(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) {
getBackwardSlice(padTensorOp.getOperation(), &backwardSlice,
                    });
   if (backwardSlice.empty())
     return;
   backwardSlice.insert(padTensorOp.getOperation());
 
   // Remove all ops in the backward slice that are not used to index the padded
   // tensor. In particular, keep `padTensorOp`, `sliceOp`, and the loop and
   // affine operations used for the index computation.
-  dropNonIndexDependencies(padTensorOp, sliceOp);
+  if (failed(dropNonIndexDependencies(padTensorOp, sliceOp)))
+    return;
 
   // Add only the loops part of the filtered `backwardSlice` 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 (backwardSlice.contains(forOp))
       packingLoops.push_back(forOp);
   if (packingLoops.empty()) {
     LLVM_DEBUG(DBGS() << "Cannot find a packing loop -> skip\n");
     return;
   }
 
   // The analysis is valid and hoisting can occur.
   valid = true;
 }
 
-void HoistingAnalysis::dropNonIndexDependencies(
-    PadTensorOp padTensorOp, tensor::ExtractSliceOp sliceOp) {
+LogicalResult
+HoistingAnalysis::dropNonIndexDependencies(PadTensorOp padTensorOp,
+                                           tensor::ExtractSliceOp sliceOp) {
   // Set of all values used for index computation.
   SetVector<Value> indexEdges;
 
   // Add all index operands of `operation` to `indexEdges`. An index operand is
   // an operand of type index.
   auto addIndexOperandsToIndexEdges = [&](Operation *operation) {
     for (Value operand : operation->getOperands())
       if (operand.getType().isIndex())
for (Operation *op : llvm::reverse(backwardSlice)) {
     // exploration of the index computation.
     if (op == padTensorOp || op == sliceOp) {
       addIndexOperandsToIndexEdges(op);
       continue;
     }
     // Add the index operands of the loop if its induction variable is
     // used for index computation.
     if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-      if (indexEdges.contains(forOp.getInductionVar())) {
+      if (!hasIndexResult(op) && indexEdges.contains(forOp.getInductionVar())) {
         addIndexOperandsToIndexEdges(op);
         continue;
       }
     }
     // Add the index operands of all other operations if at least one result is
     // used for index computation.
     if (hasIndexResult(op)) {
       addIndexOperandsToIndexEdges(op);
+      // Check the operands of the remaining operations all have index type.
+      if (llvm::any_of(op->getOperandTypes(),
+                       [](Type type) { return !type.isIndex(); })) {
+        LLVM_DEBUG(DBGS() << "Unsupported op with non index type operands: "
+                          << op << " -> skip\n");
+        return failure();
+      }
+      // Check the remaining operations do not have regions or memory effects.
+      auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op);
+      bool hasMemoryEffect = effectInterface && !effectInterface.hasNoEffect();
+      if (hasMemoryEffect || op->getNumRegions() != 0) {
+        LLVM_DEBUG(DBGS() << "Unsupported op with region or memory effect: "
+                          << op << " -> skip\n");
+        return failure();
+      }
       continue;
     }
     // Remove all other operation not used by the index computation except for
     // constant operations that may be padding values used by `padTensorOp`.
     if (!isa<arith::ConstantOp>(op))
       backwardSlice.remove(op);
   }
+  return success();
 }
 
 SmallVector<Value>
 HoistingAnalysis::getPackedTensorSizes(ImplicitLocOpBuilder &b) {
   SmallVector<Value> dynamicTensorSizes;
 
   // Upper bound the packing loop lengths to size the packed tensor. Taking
   // upper bounds can make the sizes of the packed tensor independent of the
FailureOr<Value> mlir::linalg::hoistPaddingOnTensors(PadTensorOp opToHoist,
   BlockAndValueMapping bvm;
   // Stack step 1. iteratively clone loops and push `packedTensor`.
   for (Operation *op : analysis.backwardSlice) {
     // Specifically sit out in the extract_slice(packedTensor) case: this is the
     // piece we seek to replace.
     if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(op))
       if (bvm.lookupOrDefault(sliceOp.source()) == packedTensor)
         continue;
-    auto effects = dyn_cast<MemoryEffectOpInterface>(op);
-    bool hasNoEffects = !effects || effects.hasNoEffect();
-    if (hasNoEffects &&
-        (op->getNumRegions() == 0 || isa<linalg::PadTensorOp>(op))) {
+    // Clone all operations except it is a loop.
+    auto forOp = dyn_cast<scf::ForOp>(op);
+    if (!forOp) {
       b.clone(*op, bvm);
       continue;
     }
-    // TODO: support more cases as they appear.
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    assert(forOp && llvm::is_contained(analysis.packingLoops, forOp) &&
-           "expect an scf::ForOp that is a packing loop");
-
+    // Create a packing loop that takes `packedTensor` as iteration argument.
     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());
     bvm.map(forOp.getRegionIterArgs(), clonedForOp.getRegionIterArgs());
     bvm.map(forOp.getResults(), clonedForOp.getResults());

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

%1 = scf.for %arg5 = %c0 to %c25 step %c7 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
       }
       scf.yield %2 : tensor<24x25xf32>
     }
     scf.yield %1 : tensor<24x25xf32>
   }
   return %0 : tensor<24x25xf32>
 }
 
+// -----
+
+#map0 = affine_map<(d0) -> (5, -d0 + 24)>
+#map1 = affine_map<(d0) -> (7, -d0 + 25)>
+#map2 = affine_map<(d0) -> (-d0 + 5)>
+#map3 = affine_map<(d0) -> (-d0 + 7)>
+
+//      CHECK:  unexpected_operation
+//      CHECK-DOUBLE:  unexpected_operation
+// CHECK-SAME:    %[[ARG3:[0-9a-zA-Z]*]]: memref<?xindex>
+// CHECK-SAME:    %[[ARG4:[0-9a-zA-Z]*]]: i32
+func @unexpected_operation(%arg0: tensor<24x12xf32>,
+                           %arg1: tensor<12x25xf32>,
+                           %arg2: tensor<24x25xf32>,
+                           %arg3: memref<?xindex>,
+                           %arg4: i32) -> tensor<24x25xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c5 = arith.constant 5 : index
+  %c7 = arith.constant 7 : index
+  %c6 = arith.constant 6 : index
+  %c24 = arith.constant 24 : index
+  %c25 = arith.constant 25 : index
+  %c12 = arith.constant 12 : index
+  %c0 = arith.constant 0 : index
+
+  //      CHECK:  scf.for %[[IV0:[0-9a-zA-Z]*]] =
+  %0 = scf.for %arg5 = %c0 to %c24 step %c5 iter_args(%arg6 = %arg2) -> (tensor<24x25xf32>) {
+
+    // CHECK-NEXT:  scf.for %[[IV1:[0-9a-zA-Z]*]] =
+    %1 = scf.for %arg7 = %c0 to %c25 step %c7 iter_args(%arg8 = %arg6) -> (tensor<24x25xf32>) {
+
+      // CHECK-NEXT:  scf.for %[[IV2:[0-9a-zA-Z]*]] =
+      %2 = scf.for %arg9 = %c0 to %c12 step %c6 iter_args(%arg10 = %arg8) -> (tensor<24x25xf32>) {
+        %3 = affine.min #map0(%arg5)
+        %4 = tensor.extract_slice %arg0[%arg5, %arg9] [%3, 6] [1, 1] : tensor<24x12xf32> to tensor<?x6xf32>
+        %5 = affine.min #map1(%arg7)
+        %6 = tensor.extract_slice %arg1[%arg9, %arg7] [6, %5] [1, 1] : tensor<12x25xf32> to tensor<6x?xf32>
+        %7 = tensor.extract_slice %arg10[%arg5, %arg7] [%3, %5] [1, 1] : tensor<24x25xf32> to tensor<?x?xf32>
+        %8 = affine.apply #map2(%3)
+
+        // Check cannot hoist due to unexpected operation with memory effect.
+        //      CHECK: %[[IDX0:.*]] = memref.load %[[ARG3]]
+        //      CHECK: %[[T0:.*]] = linalg.pad_tensor {{.*}}, %[[IDX0]]
+        %9 = memref.load %arg3[%c0] : memref<?xindex>
+        %10 = linalg.pad_tensor %4 nofold low[%c0, %c0] high[%8, %9]  {
+        ^bb0(%arg11: index, %arg12: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x6xf32> to tensor<5x6xf32>
+        %11 = affine.apply #map3(%5)
+
+        // Check cannot hoist due to unexpected operation with non index operand.
+        //      CHECK: %[[IDX1:.*]] = arith.index_cast %[[ARG4]]
+        //      CHECK: %[[T1:.*]] = linalg.pad_tensor {{.*}}[%[[IDX1]]
+        %12 = arith.index_cast %arg4 : i32 to index
+        %13 = linalg.pad_tensor %6 nofold low[%c0, %c0] high[%12, %11]  {
+        ^bb0(%arg11: index, %arg12: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<6x?xf32> to tensor<6x7xf32>
+        %14 = linalg.pad_tensor %7 low[%c0, %c0] high[%8, %11]  {
+        ^bb0(%arg11: index, %arg12: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x?xf32> to tensor<5x7xf32>
+
+        //      CHECK:  = linalg.matmul ins(%[[T0]], %[[T1]]
+        %15 = linalg.matmul ins(%10, %13 : tensor<5x6xf32>, tensor<6x7xf32>) outs(%14 : tensor<5x7xf32>) -> tensor<5x7xf32>
+        %16 = tensor.extract_slice %15[0, 0] [%3, %5] [1, 1] : tensor<5x7xf32> to tensor<?x?xf32>
+        %17 = tensor.insert_slice %16 into %arg10[%arg5, %arg7] [%3, %5] [1, 1] : tensor<?x?xf32> into tensor<24x25xf32>
+        scf.yield %17 : tensor<24x25xf32>
+      }
+      scf.yield %2 : tensor<24x25xf32>
+    }
+    scf.yield %1 : tensor<24x25xf32>
+  }
+  return %0 : tensor<24x25xf32>
+}
+
+// -----
+
+#map0 = affine_map<(d0) -> (5, -d0 + 24)>
+#map1 = affine_map<(d0) -> (7, -d0 + 25)>
+#map2 = affine_map<(d0) -> (-d0 + 5)>
+#map3 = affine_map<(d0) -> (-d0 + 7)>
+
+//      CHECK:  unexpected_loop
+//      CHECK-DOUBLE:  unexpected_loop
+// CHECK-SAME:    %[[ARG3:[0-9a-zA-Z]*]]: index
+func @unexpected_loop(%arg0: tensor<24x12xf32>,
+                      %arg1: tensor<12x25xf32>,
+                      %arg2: tensor<24x25xf32>,
+                      %arg3: index) -> tensor<24x25xf32> {
+  %c0 = arith.constant 0 : index
+  %c12 = arith.constant 12 : index
+  %c25 = arith.constant 25 : index
+  %c24 = arith.constant 24 : index
+  %c6 = arith.constant 6 : index
+  %c7 = arith.constant 7 : index
+  %c5 = arith.constant 5 : index
+  %cst = arith.constant 0.000000e+00 : f32
+
+  //      CHECK:  scf.for %[[IV0:[0-9a-zA-Z]*]] =
+  %0 = scf.for %arg4 = %c0 to %c24 step %c5 iter_args(%arg5 = %arg2) -> (tensor<24x25xf32>) {
+
+    // CHECK-NEXT:  scf.for %[[IV1:[0-9a-zA-Z]*]] =
+    %1 = scf.for %arg6 = %c0 to %c25 step %c7 iter_args(%arg7 = %arg5) -> (tensor<24x25xf32>) {
+
+      // Check the padding of the first input operand is hoisted.
+      //      CHECK:  = linalg.pad_tensor
+
+      //      CHECK:  scf.for %[[IV2:[0-9a-zA-Z]*]] =
+      %2 = scf.for %arg8 = %c0 to %c12 step %c6 iter_args(%arg9 = %arg7) -> (tensor<24x25xf32>) {
+        %3 = affine.min #map0(%arg4)
+        %4 = tensor.extract_slice %arg0[%arg4, %arg8] [%3, 6] [1, 1] : tensor<24x12xf32> to tensor<?x6xf32>
+        %5 = affine.min #map1(%arg6)
+        %6 = tensor.extract_slice %arg1[%arg8, %arg6] [6, %5] [1, 1] : tensor<12x25xf32> to tensor<6x?xf32>
+        %7 = tensor.extract_slice %arg9[%arg4, %arg6] [%3, %5] [1, 1] : tensor<24x25xf32> to tensor<?x?xf32>
+        %8 = affine.apply #map2(%3)
+
+        // Check cannot hoist due to unexpected operation that has a region.
+        //      CHECK: %[[IDX0:.*]] = scf.for {{.*}} step %[[ARG3]]
+        //      CHECK: %[[T0:.*]] = linalg.pad_tensor {{.*}}, %[[IDX0]]
+        %9 = scf.for %arg10 = %c0 to %c24 step %arg3 iter_args(%arg11 = %c0) -> (index) {
+          %17 = arith.addi %arg3, %arg11 : index
+          scf.yield %17 : index
+        }
+        %10 = linalg.pad_tensor %4 nofold low[%c0, %c0] high[%8, %9]  {
+        ^bb0(%arg10: index, %arg11: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x6xf32> to tensor<5x6xf32>
+        %11 = affine.apply #map3(%5)
+        %12 = linalg.pad_tensor %6 nofold low[%c0, %c0] high[%c0, %11]  {
+        ^bb0(%arg10: index, %arg11: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<6x?xf32> to tensor<6x7xf32>
+        %13 = linalg.pad_tensor %7 low[%c0, %c0] high[%8, %11]  {
+        ^bb0(%arg10: index, %arg11: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?x?xf32> to tensor<5x7xf32>
+
+        //      CHECK:  = linalg.matmul ins(%[[T0]]
+        %14 = linalg.matmul ins(%10, %12 : tensor<5x6xf32>, tensor<6x7xf32>) outs(%13 : tensor<5x7xf32>) -> tensor<5x7xf32>
+        %15 = tensor.extract_slice %14[0, 0] [%3, %5] [1, 1] : tensor<5x7xf32> to tensor<?x?xf32>
+        %16 = tensor.insert_slice %15 into %arg9[%arg4, %arg6] [%3, %5] [1, 1] : tensor<?x?xf32> into tensor<24x25xf32>
+        scf.yield %16 : tensor<24x25xf32>
+      }
+      scf.yield %2 : tensor<24x25xf32>
+    }
+    scf.yield %1 : tensor<24x25xf32>
+  }
+  return %0 : tensor<24x25xf32>
+}
+