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

Incorporating suggestions from @hanhanW

Author: banach-space

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

@@ -1054,10 +1054,10 @@ static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,
   // Collect dims for the padded shape.
   SmallVector<int64_t> paddedShape;
   for (int64_t dimIdx = 0; dimIdx < inputRank; ++dimIdx) {
-    int64_t inputDimSize = inputType.getDimSize(dimIdx);
     // 1. Non-tiled outer dims.
     // These dims should be 1 and we simply preserve them.
     if (!tileAndPosMapping.count(dimIdx)) {
+      int64_t inputDimSize = inputType.getDimSize(dimIdx);
       assert(inputDimSize == 1 &&
              "with all outer dims == 1, this non-tiled input dim should be 1!");
       paddedShape.push_back(inputDimSize);
@@ -1214,16 +1214,6 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
 
   applyPermutationToVector<OpFoldResult>(transShapeForEmpty, perm);
 
-  // If there's a tile with a dynamic size, retrieve its size. ATM only 1
-  // dynamic tile is allowed.
-  Value dynDimSize;
-  for (auto tile : packOp.getMixedTiles()) {
-    if (tile.is<Value>()) {
-      assert(!dynDimSize && "Only one scalable size is supported ATM.");
-      dynDimSize = cast<Value>(tile);
-    }
-  }
-
   Value empty =
       rewriter.create<tensor::EmptyOp>(loc, transShapeForEmpty, elemType);
   auto transposedOp =

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

@@ -51,9 +51,8 @@ PadOp mlir::tensor::createPadHighOp(RankedTensorType resType, Value source,
     AffineExpr d0, d1;
     bindDims(b.getContext(), d0, d1);
     OpFoldResult sourceDim = tensor::getMixedSize(b, loc, source, idx);
-    Value outDim = isDimDynamic
-                       ? dynOutDims[outDimIdx++]
-                       : b.create<arith::ConstantIndexOp>(loc, val).getResult();
+    OpFoldResult outDim = isDimDynamic ? OpFoldResult(dynOutDims[outDimIdx++])
+                                       : OpFoldResult(b.getIndexAttr(val));
 
     high[idx] = affine::makeComposedFoldedAffineApply(b, loc, d0 - d1,
                                                       {outDim, sourceDim});

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

@@ -76,9 +76,9 @@ func.func @simple_pad_and_pack_scalable(%input: tensor<5x1xf32>, %output: tensor
 // CHECK-SAME:      %[[SRC:[a-zA-Z0-9]+]]: tensor<5x1xf32>,
 // CHECK-SAME:      %[[DEST:[a-zA-Z0-9]+]]: tensor<1x1x?x2xf32>,
 // CHECK-SAME:      %[[PAD_VAL:[a-zA-Z0-9]+]]: f32) -> tensor<1x1x?x2xf32> {
-// CHECK:           %[[C2:.+]] = arith.constant 2 : index
-// CHECK:           %[[C8:.+]] = arith.constant 8 : index
-// CHECK:           %[[VS:.+]] = vector.vscale
+// CHECK-DAG:       %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG:       %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG:       %[[VS:.+]] = vector.vscale
 // CHECK:           %[[C8_VS:.+]] = arith.muli %[[VS]], %[[C8]] : index
 // CHECK:           %[[PAD_HIGH:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[C8_VS]]]
 // CHECK:           %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {