[mlir][linalg] Use AffineApplyOp to compute padding width (NFC).

[gysit]

Summary

Use AffineApplyOp instead of SubIOp to compute the padding width when creating a pad tensor operation.

Depends On D113382

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

PadTensorOp PadTensorOp::createPadScalarOp(Type type, Value source, Value pad,
   OpBuilder::InsertionGuard guard(builder);
   builder.createBlock(&region, region.end(), blockArgTypes);
   builder.create<linalg::YieldOp>(loc, pad);
   return padTensorOp;
 }
 
 PadTensorOp PadTensorOp::createPadHighOp(Type type, Value source, Value pad,
                                          bool nofold, Location loc,
-                                         OpBuilder &builder) {
+                                         OpBuilder &b) {
   SmallVector<OpFoldResult, 4> low, high;
   auto rankedTensorType = type.cast<RankedTensorType>();
   assert(rankedTensorType.hasStaticShape());
-  int rank = rankedTensorType.getRank();
-  for (int i = 0; i < rank; ++i) {
-    auto dimOp = builder.createOrFold<tensor::DimOp>(loc, source, i);
-    auto resultDimSize = builder.createOrFold<arith::ConstantIndexOp>(
-        loc, rankedTensorType.getDimSize(i));
-    auto highValue =
-        builder.createOrFold<arith::SubIOp>(loc, resultDimSize, dimOp);
-    high.push_back(highValue);
-    low.push_back(builder.createOrFold<arith::ConstantIndexOp>(loc, 0));
+  for (auto en : enumerate(rankedTensorType.getShape())) {
+    AffineExpr d0;
+    bindDims(b.getContext(), d0);
+    auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
+    Value paddingWidth =
+        makeComposedAffineApply(b, loc, en.value() - d0, {dimOp});
+    high.push_back(paddingWidth);
+    low.push_back(b.createOrFold<arith::ConstantIndexOp>(loc, 0));
   }
   return PadTensorOp::createPadScalarOp(type, source, pad, low, high, nofold,
-                                        loc, builder);
+                                        loc, b);
 }
 
 LogicalResult PadTensorOp::reifyResultShapes(
     OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   Location loc = getLoc();
   auto lowPad = getMixedLowPad();
   auto highPad = getMixedHighPad();
   SmallVector<Value> shapes;

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=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) -> (7, -d0 + 25)>
+// CHECK-DAG: #[[MAP2:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 5)>
+// CHECK-DAG: #[[MAP3:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 7)>
 // CHECK-DAG: #[[DIV6:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 6)>
 #map0 = affine_map<(d0) -> (5, -d0 + 24)>
 #map1 = affine_map<(d0) -> (7, -d0 + 25)>
 
 //      CHECK:  static_sizes
 //      CHECK-DOUBLE:  static_sizes
 // CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<24x12xf32>
 // CHECK-SAME:    %[[ARG1:[0-9a-zA-Z]*]]: tensor<12x25xf32>
%0 = scf.for %arg3 = %c0 to %c24 step %c5 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
     // Packing the first input operand for all values of IV2 (IV2x5x6).
     //      CHECK:  = linalg.init_tensor [2, 5, 6]
     //      CHECK:  %[[PT0:.*]] = scf.for %[[PIV0:[0-9a-z]+]] =
     //        CHECK:   %[[PIDX0:.*]] = affine.apply #[[DIV6]](%[[PIV0]])
     //        CHECK:   %[[TS0:.*]] = affine.min #[[MAP0]](%[[IV0]])
     //        CHECK:   %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
     //   CHECK-SAME:                                     %[[IV0]], %[[PIV0]]
     //   CHECK-SAME:                                     %[[TS0]], 6
-    //        CHECK:   %[[V0:.*]] = arith.subi %[[C5]], %[[TS0]]
+    //        CHECK:   %[[V0:.*]] = affine.apply #[[MAP2]](%[[TS0]])
     //        CHECK:   %[[T1:.*]] = linalg.pad_tensor %[[T0]] nofold {{.*}} high[%[[V0]]
     //        CHECK:   %[[T2:.*]] = tensor.insert_slice %[[T1:.*]] into %{{.*}}[%[[PIDX0]], 0, 0]
     //        CHECK:   scf.yield %[[T2:.*]]
 
     //      CHECK:  scf.for %[[IV1:[0-9a-zA-Z]*]] =
     %1 = scf.for %arg5 = %c0 to %c25 step %c7 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
 
       // Packing the second input operand for all values of IV2 (IV2x6x7).
       //      CHECK:  = linalg.init_tensor [2, 6, 7]
       //      CHECK:  %[[PT1:.*]] = scf.for %[[PIV1:[0-9a-z]+]] =
       //        CHECK:   %[[PIDX1:.*]] = affine.apply #[[DIV6]](%[[PIV1]])
       //        CHECK:   %[[TS1:.*]] = affine.min #[[MAP1]](%[[IV1]])
       //        CHECK:   %[[T3:.*]] = tensor.extract_slice %[[ARG1]]
       //   CHECK-SAME:                                     %[[PIV1]], %[[IV1]]
       //   CHECK-SAME:                                     6, %[[TS1]]
-      //        CHECK:   %[[V1:.*]] = arith.subi %[[C7]], %[[TS1]]
+      //        CHECK:   %[[V1:.*]] = affine.apply #[[MAP3]](%[[TS1]])
       //        CHECK:   %[[T4:.*]] = linalg.pad_tensor %[[T3]] nofold {{.*}} high[%[[C0]], %[[V1]]
       //        CHECK:   %[[T5:.*]] = tensor.insert_slice %[[T4:.*]] into %{{.*}}[%[[PIDX1]], 0, 0]
       //        CHECK:   scf.yield %[[T5:.*]]
 
       //      CHECK:  scf.for %[[IV2:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG4:.*]] =
       %2 = scf.for %arg7 = %c0 to %c12 step %c6 iter_args(%arg8 = %arg6) -> (tensor<24x25xf32>) {
         %3 = affine.min #map0(%arg3)
 
func @static_sizes(%arg0: tensor<24x12xf32>,
   return %0 : tensor<24x25xf32>
 }
 
 // -----
 
 // CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0)[s0] -> (5, -d0 + s0)>
 // CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0)[s0] -> (6, -d0 + s0)>
 // CHECK-DAG: #[[MAP2:[0-9a-z]+]] = affine_map<(d0)[s0] -> (7, -d0 + s0)>
+// CHECK-DAG: #[[MAP3:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 5)>
+// CHECK-DAG: #[[MAP4:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 6)>
+// CHECK-DAG: #[[MAP5:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 7)>
 // CHECK-DAG: #[[SDIV6:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 6)>
 // CHECK-DAG: #[[DDIV6:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 6)>
 #map0 = affine_map<(d0)[s0] -> (5, -d0 + s0)>
 #map1 = affine_map<(d0)[s0] -> (6, -d0 + s0)>
 #map2 = affine_map<(d0)[s0] -> (7, -d0 + s0)>
 
 //      CHECK:  dynamic_sizes
 //      CHECK-DOUBLE:  dynamic_sizes
%3 = scf.for %arg3 = %c0 to %0 step %c5 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
     //      CHECK:  = linalg.init_tensor [%[[PS0]], 5, 6]
     //      CHECK:  %[[PT0:.*]] = scf.for %[[PIV0:[0-9a-z]+]] =
     //        CHECK:   %[[PIDX0:.*]] = affine.apply #[[DDIV6]](%[[PIV0]])
     //        CHECK:   %[[TS0:.*]] = affine.min #[[MAP0]](%[[IV0]])[%[[D0]]
     //        CHECK:   %[[TS1:.*]] = affine.min #[[MAP1]](%[[PIV0]])[%[[D1]]
     //        CHECK:   %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
     //   CHECK-SAME:                                     %[[IV0]], %[[PIV0]]
     //   CHECK-SAME:                                     %[[TS0]], %[[TS1]]
-    //        CHECK:   %[[V0:.*]] = arith.subi %[[C5]], %[[TS0]]
-    //        CHECK:   %[[V1:.*]] = arith.subi %[[C6]], %[[TS1]]
+    //        CHECK:   %[[V0:.*]] = affine.apply #[[MAP3]](%[[TS0]])
+    //        CHECK:   %[[V1:.*]] = affine.apply #[[MAP4]](%[[TS1]])
     //        CHECK:   %[[T1:.*]] = linalg.pad_tensor %[[T0]] nofold {{.*}} high[%[[V0]], %[[V1]]
     //        CHECK:   %[[T2:.*]] = tensor.insert_slice %[[T1:.*]] into %{{.*}}[%[[PIDX0]], 0, 0]
     //        CHECK:   scf.yield %[[T2:.*]]
 
     //      CHECK:  scf.for %[[IV1:[0-9a-zA-Z]*]] =
     %4 = scf.for %arg5 = %c0 to %2 step %c7 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
 
       // Packing the second input operand for all values of IV2 (IV2x6x7).
       //      CHECK:  = linalg.init_tensor [%[[PS0]], 6, 7]
       //      CHECK:  %[[PT1:.*]] = scf.for %[[PIV1:[0-9a-z]+]] =
       //        CHECK:   %[[PIDX1:.*]] = affine.apply #[[DDIV6]](%[[PIV1]])
       //        CHECK:   %[[TS2:.*]] = affine.min #[[MAP1]](%[[PIV1]])[%[[D1]]
       //        CHECK:   %[[TS3:.*]] = affine.min #[[MAP2]](%[[IV1]])[%[[D2]]
       //        CHECK:   %[[T3:.*]] = tensor.extract_slice %[[ARG1]]
       //   CHECK-SAME:                                     %[[PIV1]], %[[IV1]]
       //   CHECK-SAME:                                     %[[TS2]], %[[TS3]]
-      //        CHECK:   %[[V2:.*]] = arith.subi %[[C6]], %[[TS2]]
-      //        CHECK:   %[[V3:.*]] = arith.subi %[[C7]], %[[TS3]]
+      //        CHECK:   %[[V2:.*]] = affine.apply #[[MAP4]](%[[TS2]])
+      //        CHECK:   %[[V3:.*]] = affine.apply #[[MAP5]](%[[TS3]])
       //        CHECK:   %[[T4:.*]] = linalg.pad_tensor %[[T3]] nofold {{.*}} high[%[[V2]], %[[V3]]
       //        CHECK:   %[[T5:.*]] = tensor.insert_slice %[[T4:.*]] into %{{.*}}[%[[PIDX1]], 0, 0]
       //        CHECK:   scf.yield %[[T5:.*]]
 
       //      CHECK:  scf.for %[[IV2:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG4:.*]] =
       %5 = scf.for %arg7 = %c0 to %1 step %c6 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
         %6 = affine.min #map0(%arg3)[%0]
         %7 = affine.min #map1(%arg7)[%1]

mlir/test/Dialect/Linalg/pad.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns="test-pad-pattern pack-paddings=1,1,0 hoist-paddings=0,0,0" -cse -canonicalize -split-input-file | FileCheck %s
 
 // CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (7, -d0 + 12)>
+// CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 7)>
 #map = affine_map<(d0) -> (7, -d0 + 12)>
 
 //      CHECK:  static_sizes_output_divisible
 // CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<24x12xf32>
 // CHECK-SAME:    %[[ARG1:[0-9a-zA-Z]*]]: tensor<12x25xf32>
 func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
                                     %arg1: tensor<12x25xf32>,
                                     %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
%1 = scf.for %arg5 = %c0 to %c25 step %c5 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
         //      CHECK:   %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
         //      CHECK:   %[[T1:.*]] = tensor.extract_slice %[[ARG1]]
         //      CHECK:   %[[T2:.*]] = tensor.extract_slice %[[ARG4]]
         %4 = tensor.extract_slice %arg0[%arg3, %arg7] [4, %3] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
         %5 = tensor.extract_slice %arg1[%arg7, %arg5] [%3, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
         %6 = tensor.extract_slice %arg8[%arg3, %arg5] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
 
         // Check statically sized matmul inputs with partially divisible sizes are padded.
-        //      CHECK:   %[[V0:.*]] = arith.subi %[[C7]], %[[TS2]]
+        //      CHECK:   %[[V0:.*]] = affine.apply #[[MAP1]](%[[TS2]])
         //      CHECK:   %[[T3:.*]] = linalg.pad_tensor %[[T0]] nofold
         // CHECK-SAME:                  [%[[C0]], %[[C0]]]
         // CHECK-SAME:                  [%[[C0]], %[[V0]]
         //      CHECK:   %[[T4:.*]] = linalg.pad_tensor %[[T1]] nofold
 
         // Check the statically sized matmul output with fully divisible sizes is not padded.
         //      CHECK:   %[[T5:.*]] = linalg.matmul
         // CHECK-SAME:                  ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>)
%0 = scf.for %arg3 = %c0 to %c24 step %c4 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
     scf.yield %1 : tensor<24x25xf32>
   }
   return %0 : tensor<24x25xf32>
 }
 
 // -----
 
 // CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (7, -d0 + 25)>
+// CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 7)>
 #map = affine_map<(d0) -> (7, -d0 + 25)>
 
 //      CHECK:  static_sizes_input_divisible
 func @static_sizes_input_divisible(%arg0: tensor<24x12xf32>,
                                    %arg1: tensor<12x25xf32>,
                                    %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   //  CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
   //  CHECK-DAG: %[[C7:.*]] = arith.constant 7
%1 = scf.for %arg5 = %c0 to %c25 step %c7 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
         //      CHECK:   %[[TS1:.*]] = affine.min #[[MAP0]](%[[IV1]])
         %4 = affine.min #map(%arg5)
         %5 = tensor.extract_slice %arg1[%arg7, %arg5] [6, %4] [1, 1] : tensor<12x25xf32> to tensor<6x?xf32>
 
         //      CHECK:   %[[T0:.*]] = tensor.extract_slice %[[ARG4]]
         %6 = tensor.extract_slice %arg8[%arg3, %arg5] [4, %4] [1, 1] : tensor<24x25xf32> to tensor<4x?xf32>
 
         // Check the statically sized matmul output with partially divisible sizes is padded.
-        //      CHECK:   %[[V0:.*]] = arith.subi %[[C7]], %[[TS1]]
+        //      CHECK:   %[[V0:.*]] = affine.apply #[[MAP1]](%[[TS1]])
         //      CHECK:   %[[T1:.*]] = linalg.pad_tensor %[[T0]] low
         // CHECK-SAME:                  [%[[C0]], %[[C0]]]
         // CHECK-SAME:                  [%[[C0]], %[[V0]]
 
         //      CHECK:   %[[T2:.*]] = linalg.matmul
         // CHECK-SAME:                  outs(%[[T1]] : tensor<4x7xf32>)
         //      CHECK:   %[[T3:.*]] = tensor.extract_slice %[[T2]]
         //      CHECK:   %[[T4:.*]] = tensor.insert_slice %[[T3]]
func @static_sizes_input_divisible(%arg0: tensor<24x12xf32>,
   return %0 : tensor<24x25xf32>
 }
 
 // -----
 
 // CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0)[s0] -> (5, -d0 + s0)>
 // CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0)[s0] -> (7, -d0 + s0)>
 // CHECK-DAG: #[[MAP2:[0-9a-z]+]] = affine_map<(d0)[s0] -> (6, -d0 + s0)>
+// CHECK-DAG: #[[MAP3:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 5)>
+// CHECK-DAG: #[[MAP4:[0-9a-z]+]] = affine_map<(d0) -> (-d0 + 6)>
+
 #map0 = affine_map<(d0)[s0] -> (5, -d0 + s0)>
 #map1 = affine_map<(d0)[s0] -> (6, -d0 + s0)>
 #map2 = affine_map<(d0)[s0] -> (7, -d0 + s0)>
 
 //      CHECK:  dynamic_sizes
 // CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<?x?xf32>
 // CHECK-SAME:    %[[ARG1:[0-9a-zA-Z]*]]: tensor<?x?xf32>
 // CHECK-SAME:    %[[ARG2:[0-9a-zA-Z]*]]: tensor<?x?xf32>
%4 = scf.for %arg5 = %c0 to %2 step %c7 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
         %6 = affine.min #map0(%arg3)[%0]
         %7 = affine.min #map1(%arg7)[%1]
         %8 = tensor.extract_slice %arg0[%arg3, %arg7] [%6, %7] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
         %9 = affine.min #map2(%arg5)[%2]
         %10 = tensor.extract_slice %arg1[%arg7, %arg5] [%7, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
         %11 = tensor.extract_slice %arg8[%arg3, %arg5] [%6, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
 
         // Check all matmul operands are padded.
-        //      CHECK:   %[[V0:.*]] = arith.subi %[[C5]], %[[TS0]]
-        //      CHECK:   %[[V1:.*]] = arith.subi %[[C6]], %[[TS2]]
+        //      CHECK:   %[[V0:.*]] = affine.apply #[[MAP3]](%[[TS0]])
+        //      CHECK:   %[[V1:.*]] = affine.apply #[[MAP4]](%[[TS2]])
         //      CHECK:   %[[T3:.*]] = linalg.pad_tensor %{{.*}} nofold
         // CHECK-SAME:                  [%[[C0]], %[[C0]]]
         // CHECK-SAME:                  [%[[V0]], %[[V1]]
         //      CHECK:   %[[T4:.*]] = linalg.pad_tensor %{{.*}} nofold
         //      CHECK:   %[[T5:.*]] = linalg.pad_tensor %{{.*}} low
 
         // Check the dynamic matmul has been erased.
         //  CHECK-NOT:   = linalg.matmul {{.*}} tensor<?x?xf32>