fixup! fixup! [mlir][tensor] Generalize/restrict GeneralizeOuterUnitDimsPackOpPattern

Raname and move getSimplifiedDimSizePair

Author: banach-space

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

@@ -60,6 +60,25 @@ void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
                                 SmallVectorImpl<Value> &dynamicVec,
                                 SmallVectorImpl<int64_t> &staticVec);
 
+/// Given OpFoldResult representing dim size value (*), generates a pair of
+/// sizes:
+///   * 1st result, static value, contains an int64_t dim size that can be used
+///   to build ShapedType (ShapedType::kDynamic is used for truly dynamic dims),
+///   * 2nd result, dynamic value, contains OpFoldResult encapsulating the
+///   actual dim size (either original or updated input value).
+/// For input sizes for which it is possible to extract a constant Attribute,
+/// replaces the original size value with an integer attribute (unless it's
+/// already a constant Attribute). The 1st return value also becomes the actual
+/// integer size (as opposed ShapedType::kDynamic).
+///
+/// (*) This hook is usually used when, given input sizes as OpFoldResult,
+/// it's required to generate two vectors:
+///   * sizes as int64_t to generate a shape,
+///   * sizes as OpFoldResult for sizes-like attribute.
+/// Please update this comment if you identify other use cases.
+std::pair<int64_t, OpFoldResult>
+getSimplifiedOfrAndStaticSizePair(OpFoldResult ofr, Builder &b);
+
 /// Extract integer values from the assumed ArrayAttr of IntegerAttr.
 template <typename IntTy>
 SmallVector<IntTy> extractFromIntegerArrayAttr(Attribute attr) {

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

@@ -1139,37 +1139,14 @@ getPackUnpackRankReducedPerm(ArrayRef<int64_t> shape,
   return perm;
 }
 
-// A helper function to generate a dim-and-size pair for Ops like
-// ExtractSliceOp that require both:
-//  * dims to specify the output shape, and
-//  * sizes for the sizes attribute (or similar).
-// For dynamic sizes, if the corresponding size is a compile time constant:
-//  * the return size becomes the attribute encapsulating the known size, and
-//  * dim is updated from kDynamic to its actual known value.
-static std::pair<int64_t, OpFoldResult>
-getSimplifiedDimSizePair(OpFoldResult tileSizeOfr, Builder &b) {
-  int64_t tileSizeForShape =
-      getConstantIntValue(tileSizeOfr).value_or(ShapedType::kDynamic);
-
-  OpFoldResult tileSizeOfrSimplified;
-  if (tileSizeForShape != ShapedType::kDynamic) {
-    tileSizeOfrSimplified = b.getIndexAttr(tileSizeForShape);
-  } else {
-    tileSizeOfrSimplified = tileSizeOfr;
-  }
-
-  return std::pair<int64_t, OpFoldResult>(tileSizeForShape,
-                                          tileSizeOfrSimplified);
-}
-
 LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
     tensor::PackOp packOp, PatternRewriter &rewriter) const {
   // TODO: support the case that outer dimensions are not all 1s. A
   // tensor.expand_shape will be generated in this case.
-  if (llvm::any_of(packOp.getTiledOuterDims(),
+  if (llvm::any_of(packOp.getAllOuterDims(),
                    [](int64_t dim) { return dim != 1; })) {
     return rewriter.notifyMatchFailure(
-        packOp, "require the tiled outer dimensions of the result are all 1s");
+        packOp, "not all outer dimensions of the result are 1s");
   }
 
   Attribute zeroIdxAttr = rewriter.getIndexAttr(0);
@@ -1202,7 +1179,7 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
   for (auto i : llvm::seq<unsigned>(0, srcRank)) {
     if (dimAndTileMapping.count(i)) {
       auto [tileSize, tileSizeOfr] =
-          getSimplifiedDimSizePair(dimAndTileMapping[i], rewriter);
+          getSimplifiedOfrAndStaticSizePair(dimAndTileMapping[i], rewriter);
       extractSliceSizes.push_back(tileSizeOfr);
       outputShapeForExtractSlice.push_back(tileSize);
     }
@@ -1236,8 +1213,8 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
   }
 
   applyPermutationToVector<OpFoldResult>(transShapeForEmptyOpDynamic, perm);
-  Value empty =  rewriter.create<tensor::EmptyOp>(
-                          loc, transShapeForEmptyOpDynamic, elemType);
+  Value empty = rewriter.create<tensor::EmptyOp>(
+      loc, transShapeForEmptyOpDynamic, elemType);
 
   // 2.2 Create linalg.transpose
   auto transposedOp =
@@ -1254,7 +1231,7 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
 
   for (auto tileSize : packOp.getMixedTiles()) {
     auto [tileSizeStatic, tileSizeOfr] =
-        getSimplifiedDimSizePair(tileSize, rewriter);
+        getSimplifiedOfrAndStaticSizePair(tileSize, rewriter);
     writeSizes.push_back(tileSizeOfr);
     writeShape.push_back(tileSizeStatic);
   }

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

@@ -58,6 +58,20 @@ void dispatchIndexOpFoldResult(OpFoldResult ofr,
   staticVec.push_back(ShapedType::kDynamic);
 }
 
+std::pair<int64_t, OpFoldResult>
+getSimplifiedOfrAndStaticSizePair(OpFoldResult tileSizeOfr, Builder &b) {
+  int64_t tileSizeForShape =
+      getConstantIntValue(tileSizeOfr).value_or(ShapedType::kDynamic);
+
+  OpFoldResult tileSizeOfrSimplified =
+      (tileSizeForShape != ShapedType::kDynamic)
+          ? b.getIndexAttr(tileSizeForShape)
+          : tileSizeOfr;
+
+  return std::pair<int64_t, OpFoldResult>(tileSizeForShape,
+                                          tileSizeOfrSimplified);
+}
+
 void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
                                 SmallVectorImpl<Value> &dynamicVec,
                                 SmallVectorImpl<int64_t> &staticVec) {