[mlir][tensor] More static padding sizes for PadOp

[springerm]

Summary

Improve tensor::createPadHighOp: more static padding sizes when possible.

Depends On: D154199

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
-  auto zero = b.createOrFold<arith::ConstantIndexOp>(loc, 0);
-  SmallVector<OpFoldResult> low(type.getRank(), zero);
-  SmallVector<OpFoldResult> high(type.getRank(), zero);
+  SmallVector<OpFoldResult> low(type.getRank(), b.getIndexAttr(0));
+  SmallVector<OpFoldResult> high(type.getRank(), b.getIndexAttr(0));
   for (const auto &en : enumerate(type.getShape())) {
     // Pad only the static dimensions of the result tensor type.
     if (ShapedType::isDynamic(en.value()))
       continue;
     // Compute the padding width.
     AffineExpr d0;
     bindDims(b.getContext(), d0);
     OpFoldResult sz = tensor::getMixedSize(b, loc, source, en.index());
     high[en.index()] =
-        affine::makeComposedAffineApply(b, loc, en.value() - d0, {sz})
-            .getResult();
+        affine::makeComposedFoldedAffineApply(b, loc, en.value() - d0, {sz});
   }
   return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
 }
 
 SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
                                                         Location loc,
                                                         Value rankedTensor) {
   auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());

mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir

 func.func @simple_pad_and_pack(%input: tensor<5x1xf32>, %output: tensor<1x1x8x2xf32>, %pad: f32) -> tensor<1x1x8x2xf32> {
   %0 = tensor.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<5x1xf32> -> tensor<1x1x8x2xf32>
   return %0 : tensor<1x1x8x2xf32>
 }
 // CHECK-LABEL: func.func @simple_pad_and_pack
 // CHECK-SAME:    %[[SRC:[a-zA-Z0-9]+]]
 // CHECK-SAME:    %[[DEST:[a-zA-Z0-9]+]]
 // CHECK-SAME:    %[[PAD_VAL:[a-zA-Z0-9]+]]
-// CHECK-DAG:     %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG:     %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG:     %[[C3:.+]] = arith.constant 3 : index
-// CHECK:         %[[PAD:.+]] = tensor.pad %[[SRC]] low[%[[C0]], %[[C0]]] high[%[[C3]], %[[C1]]]
+// CHECK:         %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[3, 1]
 // CHECK:           tensor.yield %[[PAD_VAL]]
 // CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<8x2xf32>
 // CHECK:         %[[TRANSP:.+]] = linalg.transpose
 // CHECK-SAME:      ins(%[[PAD]] : tensor<8x2xf32>)
 // CHECK-SAME:      outs(%[[EMPTY]] : tensor<8x2xf32>)
 // CHECK-SAME:      permutation = [0, 1]
 // CHECK:         %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]
 // CHECK-SAME:      [0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1]

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

func.func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
   //      CHECK: %[[T0:.*]] = tensor.extract_slice %
   //      CHECK: %[[T1:.*]] = tensor.extract_slice %
   //      CHECK: %[[T2:.*]] = tensor.extract_slice %
   %1 = tensor.extract_slice %arg0[%iv0, %iv2] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
   %2 = tensor.extract_slice %arg1[%iv2, %iv1] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
   %3 = tensor.extract_slice %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<24x25xf32> to 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>)
func.func @static_sizes_output_divisible_on_empty_op(%arg0: tensor<24x12xf32>,
   //      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>)

 }
 
 // -----
 
 // Check that the padding can be applied even when the output argument of the
 // linalg op is not produced by an empty op or an extract_slice op.
 
 // CHECK-DAG: #[[$MAP_MIN:.*]] = affine_map<(d0) -> (-d0 + 2044, 16)>
-// CHECK-DAG: #[[$MAP_C0:.*]] = affine_map<() -> (0)>
 // CHECK-DAG: #[[$MAP_TO_16:.*]] = affine_map<(d0) -> (-d0 + 16)>
 // CHECK-LABEL: @outs_not_produced_by_empty_or_extract_slice(
 // CHECK-SAME: %[[A:[^: ]*]]: tensor<128x2044xf32>,
 // CHECK-SAME: %[[B:[^: ]*]]: tensor<2044x128xf32>)
 func.func @outs_not_produced_by_empty_or_extract_slice(%a : tensor<128x2044xf32>, %b : tensor<2044x128xf32>) -> tensor<128x128xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %0 = tensor.empty() : tensor<128x128xf32>
   %9 = linalg.fill ins(%cst : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32>
 
   %c0 = arith.constant 0 : index
   %c16 = arith.constant 16 : index
   %c2044 = arith.constant 2044 : index
   // CHECK: scf.for %[[ARG3:.*]] = {{.*}} iter_args(%[[ARG4:.*]] = %{{.*}})
   %10 = scf.for %arg3 = %c0 to %c2044 step %c16 iter_args(%arg4 = %9) -> (tensor<128x128xf32>) {
     // CHECK: %[[MIN:.*]] = affine.min #[[$MAP_MIN]](%[[ARG3]])
     %11 = affine.min affine_map<(d0) -> (-d0 + 2044, 16)>(%arg3)
     // CHECK: %[[A_SLICE:.*]] = tensor.extract_slice %[[A]]
     // CHECK: %[[B_SLICE:.*]] = tensor.extract_slice %[[B]]
     %extracted_slice_2 = tensor.extract_slice %a[0, %arg3] [128, %11] [1, 1] : tensor<128x2044xf32> to tensor<128x?xf32>
     %extracted_slice_3 = tensor.extract_slice %b[%arg3, 0] [%11, 128] [1, 1] : tensor<2044x128xf32> to tensor<?x128xf32>
     // CHECK-DAG: %[[CST:.*]] = arith.constant 0.
-    // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 
-    // CHECK-DAG: %[[ZERO:.*]] = affine.apply #[[$MAP_C0]]()
     // CHECK-DAG: %[[TO_16:.*]] = affine.apply #[[$MAP_TO_16]](%[[MIN]])
-    // CHECK: %[[PADDED_A_SLICE:.*]] = tensor.pad %[[A_SLICE]] nofold low[%[[C0]], %[[C0]]] high[%[[ZERO]], %[[TO_16]]]
+    // CHECK: %[[PADDED_A_SLICE:.*]] = tensor.pad %[[A_SLICE]] nofold low[0, 0] high[0, %[[TO_16]]]
     // CHECK: tensor.yield %[[CST]]
     // CHECK: %[[PADDED_B_SLICE:.*]] = tensor.pad %[[B_SLICE]] nofold
     // The output shape is already padded, so actually we shouldn't
     // add anything to the upper bound.
-    // CHECK: %[[ZERO0:.*]] = affine.apply #[[$MAP_C0]]()
-    // CHECK: %[[ZERO1:.*]] = affine.apply #[[$MAP_C0]]()
-    // CHECK: %[[PADDED_ARG4:.*]] = tensor.pad %[[ARG4]] nofold low[{{.*}}] high[%[[ZERO0]], %[[ZERO1]]]
+    // CHECK: %[[PADDED_ARG4:.*]] = tensor.pad %[[ARG4]] nofold low[{{.*}}] high[0, 0]
 
     //      CHECK: %[[T5:.*]] = linalg.matmul
     // CHECK-SAME:              ins(%[[PADDED_A_SLICE]], %[[PADDED_B_SLICE]] : tensor<128x16xf32>, tensor<16x128xf32>)
     // CHECK-SAME:              outs(%[[PADDED_ARG4]] : tensor<128x128xf32>)
     %res = linalg.matmul ins(%extracted_slice_2, %extracted_slice_3 : tensor<128x?xf32>, tensor<?x128xf32>) outs(%arg4 : tensor<128x128xf32>) -> tensor<128x128xf32>
     scf.yield %res : tensor<128x128xf32>
   }
   return %10 : tensor<128x128xf32>
func.func @pack_everything(%arg0: tensor<24x12xf32>,
   //      CHECK: %[[T0:.*]] = tensor.extract_slice %
   //      CHECK: %[[T1:.*]] = tensor.extract_slice %
   //      CHECK: %[[T2:.*]] = tensor.extract_slice %
   %1 = tensor.extract_slice %arg0[%iv0, %iv2] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
   %2 = tensor.extract_slice %arg1[%iv2, %iv1] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
   %3 = tensor.extract_slice %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
 
   //  CHECK-DAG: %[[CST:.*]] = arith.constant 0.
-  //  CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 
   //      CHECK: %[[PAD0:.*]] = tensor.pad %[[T0]] nofold
   //      CHECK: %[[PAD1:.*]] = tensor.pad %[[T1]] nofold
   //      CHECK: %[[PAD2:.*]] = tensor.pad %[[T2]] nofold
 
   //      CHECK: %[[T5:.*]] = linalg.matmul
   // CHECK-SAME:              ins(%[[PAD0]], %[[PAD1]] : tensor<4x7xf32>, tensor<7x5xf32>)
   // CHECK-SAME:              outs(%[[PAD2]] : tensor<4x5xf32>)