fixup! fixup! [mlir][tensor] Extend the logic to generalise tensor.pack

Allow multiple dynamic dims and remove the dependency on #109667

Author: banach-space

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

@@ -14,13 +14,22 @@
 namespace mlir {
 namespace tensor {
 
-// Return a PadOp that pads `source` to `type` size. Output sizes (from `type`)
-// are assumed to be static and greater than the potentially dynamic input sizes
-// (from `source`). The op performs "high" padding (i.e. it adds trailing
-// padding values until the desired size is met).
-PadOp createPadHighOp(RankedTensorType type, Value source, Value pad,
+// Return a PadOp that pads `source` to `resType` size. The op performs "high"
+// padding, i.e. it adds trailing padding values until the desired size is met.
+// Output sizes are assumed to be greater than the input sizes. The padding
+// width is calculated as: resDim - sourceDim.
+//
+// Handling static sizes is trivial. Dynamic dimensions are trickier (*):
+//  1. dynamic input sizes are extracted from `source`
+//  2. for dynamic output dims, there are two options:
+//    2.1 all output dynamic dim sizes are specified in `dynOutDim`,
+//    2.2 `dynOutDim` is empty and the corresponding padding width is set to 0.
+//
+// (*) Note that `resType` is just a shape and it only encodes the actual sizes
+// for _static_ dimensions.
+PadOp createPadHighOp(RankedTensorType resType, Value source, Value pad,
                       bool nofold, Location loc, OpBuilder &builder,
-                      std::optional<Value> dynOutDim = {});
+                      SmallVector<Value> dynOutDim = {});
 
 // Creates dim ops for each dynamic dimension of the ranked tensor argument and
 // returns these as values.

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

@@ -1048,8 +1048,8 @@ static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,
   DenseMap<int64_t, OpFoldResult> tileAndPosMapping =
       packOp.getDimAndTileMapping();
 
-  // The size of a dynamic tile (if present).
-  Value dynamicTileSize;
+  // The sizes of dynamic tiles
+  SmallVector<Value> dynamicTileSizes;
 
   // Collect dims for the padded shape.
   SmallVector<int64_t> paddedShape;
@@ -1080,14 +1080,13 @@ static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,
     paddedShape.push_back(ShapedType::kDynamic);
 
     // Get the value that holds the dynamic size.
-    assert(!dynamicTileSize && "Only one dynamic tile is supported ATM.");
-    dynamicTileSize = llvm::dyn_cast_if_present<Value>(tileSizeForDim);
+    dynamicTileSizes.push_back(llvm::dyn_cast<Value>(tileSizeForDim));
   }
   auto resultType =
       RankedTensorType::get(paddedShape, inputType.getElementType());
   return tensor::createPadHighOp(resultType, input, packOp.getPaddingValue(),
                                  /*nofold=*/false, loc, builder,
-                                 dynamicTileSize);
+                                 dynamicTileSizes);
 }
 
 // Normalizes a permutation on a higher rank space to its actual size, e.g.

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

@@ -22,36 +22,43 @@
 using namespace mlir;
 using namespace mlir::tensor;
 
-PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
+PadOp mlir::tensor::createPadHighOp(RankedTensorType resType, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b,
-                                    std::optional<Value> dynOutDim) {
+                                    SmallVector<Value> dynOutDims) {
 
-  assert(type.getNumDynamicDims() <= 1 &&
-         "At most one output dim can be dynamic!");
+  assert((resType.getNumDynamicDims() == dynOutDims.size()) ||
+         dynOutDims.empty() &&
+             "Either none or all output dynamic dims must be specified!");
 
   // Init "low" and "high" padding values ("low" is kept as is, "high" is
   // computed below).
-  SmallVector<OpFoldResult> low(type.getRank(), b.getIndexAttr(0));
-  SmallVector<OpFoldResult> high(type.getRank(), b.getIndexAttr(0));
+  SmallVector<OpFoldResult> low(resType.getRank(), b.getIndexAttr(0));
+  SmallVector<OpFoldResult> high(resType.getRank(), b.getIndexAttr(0));
 
-  for (const auto [idx, val] : enumerate(type.getShape())) {
-    bool isOutDimDynamic = ShapedType::isDynamic(val);
-    assert((!isOutDimDynamic || dynOutDim.has_value()) &&
-           "dynamic output dim requires dynOutDim to be set");
+  size_t outDimIdx = 0;
 
-    // Compute the padding width: outDim - srcDim.
+  for (const auto [idx, val] : enumerate(resType.getShape())) {
+    bool isDimDynamic = ShapedType::isDynamic(val);
+    bool updatePadHigh = !isDimDynamic || !dynOutDims.empty();
+
+    // Keep the default padding width (i.e. "0") when the output dim is dynamic
+    // and no actual output sizes have been provided.
+    if (!updatePadHigh)
+      continue;
+
+    // Compute the padding width: resDim - sourceDim.
     AffineExpr d0, d1;
     bindDims(b.getContext(), d0, d1);
-    OpFoldResult srcDim = tensor::getMixedSize(b, loc, source, idx);
-    Value outDim = isOutDimDynamic
-                       ? dynOutDim.value()
+    OpFoldResult sourceDim = tensor::getMixedSize(b, loc, source, idx);
+    Value outDim = isDimDynamic
+                       ? dynOutDims[outDimIdx++]
                        : b.create<arith::ConstantIndexOp>(loc, val).getResult();
 
     high[idx] = affine::makeComposedFoldedAffineApply(b, loc, d0 - d1,
-                                                      {outDim, srcDim});
+                                                      {outDim, sourceDim});
   }
-  return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
+  return b.create<PadOp>(loc, resType, source, low, high, pad, nofold);
 }
 
 SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,