[llvm] [mlir] Revert "[MLIR] Generalize expand_shape to take shape as explicit input" (PR #89540)

via llvm-commits llvm-commits at lists.llvm.org
Sun Apr 21 05:34:17 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-memref

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

Reverts llvm/llvm-project#<!-- -->69267

this broke some bots.

---

Patch is 259.47 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/89540.diff


52 Files Affected:

- (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+22-58) 
- (modified) mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td (+23-56) 
- (modified) mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h (+12-46) 
- (modified) mlir/include/mlir/Dialect/Utils/StaticValueUtils.h (+3-2) 
- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+1) 
- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+4-12) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp (+1-1) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp (+4-6) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp (+5-8) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+18-49) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp (-1) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+3-5) 
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+7-75) 
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp (+2-9) 
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+14-71) 
- (modified) mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp (-3) 
- (modified) mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp (+8-16) 
- (modified) mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp (+14-69) 
- (modified) mlir/lib/Dialect/Utils/StaticValueUtils.cpp (+4-3) 
- (modified) mlir/test/Conversion/MemRefToLLVM/expand-then-convert-to-llvm.mlir (+8-8) 
- (modified) mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir (+2-2) 
- (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir (+10-19) 
- (modified) mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir (+33-81) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir (+1-1) 
- (modified) mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir (+1-3) 
- (modified) mlir/test/Dialect/Linalg/collapse-dim.mlir (+3-3) 
- (modified) mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir (+10-10) 
- (modified) mlir/test/Dialect/Linalg/data-layout-propagation.mlir (+12-18) 
- (modified) mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir (+37-71) 
- (modified) mlir/test/Dialect/Linalg/flatten-elementwise.mlir (+1-1) 
- (modified) mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir (+44-57) 
- (modified) mlir/test/Dialect/Linalg/fusion-push-reshape.mlir (+10-14) 
- (modified) mlir/test/Dialect/Linalg/reshape_control_fusion.mlir (+1-1) 
- (modified) mlir/test/Dialect/Linalg/reshape_fusion.mlir (+100-192) 
- (modified) mlir/test/Dialect/Linalg/resolve-shaped-type-result-dims.mlir (+3-2) 
- (modified) mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir (+14-14) 
- (modified) mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir (+1-3) 
- (modified) mlir/test/Dialect/MemRef/canonicalize.mlir (+18-17) 
- (modified) mlir/test/Dialect/MemRef/expand-strided-metadata.mlir (+7-9) 
- (modified) mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir (+10-12) 
- (modified) mlir/test/Dialect/MemRef/invalid.mlir (+24-14) 
- (modified) mlir/test/Dialect/MemRef/ops.mlir (+32-40) 
- (modified) mlir/test/Dialect/MemRef/runtime-verification.mlir (+2-3) 
- (modified) mlir/test/Dialect/SparseTensor/sparse_reshape.mlir (+6-6) 
- (modified) mlir/test/Dialect/Tensor/bufferize.mlir (+10-14) 
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+59-53) 
- (modified) mlir/test/Dialect/Tensor/fold-empty-op.mlir (+3-2) 
- (modified) mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir (+2-4) 
- (modified) mlir/test/Dialect/Tensor/invalid.mlir (+16-5) 
- (modified) mlir/test/Dialect/Tensor/ops.mlir (+2-16) 
- (modified) mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir (+7-7) 
- (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (-1) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 14b8d95ea15b41..39e66cd9e6e5ab 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1548,6 +1548,7 @@ def MemRef_ReshapeOp: MemRef_Op<"reshape", [
 class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     MemRef_Op<mnemonic, !listconcat(traits,
       [Pure, ViewLikeOpInterface])>,
+    Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyStridedMemRef:$result)>{
 
   code commonExtraClassDeclaration = [{
@@ -1572,6 +1573,10 @@ class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     Value getViewSource() { return getSrc(); }
   }];
 
+  let assemblyFormat = [{
+    $src $reassociation attr-dict `:` type($src) `into` type($result)
+  }];
+
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
@@ -1593,10 +1598,14 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
     Example:
 
     ```mlir
-    %r = memref.expand_shape %0 [[0, 1], [2]] output_shape [%sz0, %sz1, 32]
-        : memref<?x32xf32> into memref<?x?x32xf32>
+    %r = memref.expand_shape %0 [[0, 1], [2]]
+        : memref<?x?xf32> into memref<?x5x?xf32>
     ```
 
+    At most one dimension of a reassociation group (e.g., [0, 1] above) may be
+    dynamic in the result type. Otherwise, the op would be ambiguous, as it
+    would not be clear how the source dimension is extended.
+
     If an op can be statically proven to be invalid (e.g, an expansion from
     `memref<10xf32>` to `memref<2x6xf32>`), it is rejected by the verifier. If
     it cannot statically be proven invalid (e.g., the full example above; it is
@@ -1613,80 +1622,41 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
     there must be a dynamic result dimension in the corresponding reassociation
     group. Same for strides.
 
-    The representation for the output shape supports a partially-static
-    specification via attributes specified through the `static_output_shape`
-    argument.  A special sentinel value `ShapedType::kDynamic` encodes that the
-    corresponding entry has a dynamic value.  There must be exactly as many SSA
-    inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
-    `static_output_shape`.
-
     Note: This op currently assumes that the inner strides are of the
     source/result layout map are the faster-varying ones.
   }];
 
-  let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation,
-                       Variadic<Index>:$output_shape,
-                       DenseI64ArrayAttr:$static_output_shape);
-
-  let assemblyFormat = [{
-    $src $reassociation `output_shape`
-    custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
-    type($src) `into` type($result)
-  }];
-
   let builders = [
     // Builders using ReassociationIndices.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
-      "ArrayRef<OpFoldResult>":$outputShape)>,
-
-    // It will infer output shape using inferOutputShape() method.
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationIndices>":$reassociation)>,
-
-    // Builder using ReassociationExprs.
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationExprs>":$reassociation),
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
-      auto reassociationIndices =
-          convertReassociationMapsToIndices(reassociation);
-      build($_builder, $_state, resultType, src, reassociationIndices);
+      build($_builder, $_state, resultType, src, attrs);
+      $_state.addAttribute("reassociation",
+                          getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
 
+    // Builder using ReassociationExprs.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      "ArrayRef<OpFoldResult>":$outputShape),
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices(reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps,
-            outputShape);
+          convertReassociationMapsToIndices($_builder, reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>,
 
-    // Builder that infers the result layout map. The result shape must be
-    // specified. Otherwise, the op may be ambiguous. The output shape for 
-    // the op will be inferred using the inferOutputShape() method.
-    OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
-               "ArrayRef<ReassociationIndices>":$reassociation)>,
-
     // Builder that infers the result layout map. The result shape must be
     // specified. Otherwise, the op may be ambiguous.
     OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
-               "ArrayRef<ReassociationIndices>":$reassociation,
-               "ArrayRef<OpFoldResult>":$outputShape)>
+               "ArrayRef<ReassociationIndices>":$reassociation)>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     static FailureOr<MemRefType> computeExpandedType(
         MemRefType srcType, ArrayRef<int64_t> resultShape,
         ArrayRef<ReassociationIndices> reassociation);
-
-    // Infer the output shape for a memref.expand_shape when it is possible
-    // to do so.
-    static FailureOr<SmallVector<OpFoldResult>> inferOutputShape(
-        OpBuilder &b, Location loc, MemRefType expandedType,
-        ArrayRef<ReassociationIndices> reassociation,
-        ArrayRef<OpFoldResult> inputShape);
   }];
 
   let hasVerifier = 1;
@@ -1737,12 +1707,6 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
     source/result layout map are the faster-varying ones.
   }];
 
-  let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation);
-
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.
@@ -1754,7 +1718,7 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices(reassociation);
+          convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
@@ -1772,7 +1736,7 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices(reassociation);
+          convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index a403e89a39f98c..cf7f3e89079c1c 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1062,7 +1062,8 @@ class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     Tensor_Op<mnemonic, !listconcat(traits, [
       DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
       Pure])>,
-    Results<(outs AnyTensor:$result)> {
+    Arguments<(ins AnyRankedTensor:$src, IndexListArrayAttr:$reassociation)>,
+    Results<(outs AnyRankedTensor:$result)> {
 
   code commonExtraClassDeclaration = [{
     static StringRef getReassociationAttrStrName() { return "reassociation"; }
@@ -1085,6 +1086,10 @@ class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     }
   }];
 
+  let assemblyFormat = [{
+    $src $reassociation attr-dict `:` type($src) `into` type($result)
+  }];
+
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
@@ -1097,75 +1102,43 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
     rank than the operand `src` whose dimension sizes are a reassociation of
     `src`.
 
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of DenseI64ArrayAttr attribute.  The reassociation
-    maps applied to the result tensor with the higher rank must result in the
-    operand tensor with the smaller rank.
+    A reassociation is defined as a continuous grouping of dimensions. It is
+    represented with an array of DenseI64ArrayAttr attribute. Entries in the
+    array are referred to as reassociation maps.
 
-    The representation for the output shape supports a partially-static
-    specification via attributes specified through the `static_output_shape`
-    argument.  A special sentinel value `ShapedType::kDynamic` encodes that the
-    corresponding entry has a dynamic value.  There must be exactly as many SSA
-    inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
-    `static_output_shape`.
+    The reassociation maps are applied to the result shape to obtain the operand
+    shape.
 
     Example:
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
-    %b = tensor.expand_shape %a [[0, 1], [2]] output_shape [%sz0, %sz1, 32]
-        : tensor<?x32xf32> into tensor<?x?x32xf32>
+    %b = tensor.expand_shape %a [[0, 1], [2]]
+        : tensor<?x?xf32> into tensor<?x?x?xf32>
     ```
   }];
-
-  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation,
-                       Variadic<Index>:$output_shape,
-                       DenseI64ArrayAttr:$static_output_shape);
-
-  let assemblyFormat = [{
-    $src $reassociation `output_shape`
-    custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
-    type($src) `into` type($result)
-  }];
-
   let builders = [
     // Builders using ReassociationIndices.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
-      "ArrayRef<OpFoldResult>":$outputShape)>,
-
-    // It will infer output shape using inferOutputShape() method.
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationIndices>":$reassociation)>,
-
-    // Builder using ReassociationExprs.
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationExprs>":$reassociation),
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
-      auto reassociationIndices =
-          convertReassociationMapsToIndices(reassociation);
-      build($_builder, $_state, resultType, src, reassociationIndices);
+      build($_builder, $_state, resultType, src, attrs);
+      $_state.addAttribute("reassociation",
+          getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      "ArrayRef<OpFoldResult>":$outputShape),
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
-      auto reassociationIndices =
-          convertReassociationMapsToIndices(reassociation);
-      build($_builder, $_state, resultType, src, reassociationIndices,
-            outputShape);
+      auto reassociationMaps =
+          convertReassociationMapsToIndices($_builder, reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     int64_t getCorrespondingSourceDim(int64_t resultDim);
-
-    // Infer the output shape for a tensor.expand_shape when it is possible
-    // to do so.
-    static FailureOr<SmallVector<OpFoldResult>> inferOutputShape(
-        OpBuilder &b, Location loc, RankedTensorType expandedType,
-        ArrayRef<ReassociationIndices> reassociation,
-        ArrayRef<OpFoldResult> inputShape);
   }];
 
   let hasVerifier = 1;
@@ -1173,7 +1146,6 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
 
 def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
   let summary = "operation to produce a tensor with a smaller rank";
-  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation);
   let description = [{
     The `tensor.collapse_shape` op produces a new tensor of lower (or equal)
     rank whose dimension sizes are a reassociation of the original `src` dimensions.
@@ -1191,11 +1163,6 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
         : tensor<?x?x?xf32> into tensor<?x?xf32>
     ```
   }];
-
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.
@@ -1207,7 +1174,7 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices(reassociation);
+          convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
@@ -1225,7 +1192,7 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices(reassociation);
+          convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 8a41a0a18b0ab3..ae9824f728da4d 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -30,27 +30,6 @@ using ReassociationExprs = SmallVector<AffineExpr, 2>;
 /// Attribute name for the ArrayAttr which encodes reassociation indices.
 constexpr StringRef getReassociationAttrName() { return "reassociation"; }
 
-// Infer the output shape for a {memref|tensor}.expand_shape when it is possible
-// to do so.
-//
-// Note: This should *only* be used to implement
-// `ExpandShapeOp::inferOutputShape` in both the memref and tensor namespaces.
-// If you need to infer the output shape you should use the static method of
-// `ExpandShapeOp` instead of calling this.
-//
-// `inputShape` is the shape of the tensor or memref being expanded as a
-// sequence of SSA values or constants. `expandedType` is the output shape of
-// the expand_shape operation. `reassociation` is the reassociation denoting
-// the output dims each input dim is mapped to.
-//
-// Returns the output shape in `outputShape` and `staticOutputShape`, following
-// the conventions for the output_shape and static_output_shape inputs to the
-// expand_shape ops.
-std::optional<SmallVector<OpFoldResult>>
-inferExpandShapeOutputShape(OpBuilder &b, Location loc, ShapedType expandedType,
-                            ArrayRef<ReassociationIndices> reassociation,
-                            ArrayRef<OpFoldResult> inputShape);
-
 /// Compose reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when
 /// both the producer and consumer are collapsing dimensions or both are
@@ -83,7 +62,7 @@ getReassociationIndicesAttribute(OpBuilder &b,
 
 /// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
 SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
-    ArrayRef<ReassociationExprs> reassociationExprs);
+    OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
 
 /// Return the reassociations maps to use to reshape given the source type and
 /// the target type when possible. Return std::nullopt when this computation
@@ -161,11 +140,14 @@ static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
       op.getReassociationIndices(), isExpansion);
 }
 
-/// Verify that shapes of the reshaped types using following rule:
-/// if a dimension in the collapsed type is static, then the corresponding
-/// dimensions in the expanded shape should be
+/// Verify that shapes of the reshaped types using following rules
+/// 1) if a dimension in the collapsed type is static, then the corresponding
+///    dimensions in the expanded shape should be
 ///    a) static
 ///    b) the product should be same as the collaped shape.
+/// 2) if a dimension in the collaped type is dynamic, one and only one of the
+///    corresponding dimensions in the expanded type should be dynamic. This
+///    rule is only needed with reshape operations that are expanding.
 LogicalResult reshapeLikeShapesAreCompatible(
     function_ref<LogicalResult(const Twine &)> emitError,
     ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
@@ -174,11 +156,9 @@ LogicalResult reshapeLikeShapesAreCompatible(
 /// Returns true iff the type is a MemRefType and has a non-identity layout.
 bool hasNonIdentityLayout(Type type);
 
-enum class ReshapeOpKind { kExpand, kCollapse };
-
 /// Pattern to collapse producer/consumer reshape ops that are both collapsing
 /// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy, ReshapeOpKind opKind>
+template <typename ReshapeOpTy>
 struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
@@ -201,18 +181,8 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
                                     rewriter.getContext());
     if (!reassociationIndices)
       return failure();
-
-    if constexpr (opKind == ReshapeOpKind::kExpand) {
-      SmallVector<OpFoldResult> outputShape(
-          getMixedValues(reshapeOp.getStaticOutputShape(),
-                         reshapeOp.getOutputShape(), rewriter));
-      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
-          outputShape);
-    } else {
-      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
-    }
+    rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+        reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
     return success();
   }
 };
@@ -245,8 +215,7 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
 //
 /// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
 /// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy,
-          typename DimOpTy, typename TensorTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
 struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
   using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
@@ -353,11 +322,8 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
     if (!composedReassociation)
       return failure();
 
-    SmallVector<OpFoldResult> outputShape(getMixedValues(
-        expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter));
     rewriter.replaceOpWithNewOp<ExpandOpTy>(
-        expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
-        outputShape);
+        expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
     return success();
   }
 
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index 594bcf5dbb399a..20f019666a2e6a 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -125,8 +125,9 @@ SmallVect...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/89540


More information about the llvm-commits mailing list