[mlir][linalg] Enhance padding LinalgOps to handle tensor.empty cases.

[hanchung]

rebase

[hanchung]

rebase and fix tests

[hanchung]

fix tests

[nicolasvasilache]

Thanks for extending!
Could you point me to your usage of this abstraction (offline is fine)?

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

131

Can you reshuffle this a bit to early-exit ?

auto sliceOp = dyn_cast...;
auto emptyOp = dyn_cast...;
if (!sliceOp && !emptyOp) return failure();
etc

[hanchung]

In D143043#4105155, @nicolasvasilache wrote:

Thanks for extending!
Could you point me to your usage of this abstraction (offline is fine)?

Thanks for reviewing it!

The example is come from IREE. A tensor.empty op takes a result of affine_min after tiling: https://gist.githubusercontent.com/hanhanW/0ff19d0817b9b41b0a6a0a846db6a53d/raw

%20 = affine.min affine_map<(d0) -> (-d0 + 72, 32)>(%arg4)
%21 = tensor.empty(%20) : tensor<8x?xi32>
%22 = linalg.fill ins(%c0_i32 : i32) outs(%21 : tensor<8x?xi32>) -> tensor<8x?xi32>
%23 = linalg.matmul ins(%extracted_slice, %extracted_slice_1 : tensor<8x24xi8>, tensor<24x?xi8>) outs(%22 : tensor<8x?xi32>) -> tensor<8x?xi32>

If the source is not tensor.empty op, it will become extract_slice(%src) -- which works well with existing patterns. If the source is tensor.empty, some patterns are kicked in and create tensor.empty op that takes bounded sizes. For more details, see https://github.com/iree-org/iree/issues/11880

[hanchung]

address comments

[hanchung]

rebase

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

static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
 
   // Follow the use-def chain if `currOpOperand` is defined by a LinalgOp.
   OpOperand *currOpOperand = opOperand;
   while (auto linalgOp = currOpOperand->get().getDefiningOp<LinalgOp>()) {
     OpResult result = currOpOperand->get().cast<OpResult>();
     currOpOperand = linalgOp.getDpsInitOperand(result.getResultNumber());
   }
 
-  // Fail if `currOpOperand` is not defined by an ExtractSliceOp.
-  auto sliceOp = currOpOperand->get().getDefiningOp<tensor::ExtractSliceOp>();
-  if (!sliceOp)
-    return failure();
-
+  // Fail if `currOpOperand` is not defined by an ExtractSliceOp or EmptyOp.
+  llvm::SmallBitVector droppedDims;
+  SmallVector<OpFoldResult> mixedSizes;
+  if (auto sliceOp =
+          currOpOperand->get().getDefiningOp<tensor::ExtractSliceOp>()) {
     // Compute the dropped dimensions if `sliceOp` is ranke-reducing.
-  llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
-  OffsetSizeAndStrideOpInterface shapedOp = sliceOp;
+    droppedDims = sliceOp.getDroppedDims();
+    mixedSizes = sliceOp.getMixedSizes();
+  } else if (auto emptyOp =
+                 currOpOperand->get().getDefiningOp<tensor::EmptyOp>()) {
+    mixedSizes = emptyOp.getMixedSizes();
+    droppedDims.resize(mixedSizes.size());
+  } else {
+    return failure();
+  }
 
-  // Upper bound the `sliceOp` sizes to obtain a static bounding box.
+  // Upper bound the sizes to obtain a static bounding box.
   SmallVector<int64_t> paddedShape(shape.begin(), shape.end());
   int64_t shapeIdx = 0;
-  for (const auto &en : enumerate(shapedOp.getMixedSizes())) {
+  for (const auto &en : enumerate(mixedSizes)) {
     // Skip dropped dimensions.
     if (droppedDims.test(en.index()))
       continue;
     // Skip dimensions that do not require padding.
     if (!shapeDimsToPad.contains(shapeIdx)) {
       shapeIdx++;
       continue;
     }

mlir/test/Dialect/Linalg/transform-op-pad.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
 }
 
 // -----
 
+#map = affine_map<()[s0] -> (-s0 + 12, 7)>
+
+// CHECK-LABEL: @static_sizes_output_divisible_on_empty_op
+func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
+    %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>, %iv0: index,
+    %iv1: index, %iv2: index) -> tensor<24x25xf32> {
+  %0 = affine.min #map()[%iv2]
+
+  //      CHECK: %[[T0:.*]] = tensor.empty
+  //      CHECK: %[[T1:.*]] = tensor.empty
+  //      CHECK: %[[T2:.*]] = tensor.empty
+  %1 = tensor.empty(%0) : tensor<4x?xf32>
+  %2 = tensor.empty(%0) : tensor<?x5xf32>
+  %3 = tensor.empty() : tensor<4x5xf32>
+
+  //  CHECK-DAG: %[[CST:.*]] = arith.constant 0.
+  //  CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+
+  //      CHECK: %[[T3:.*]] = tensor.pad %[[T0]] nofold
+  //      CHECK: tensor.yield %[[CST]]
+  //      CHECK: %[[T4:.*]] = tensor.pad %[[T1]] nofold
+
+  //      CHECK: %[[T5:.*]] = linalg.matmul
+  // CHECK-SAME:              ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>)
+  // CHECK-SAME:              outs(%[[T2]] : tensor<4x5xf32>)
+  %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
+  %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
+  func.return %5 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
+}
+
+// -----
+
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }