[mlir][linalg] Fix tensor.pad sizes computation in lowerPack.

[hanchung]

Summary

The padded sizes should be derived from destination tensor, not source
tensor. There could be more than one incomplete tile in padding domain.

[hanchung]

poke CI

[chelini]

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

541–543

nit: AffineMap::get(/*dimCount*/2, /*symbolCount*/1, d0 * s0 - d1);

mlir/test/Dialect/Linalg/transform-lower-pack.mlir

215

I would just call it: 'pack_with_pad'.

[qedawkins]

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

551

Tangentially related, the documentation for the pack op says if the padding value is absent and the tile doesn't perfectly divide then it is UB. Do we need to update the op semantics or does this need to be a failure? (https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorpack-mlirtensorpackop)

[hanchung]

rebase + address comments

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

551

It will raise an error if the padding value is absent and the tile doesn't perfectly divide the dimension. See https://github.com/llvm/llvm-project/blob/afe78db7701d15f1d69b3fa50ee05fd42a6297cd/mlir/test/Dialect/Tensor/invalid.mlir#L542-L546

For dynamic cases, it is hard to raise the signal. I think the pattern model the UB to pad with zeros when do the lowering. The padding value is needed for pad ops, even if the pad op is a no op (i.e., all the values in lows and highs are zeros).

I think the pattern can decide how to handle UB itself, so we don't have to update the op semantics. The failures is raised if the op has enough static information. In this context, we don't need to update the doc. What do you think?

[qedawkins]

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

551

Ah I forgot that the verifier will handle the static case, and for the dynamic case UB makes sense.

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

FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
 
   // 4. Pad the source of packOp to a shape we can expand into stripMinedShape.
   SmallVector<OpFoldResult> lows(packOp.getSourceRank(),
                                  rewriter.getIndexAttr(0));
   SmallVector<OpFoldResult> highs(packOp.getSourceRank(),
                                   rewriter.getIndexAttr(0));
   for (auto [pos, innerSize] :
        llvm::zip_equal(packOp.getInnerDimsPos(), packOp.getMixedTiles())) {
+    int outerPos =
+        packedToStripMinedShapePerm[packingMetadata.outerPositions[pos]];
     OpFoldResult origSize = rewriter.createOrFold<tensor::DimOp>(
         loc, packOp.getSource(),
         rewriter.create<arith::ConstantIndexOp>(loc, pos));
-    AffineExpr s0, d0;
-    bindDims(rewriter.getContext(), d0);
+    OpFoldResult outerSize = rewriter.createOrFold<tensor::DimOp>(
+        loc, packOp.getDest(),
+        rewriter.create<arith::ConstantIndexOp>(loc, outerPos));
+    AffineExpr s0, d0, d1;
+    bindDims(rewriter.getContext(), d0, d1);
     bindSymbols(rewriter.getContext(), s0);
-    auto map = AffineMap::get(1, 1, d0.ceilDiv(s0) * s0 - d0);
-    highs[pos] = affine::makeComposedFoldedAffineApply(rewriter, loc, map,
-                                                       {origSize, innerSize});
+    auto map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/1, d0 * s0 - d1);
+    highs[pos] = affine::makeComposedFoldedAffineApply(
+        rewriter, loc, map, {outerSize, origSize, innerSize});
   }
   RankedTensorType collapsed = tensor::CollapseShapeOp::inferCollapsedType(
       RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
       packingMetadata.reassociations);
   Value paddingValue = packOp.getPaddingValue();
   if (!paddingValue) {
     paddingValue = rewriter.create<arith::ConstantOp>(
         loc, rewriter.getZeroAttr(getElementTypeOrSelf(collapsed)));
   }
   auto padOp =
       rewriter.create<tensor::PadOp>(loc, collapsed, packOp.getSource(), lows,
                                      highs, paddingValue, /*nofold=*/false);
 
   LLVM_DEBUG(
       DBGSNL(); DBGSNL(); llvm::interleaveComma(packingMetadata.insertPositions,
                                                 DBGS() << "insertPositions: ");

mlir/test/Dialect/Linalg/transform-lower-pack.mlir

^bb1(%module_op: !transform.any_op):
   %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
     : (!transform.any_op) -> !transform.op<"tensor.pack">
   transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
     -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
 }
 
 // -----
 
+// CHECK-LABEL: func.func @pack_with_pad(
+func.func @pack_with_pad(%src: tensor<4225x12xf32>, %dest: tensor<265x16x16x1xf32>)
+    -> tensor<265x16x16x1xf32> {
+  //      CHECK: tensor.pad {{.*}} low[0, 0]
+  //      CHECK:   : tensor<4225x12xf32> to tensor<4240x16xf32>
+  //      CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2, 3]]
+  // CHECK-SAME:   : tensor<4240x16xf32> into tensor<265x16x16x1xf32>
+  //      CHECK: linalg.transpose
+  // CHECK-SAME:   ins(%{{[a-zA-Z0-9]*}} : tensor<265x16x16x1xf32>)
+  // CHECK-SAME:   outs(%{{[a-zA-Z0-9]*}} : tensor<265x16x16x1xf32>)
+  // CHECK-SAME:   permutation = [0, 2, 1, 3]
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.pack %src
+    padding_value(%cst : f32)
+    inner_dims_pos = [0, 1]
+    inner_tiles = [16, 1] into %dest
+    : tensor<4225x12xf32> -> tensor<265x16x16x1xf32>
+  return %0 : tensor<265x16x16x1xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+  %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
+    : (!transform.any_op) -> !transform.op<"tensor.pack">
+  transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
+    -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+}
+
+// -----
+
 // CHECK-LABEL: func.func @pack_with_pad_and_outer_dims_perm(
 func.func @pack_with_pad_and_outer_dims_perm(%src: tensor<100x200x127x255xi32>,
                                              %dest: tensor<200x4x16x100x16x32xi32>)
     -> tensor<200x4x16x100x16x32xi32> {
   //      CHECK: tensor.pad {{.*}} low[0, 0, 0, 0]
   //      CHECK:   : tensor<100x200x127x255xi32> to tensor<100x200x128x256xi32>
   //      CHECK: tensor.expand_shape %{{.*}} [{{.*}}[0], [1], [2, 3], [4, 5]]
   // CHECK-SAME:   : tensor<100x200x128x256xi32> into tensor<100x200x4x32x16x16xi32>
^bb1(%module_op: !transform.any_op):
   %pack = transform.structured.match ops{["tensor.pack"]} in %module_op
     : (!transform.any_op) -> !transform.op<"tensor.pack">
   transform.structured.lower_pack %pack : (!transform.op<"tensor.pack">)
     -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
 }
 
 // -----
 
-// CHECK-DAG:   #[[MAP0:.+]] = affine_map<()[s0] -> (-s0 + (s0 ceildiv 16) * 16)>
-// CHECK-DAG:   #[[MAP1:.+]] = affine_map<()[s0] -> (-s0 + (s0 ceildiv 32) * 32)>
+// CHECK-DAG:   #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 * 16 - s1)>
+// CHECK-DAG:   #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 * 32 - s1)>
 // CHECK:       func.func @dynamic_pack_pad_transpose_inner_and_outer_dims(
 // CHECK-SAME:    %[[SRC:[a-zA-Z0-9]+]]
 func.func @dynamic_pack_pad_transpose_inner_and_outer_dims(%source: tensor<?x?xf32>) -> tensor<?x?x16x32xf32> {
   // CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
   // CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
   // CHECK-DAG:   %[[C16:.+]] = arith.constant 16 : index
   // CHECK-DAG:   %[[C32:.+]] = arith.constant 32 : index
   // CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[SRC]], %[[C0]]
   // CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[SRC]], %[[C1]]
   // CHECK-DAG:   %[[OUT_D0:.+]] = arith.ceildivui %[[D1]], %[[C16]] : index
   // CHECK-DAG:   %[[OUT_D1:.+]] = arith.ceildivui %[[D0]], %[[C32]] : index
   // CHECK-DAG:   %[[EMPTY:.+]] = tensor.empty(%[[OUT_D0]], %[[OUT_D1]]) : tensor<?x?x16x32xf32>
-  // CHECK-DAG:   %[[H1:.+]] = affine.apply #[[MAP0]]()[%[[D1]]]
-  // CHECK-DAG:   %[[H0:.+]] = affine.apply #[[MAP1]]()[%[[D0]]]
+  // CHECK-DAG:   %[[DEST_D0:.+]] = tensor.dim %[[EMPTY]], %[[C0]]
+  // CHECK-DAG:   %[[DEST_D1:.+]] = tensor.dim %[[EMPTY]], %[[C1]]
+  // CHECK-DAG:   %[[H1:.+]] = affine.apply #[[MAP0]]()[%[[DEST_D0]], %[[D1]]]
+  // CHECK-DAG:   %[[H0:.+]] = affine.apply #[[MAP1]]()[%[[DEST_D1]], %[[D0]]]
   // CHECK:       %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[%[[H0]], %[[H1]]]
   // CHECK:         : tensor<?x?xf32> to tensor<?x?xf32>
   // CHECK:       %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]] {{\[}}[0, 1], [2, 3]]
   // CHECK-SAME:   : tensor<?x?xf32> into tensor<?x32x?x16xf32>
   // CHECK:       %[[TRANSP:.+]] = linalg.transpose
   // CHECK-SAME:    ins(%[[EXPAND]] : tensor<?x32x?x16xf32>)
   // CHECK-SAME:    outs(%[[EMPTY]] : tensor<?x?x16x32xf32>)
   // CHECK-SAME:    permutation = [2, 0, 3, 1]