[DRAFT] Generalize expand_shape to take shape as explicit input

[ramiro050]

Summary

*DO NOT SUBMIT*

(This patch is for early design feedback only. Notably, tests have not been
updated and the implementation is incomplete in some cases.)

This patch generalizes tensor.expand_shape and memref.expand_shape to consume
the output shape as a list of SSA values. This enables us to implement generic
reshape operations with dynamic shapes using collapse_shape/expand_shape pairs.

The output_shape input to expand_shape follows the static/dynamic representation
that's also used in tensor.extract_slice.

[nicolasvasilache]

Generally looks reasonable to me.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1483

drop this paragraph

1528

In practice, such APIs are quite painful to use.

Instead we use ArrayRef<OpFoldResult> (see e.g. tensor.insert_slice).
We also have SmallVector<OpFoldResult> getMixedSizes-like APIs.

Last, StaticValueUtils.cpp|h has helpers such as

SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                         ValueRange dynamicValues, Builder &b);

to do the mechanical parts for you

1609

I would add some more sugar, e.g.:

%b = memref.expand_shape %a reassociate[[0, 1], [2]] into [%sz0, %sz1, 32]

: memref<?x32xf32> into memref<?x?x32xf32>

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1021

nit: must result in ?

obtain does not work anymore in your new phrasing.

1049

I would add some more sugar, e.g.:

%b = tensor.expand_shape %a reassociate[[0, 1], [2]] into [%sz0, %sz1, 32]
        : tensor<?x32xf32> into tensor<?x?x32xf32>

1062

same comment re APIs

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

43

Make a static member of ExpandShapeOp and attach it to Tablegen, otherwise we easily forget what is available where and people end up reimplementing the same functionality in different places.

47

OpFoldResult-based API should be enough

49

delete this? I don't see why 2 versions

196

seems redundant for no good reason, just use std::is_same?

[mravishankar]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1520

Nit: Maybe just add a keyword output_shape (or something) after the reassociation list to better show what this is.

1528

+1.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1049

Same comment as mine above, this form looks good to me.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

311

I am guessing you wanted to get comments on whether you just fail on this path?

The additional use here you are mentioning, I am guessing are uses in tensor.dim ops. Those need to be handled separately through -resolve-shaped-type-result-dims. This pass is basically a pattern for DimOp, so if you have

%1 = tensor.collapse_shape %0 ...
%2 = tensor.dim %1, ...

It replaces the tensor.dim %1 in terms of tensor.dim %0. This is done through the ReifyRankedShapedTypeOpInterface that is implemented by the collapse/expand shape ops.

[sanjoy]

Thanks for the review! Replied to a few things inline as I continue to work on this.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1609

This bit is for collapse_shape -- I needed to change it because I moved out assemblyFormat from MemRef_ReassociativeReshapeOp. The behavior for collapse_shape should be unchanged, LMK if you see otherwise.

But I added the output_shape keyword to expand_shape as you suggested above.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1049

Added an output_shape keyword (like you suggested for memref.expand_shape) since we already use into when we print the types.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

196

Unless I'm missing some C++ trick, that requires to add a dependency from this header to either the memref or the tensor dialect. I assumed that was not OK, but please let me know if it is.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1049

would wfm but I don't see output_shape spelled out anywhere atm.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

196

ugh you're right, I always forget about the unfortunate duplication of ops in 2 dialects

[ramiro050]

After talking with @sanjoy, we decided that I will finish the implementation of this patch.

[ramiro050]

Updating D140821: [DRAFT] Generalize expand_shape to take shape as explicit input

Rebase, address comments, fix crashes when running integration tests.
All that remains is updating lit tests to use new assembly for
expand_shape.

[ramiro050]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

311

I've updated this part to handle the creation of the expand op. Should I be adding the handling for the collapse_shape + tensor.dim pattern, or is this something that should already happen during -resolve-shaped-type-result-dims?

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 
 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
 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 = [{
     SmallVector<AffineMap, 4> getReassociationMaps();
 
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
 
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
     };
 
     MemRefType getSrcType() { return getSrc().getType().cast<MemRefType>(); }
 
     MemRefType getResultType() { return getResult().getType().cast<MemRefType>(); }
 
     Value getViewSource() { return getSrc(); }
   }];
 
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
   let summary = "operation to produce a memref with a higher rank.";
   let description = [{
     The `memref.expand_shape` op produces a new view with a higher rank whose
     sizes are a reassociation of the original `view`. The operation is limited
     to such reassociations, where a dimension is expanded into one or multiple
     contiguous dimensions. Such reassociations never require additional allocs
     or copies.
 
     A reassociation is defined as a grouping of dimensions and is represented
     with an array of DenseI64ArrayAttr attributes.
 
     Example:
 
     ```mlir
-    %r = memref.expand_shape %0 [[0, 1], [2]]
-        : memref<?x?xf32> into memref<?x5x?xf32>
+    %r = memref.expand_shape %0 [[0, 1], [2]] [%sz0, %sz1, 32]
+        : memref<?x32xf32> into memref<?x?x32xf32>
     ```
 
     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
     unclear whether the first source dimension is divisible by 5), the op is
     accepted by the verifier. However, if the op is in fact invalid at runtime,
     the behavior is undefined.
 
     The source memref can be zero-ranked. In that case, the reassociation
     indices must be empty and the result shape may only consist of unit
     dimensions.
 
     For simplicity, this op may not be used to cast dynamicity of dimension
     sizes and/or strides. I.e., if and only if a source dimension is dynamic,
     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
+    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,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ValueRange":$output_shape, "ArrayRef<int64_t>":$static_output_shape),
     [{
-      build($_builder, $_state, resultType, src, attrs);
-      $_state.addAttribute("reassociation",
-                          getReassociationIndicesAttribute($_builder, reassociation));
+      build($_builder, $_state, resultType, src,
+            getReassociationIndicesAttribute($_builder, reassociation),
+            output_shape, static_output_shape);
     }]>,
 
     // Builder using ReassociationExprs.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ValueRange":$output_shape, "ArrayRef<int64_t>":$static_output_shape),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+          convertReassociationMapsToIndices(reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps,
+            output_shape, static_output_shape);
     }]>,
 
     // 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)>
   ];
 
let description = [{
     if at least one dimension in the corresponding reassociation group is
     dynamic. Similarly, the stride of a result dimension must be dynamic if and
     only if the corresponding start dimension in the source type is dynamic.
 
     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);
+
+  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`.
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
     // Builders for a reshape whose result type is passed explicitly.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       build($_builder, $_state, resultType, src, attrs);
       $_state.addAttribute("reassociation",
                           getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     /// Return `true` if this source MemRef type is guaranteed to be collapsible
     /// according to the given reassociation indices. In the presence of dynamic
     /// strides this is usually not the case.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
 class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     Tensor_Op<mnemonic, !listconcat(traits, [
       DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
       Pure])>,
-    Arguments<(ins AnyTensor:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
 
   code commonExtraClassDeclaration = [{
     static StringRef getReassociationAttrStrName() { return "reassociation"; }
     SmallVector<AffineMap, 4> getReassociationMaps();
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
     };
     RankedTensorType getSrcType() {
       return getSrc().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return getResult().getType().cast<RankedTensorType>();
     }
   }];
 
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
   let summary = "operation to produce a tensor with a higher rank";
   let description = [{
     The `tensor.expand_shape` op produces a new tensor with a higher
     rank whose sizes are a reassociation of the original `src`.
 
     A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of DenseI64ArrayAttr attribute.
-
-    The verification rule is that the reassociation maps are applied to the
-    result tensor with the higher rank to obtain the operand tensor with the
-    smaller rank.
+    represented with an array of DenseI64ArrayAttr attribute.  The reassociation
+    maps applied to the result tensor with the higher rank must obtain the
+    operand tensor with the smaller rank.
+
+    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 operand tensor type of a reshape can be zero-ranked if the result
     tensor type is statically shaped with all dimensions being unit extent. In
     such cases the reassociation map is empty.
 
     Examples:
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
-    %b = tensor.expand_shape %a [[0, 1], [2]]
-        : tensor<?x?xf32> into tensor<?x?x?xf32>
+    %b = tensor.expand_shape %a [[0, 1], [2]] [%sz0, %sz1, 32]
+        : tensor<?x32xf32> into tensor<?x?x32xf32>
     ```
   }];
+
+  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation,
+                       Variadic<Index>:$output_shape,
+                       DenseI64ArrayAttr:$static_output_shape);
+
+  let assemblyFormat = [{
+    $src $reassociation
+    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,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ValueRange":$output_shape, "ArrayRef<int64_t>":$static_output_shape),
     [{
-      build($_builder, $_state, resultType, src, attrs);
-      $_state.addAttribute("reassociation",
-          getReassociationIndicesAttribute($_builder, reassociation));
+      build($_builder, $_state, resultType, src,
+            getReassociationIndicesAttribute($_builder, reassociation),
+            output_shape, static_output_shape);
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ValueRange":$output_shape, "ArrayRef<int64_t>":$static_output_shape),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+          convertReassociationMapsToIndices(reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps,
+            output_shape, static_output_shape);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     int64_t getCorrespondingSourceDim(int64_t resultDim);
   }];
 
   let hasVerifier = 1;
 }
 
 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 with a smaller
     rank whose sizes are a reassociation of the original `src`.
 
     A reassociation is defined as a continuous grouping of dimensions and is
     represented with an array of DenseI64ArrayAttr attribute.
 
     The verification rule is that the reassociation maps are applied to the
     operand tensor with the higher rank to obtain the result tensor with the
     smaller rank.
 
     The result tensor type of a reshape can be zero-ranked if the operand
     tensor type is statically shaped with all dimensions being unit extent. In
     such case the reassociation map is empty.
 
     Examples:
 
     ```mlir
     // Dimension collapse (i, j) -> i' and k -> k'
     %b = tensor.collapse_shape %a [[0, 1], [2]]
         : 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`.
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
     // Builders for a reshape whose result type is passed explicitly.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       build($_builder, $_state, resultType, src, attrs);
       $_state.addAttribute("reassociation",
           getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration;
   let hasVerifier = 1;
 }
 

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 
 using ReassociationIndices = SmallVector<int64_t, 2>;
 using ReassociationIndicesRef = ArrayRef<int64_t>;
 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.
+//
+// `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.
+void inferExpandShapeOutputShape(RankedTensorType expandedType,
+                                 ArrayRef<ReassociationIndices> reassociation,
+                                 ArrayRef<OpFoldResult> inputShape,
+                                 SmallVectorImpl<Value> &outputShape,
+                                 SmallVector<int64_t> &staticOutputShape);
+
+void inferExpandShapeOutputShape(RankedTensorType expandedType,
+                                 ArrayRef<ReassociationExprs> reassociation,
+                                 ArrayRef<OpFoldResult> inputShape,
+                                 SmallVectorImpl<Value> &outputShape,
+                                 SmallVector<int64_t> &staticOutputShape);
+
 /// 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
 /// expanding dimensions.
 ///
 /// For example,
 ///   producerReassociation = [[0, 1], [2], [3, 4]]
 ///   consumerReassociation = [[0, 1], [2]]

 
 /// Wraps a list of reassociations in an ArrayAttr.
 ArrayAttr
 getReassociationIndicesAttribute(OpBuilder &b,
                                  ArrayRef<ReassociationIndices> reassociation);
 
 /// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
 SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
-    OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
+    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
 /// failed.
 Optional<SmallVector<ReassociationIndices>>
 getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType);
 
 /// Returns the reassociation maps to collapse `sourceShape` to `targetShape` if
return reshapeLikeShapesAreCompatible(
       [&](const Twine &msg) { return op->emitOpError(msg); },
       collapsedType.getShape(), expandedType.getShape(),
       op.getReassociationIndices(), isExpandingReshape);
 }
 
 /// 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>
+template <typename ReshapeOpTy, ReshapeOpKind opKind>
 struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
                                 PatternRewriter &rewriter) const override {
     auto srcReshapeOp =
         reshapeOp.getSrc().template getDefiningOp<ReshapeOpTy>();
     if (!srcReshapeOp)
       return failure();
 
     ShapedType resultType = reshapeOp.getResultType();
 
     if (hasNonIdentityLayout(srcReshapeOp.getSrc().getType()) ||
         hasNonIdentityLayout(reshapeOp.getSrc().getType()) ||
         hasNonIdentityLayout(reshapeOp.getResult().getType()))
       return failure();
 
     Optional<SmallVector<ReassociationIndices>> reassociationIndices =
         composeReassociationIndices(srcReshapeOp.getReassociationIndices(),
                                     reshapeOp.getReassociationIndices(),
                                     rewriter.getContext());
     if (!reassociationIndices)
       return failure();
+
+    if constexpr (opKind == ReshapeOpKind::kExpand)
+      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+          reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
+          reshapeOp.getOutputShape(), reshapeOp.getStaticOutputShape());
+    else
       rewriter.replaceOpWithNewOp<ReshapeOpTy>(
           reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
     return success();
   }
 };
 
 /// Pattern to compose
 /// `collapse_shape(expand_shape(%src, reassociation_1), reassociation_2)`.
 /// In that case both `srcType` and `resultType` can be expressed as a function
 /// of `intermediateType`.
for (const auto &lowerRankIndices : lowerRankReassociation) {
           break;
       }
       composedReassociation.push_back(composedIndices);
     }
     if (srcRank > resultRank) {
       rewriter.replaceOpWithNewOp<CollapseOpTy>(
           collapseOp, resultType, expandOp.getSrc(), composedReassociation);
     } else if (srcRank < resultRank) {
-      rewriter.replaceOpWithNewOp<ExpandOpTy>(
-          collapseOp, resultType, expandOp.getSrc(), composedReassociation);
+      abort();
+      // inferExpandShapeOutputShape won't work here since it will introduce
+      // additional uses for collapseOp.
+      //
+      // rewriter.replaceOpWithNewOp<ExpandOpTy>(
+      //     collapseOp, resultType, expandOp.getSrc(), composedReassociation);
     } else {
       // Collapses/expansions that do not change the rank are not allowed. Use
       // a cast instead.
       assert(llvm::equal(srcType.getShape(), resultType.getShape()) &&
              "expected same shape");
       rewriter.replaceOpWithNewOp<CastOpTy>(collapseOp, resultType,
                                             expandOp.getSrc());
     }
LogicalResult matchAndRewrite(ExpandOpTy expandOp,
     }
     auto composedReassociation =
         findCollapsingReassociation(resultReassociation, srcReassociation,
                                     resultType.getShape(), srcType.getShape());
     if (!composedReassociation)
       return failure();
 
     rewriter.replaceOpWithNewOp<ExpandOpTy>(
-        expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
+        expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
+        expandOp.getOutputShape(), expandOp.getStaticOutputShape());
     return success();
   }
 
 private:
   // Attempts to find a way to collapse `srcShape` to `resultShape` by
   // collapsing subshapes defined by the reassociation indices.
   Optional<SmallVector<ReassociationIndices>> findCollapsingReassociation(
       ArrayRef<ReassociationIndices> srcReassociation,

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

static LogicalResult reduceMatchAndRewriteHelper(Operation *op, uint64_t axis,
   }
 
   if (expandInputRank != 0) {
     int32_t expandedDim = axis < expandInputRank ? axis : expandInputRank - 1;
     reassociationMap[expandedDim].push_back(
         rewriter.getAffineDimExpr(expandedDim + 1));
   }
 
+  Value result = linalgOp.getResults()[0];
+  SmallVector<OpFoldResult> inputShape =
+      tensor::createDimValues(rewriter, loc, result);
+  SmallVector<Value> outputShape;
+  SmallVector<int64_t> staticOutputShape;
+  inferExpandShapeOutputShape(resultTy.cast<RankedTensorType>(),
+                              reassociationMap, inputShape, outputShape,
+                              staticOutputShape);
+
   rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-      op, resultTy, linalgOp.getResults()[0], reassociationMap);
+      op, resultTy, result, reassociationMap, outputShape, staticOutputShape);
   return success();
 }
 
 static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
                                   ArrayRef<int64_t> rhsShape,
                                   SmallVector<int64_t> &intermediateShape,
                                   bool isDynamic) {
   if (isDynamic) {
matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
                                intermediateShape, isDynamic) ||
         intermediateShape != operandTy.getShape()) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot expand into given shape");
     }
+
+    Value src = adaptor.getOperands()[0];
+    SmallVector<OpFoldResult> srcShape =
+        tensor::createDimValues(rewriter, src.getLoc(), src);
+    SmallVector<Value> outputShape;
+    SmallVector<int64_t> staticOutputShape;
+    inferExpandShapeOutputShape(resultTy.cast<RankedTensorType>(),
+                                reassociationMap, srcShape, outputShape,
+                                staticOutputShape);
+
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
+        reshape, resultTy, src, reassociationMap, outputShape,
+        staticOutputShape);
     return success();
   }
 };
 
 class ReshapeConverterCollapseExpand
     : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

struct FoldFillWithTensorReshape : OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     auto oldFill = reshapeOp.getSrc().template getDefiningOp<FillOp>();
     if (!oldFill)
       return failure();
 
     Location loc = oldFill.getLoc();
-    auto newInit = rewriter.create<TensorReshapeOp>(
+    TensorReshapeOp newInit;
+    if constexpr (std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value)
+
+      newInit = rewriter.create<TensorReshapeOp>(
           loc, reshapeOp.getResultType(), oldFill.output(),
+          reshapeOp.getReassociation(), reshapeOp.getOutputShape(),
+          reshapeOp.getStaticOutputShape());
+    else
+      newInit = rewriter.create<TensorReshapeOp>(loc, reshapeOp.getResultType(),
+                                                 oldFill.output(),
                                                  reshapeOp.getReassociation());
     rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, ValueRange{oldFill.value()},
                                         ValueRange{newInit});
 
     return success();
   }
 };
 
 /// Fold tensor.pad(linalg.fill) into linalg.fill if the padding value and the

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

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
if (rankReductionStrategy == RankReductionStrategy::ExtractInsertSlice) {
       SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
       return rewriter.createOrFold<tensor::InsertSliceOp>(
           loc, result, origOutput, offsets, sizes, strides);
     }
 
     assert(rankReductionStrategy ==
                RankReductionStrategy::ReassociativeReshape &&
            "unknown rank reduction strategy");
+    SmallVector<OpFoldResult> inputShape =
+        tensor::createDimValues(rewriter, loc, result);
+    SmallVector<Value> outputShape;
+    SmallVector<int64_t> staticOutputShape;
+    inferExpandShapeOutputShape(origResultType, reassociation, inputShape,
+                                outputShape, staticOutputShape);
     return rewriter.create<tensor::ExpandShapeOp>(loc, origResultType, result,
-                                                  reassociation);
+                                                  reassociation, outputShape,
+                                                  staticOutputShape);
   }
 
   // Collapse the given value.
   Value collapseValue(Value operand, ArrayRef<int64_t> targetShape,
                       ArrayRef<ReassociationIndices> reassociation,
                       Location loc, PatternRewriter &rewriter) const {
     if (auto memrefType = operand.getType().dyn_cast<MemRefType>()) {
       if (rankReductionStrategy == RankReductionStrategy::ExtractInsertSlice) {
auto rankReducedType =
         tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
             reassociation->size(), sliceOp.getSourceType(), offsets, sizes,
             strides)
             .cast<RankedTensorType>();
 
     Location loc = sliceOp.getLoc();
     Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
         loc, rankReducedType, sliceOp.getSource(), offsets, sizes, strides);
+
+    SmallVector<OpFoldResult> newSliceShape =
+        tensor::createDimValues(rewriter, loc, newSlice);
+    SmallVector<Value> outputShape;
+    SmallVector<int64_t> staticOutputShape;
+    inferExpandShapeOutputShape(resultType, *reassociation, newSliceShape,
+                                outputShape, staticOutputShape);
+
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        sliceOp, resultType, newSlice, *reassociation);
+        sliceOp, resultType, newSlice, *reassociation, outputShape,
+        staticOutputShape);
     return success();
   }
 };
 
 /// Convert `insert_slice` operations to rank-reduced versions.
 /// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.
 template <typename InsertOpTy>
 struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {

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

 
 #include "mlir/Dialect/Linalg/Passes.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include <utility>
 
if (auto opOperandType =
         if (failed(reshapeLikeShapesAreCompatible(
                 [&](const Twine &msg) {
                   return rewriter.notifyMatchFailure(genericOp, msg);
                 },
                 opOperandType.getShape(), expandedOperandType.getShape(),
                 reassociation,
                 /*isExpandingReshape=*/true)))
           return std::nullopt;
+        Location loc = genericOp.getLoc();
+
+        SmallVector<OpFoldResult> inputShape =
+            tensor::createDimValues(rewriter, loc, opOperand->get());
+        SmallVector<Value> outputShape;
+        SmallVector<int64_t> staticOutputShape;
+
+        inferExpandShapeOutputShape(expandedOperandType, reassociation,
+                                    inputShape, outputShape, staticOutputShape);
         expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
             genericOp.getLoc(), expandedOperandType, opOperand->get(),
-            reassociation));
+            reassociation, outputShape, staticOutputShape));
         continue;
       }
     }
     expandedOpOperands.push_back(opOperand->get());
   }
 
   Location loc = genericOp.getLoc();
   SmallVector<Value> outputs;
   for (OpOperand *opOperand : genericOp.getDpsInitOperands()) {
     AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
     auto opOperandType = opOperand->get().getType().cast<RankedTensorType>();
     RankedTensorType expandedOutputType =
         getExpandedType(opOperandType, indexingMap, expansionInfo);
     if (expandedOutputType != opOperand->get().getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(indexingMap, expansionInfo);
       if (failed(reshapeLikeShapesAreCompatible(
               [&](const Twine &msg) {
                 return rewriter.notifyMatchFailure(genericOp, msg);
               },
               opOperandType.getShape(), expandedOutputType.getShape(),
               reassociation,
               /*isExpandingReshape=*/true)))
         return std::nullopt;
+
+      Location loc = genericOp.getLoc();
+
+      Value operand = opOperand->get();
+      SmallVector<OpFoldResult> operandShape =
+          tensor::createDimValues(rewriter, loc, operand);
+      SmallVector<Value> outputShape;
+      SmallVector<int64_t> staticOutputShape;
+      inferExpandShapeOutputShape(expandedOutputType, reassociation,
+                                  operandShape, outputShape, staticOutputShape);
+
       outputs.push_back(rewriter.create<tensor::ExpandShapeOp>(
-          genericOp.getLoc(), expandedOutputType, opOperand->get(),
-          reassociation));
+          loc, expandedOutputType, operand, reassociation, outputShape,
+          staticOutputShape));
     } else {
       outputs.push_back(opOperand->get());
     }
   }
 
   // The iterator types of the expanded op are all parallel.
   SmallVector<utils::IteratorType> iteratorTypes(
       expansionInfo.getExpandedOpNumDims(), utils::IteratorType::parallel);
for (const auto &originalResult : llvm::enumerate(genericOp->getResults())) {
     auto originalResultType =
         originalResult.value().getType().cast<ShapedType>();
     auto collapsedOpResultType = collapsedOpResult.getType().cast<ShapedType>();
     if (collapsedOpResultType.getRank() != originalResultType.getRank()) {
       AffineMap indexingMap =
           genericOp.getIndexingMapMatchingResult(originalResult.value());
       SmallVector<ReassociationIndices> reassociation =
           getOperandReassociation(indexingMap, collapsingInfo);
+
+      SmallVector<OpFoldResult> collapsedOpShape =
+          tensor::createDimValues(rewriter, loc, collapsedOpResult);
+      SmallVector<Value> outputShape;
+      SmallVector<int64_t> staticOutputShape;
+
+      inferExpandShapeOutputShape(originalResultType.cast<RankedTensorType>(),
+                                  reassociation, collapsedOpShape, outputShape,
+                                  staticOutputShape);
+
       Value result = rewriter.create<tensor::ExpandShapeOp>(
-          loc, originalResultType, collapsedOpResult, reassociation);
+          loc, originalResultType, collapsedOpResult, reassociation,
+          outputShape, staticOutputShape);
       results.push_back(result);
     } else {
       results.push_back(collapsedOpResult);
     }
   }
   return results;
 }
 

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

 // canonicalizations of named ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Passes.h"
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_LINALGNAMEDOPCONVERSION
 #include "mlir/Dialect/Linalg/Passes.h.inc"
Operation *newConv =
           })
           .Default([](Operation *op) { return nullptr; });
   if (!newConv)
     return failure();
   for (auto attr : preservedAttrs)
     newConv->setAttr(attr.getName(), attr.getValue());
 
   // Expand dimensions back out to
+
+  Value newConvVal = newConv->getResult(0);
+  SmallVector<OpFoldResult> newConvShape =
+      tensor::createDimValues(rewriter, loc, newConvVal);
+  SmallVector<Value> outputShape;
+  SmallVector<int64_t> staticOutputShape;
+  inferExpandShapeOutputShape(resultTy, collapsedInitDims, newConvShape,
+                              outputShape, staticOutputShape);
   rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-      operation, resultTy, newConv->getResult(0), collapsedInitDims);
+      operation, resultTy, result, collapsedInitDims, outputShape,
+      staticOutputShape);
   return success();
 }
 
 namespace {
 struct SimplifyDepthwiseConvOp
     : public OpRewritePattern<DepthwiseConv2DNhwcHwcmOp> {
   using OpRewritePattern<DepthwiseConv2DNhwcHwcmOp>::OpRewritePattern;
 

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

for (OpOperand *operand : op.getDpsInputOperands()) {
     // If the shape is unchanged the input doesn't change.
     if (newShape == op.getShape(operand)) {
       newInputs.push_back(operand->get());
       continue;
     }
     Type newType = RankedTensorType::get(
         newShape,
         operand->get().getType().cast<RankedTensorType>().getElementType());
+
+    SmallVector<OpFoldResult> inputShape =
+        tensor::createDimValues(b, loc, operand->get());
+    SmallVector<Value> outputShape;
+    SmallVector<int64_t> staticOutputShape;
+    inferExpandShapeOutputShape(newType.cast<RankedTensorType>(), reassociation,
+                                inputShape, outputShape, staticOutputShape);
+
     Value newInput = b.create<tensor::ExpandShapeOp>(
-        loc, newType, operand->get(), reassociation);
+        loc, newType, operand->get(), reassociation, outputShape,
+        staticOutputShape);
     newInputs.push_back(newInput);
   }
 
   // Calculate the new output map and shape, we insert the new dimension based
   // on the index returned by `controlSplitReductionFn`.
   SmallVector<int64_t> newOutputShape;
   AffineMap oldOutputMap = op.getMatchingIndexingMap(op.getDpsInitOperand(0));
   ArrayRef<int64_t> oldShape = op.getShape(op.getDpsInitOperand(0));

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
                           ArrayRef<ReassociationIndices> reassociation) {
   // Only ranked memref source values are supported.
   auto srcType = src.getType().cast<MemRefType>();
   FailureOr<MemRefType> resultType =
       ExpandShapeOp::computeExpandedType(srcType, resultShape, reassociation);
   // Failure of this assertion usually indicates a problem with the source
   // type, e.g., could not get strides/offset.
   assert(succeeded(resultType) && "could not compute layout");
-  build(builder, result, *resultType, src, reassociation);
+  build(builder, result, *resultType, src, reassociation, ValueRange{},
+        resultShape);
 }
 
 LogicalResult ExpandShapeOp::verify() {
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
 
   if (srcType.getRank() >= resultType.getRank())
     return emitOpError("expected rank expansion, but found source rank ")
LogicalResult ExpandShapeOp::verify() {
   if (failed(expectedResultType))
     return emitOpError("invalid source layout map");
 
   // Check actual result type.
   if (*expectedResultType != resultType)
     return emitOpError("expected expanded type to be ")
            << *expectedResultType << " but found " << resultType;
 
+  if ((int64_t)getStaticOutputShape().size() != resultType.getRank())
+    return emitOpError("expected number of static shape bounds to be equal to "
+                       "the output rank (")
+           << resultType.getRank() << ") but found "
+           << getStaticOutputShape().size() << " inputs instead";
+
+  if ((int64_t)getOutputShape().size() !=
+      llvm::count(getStaticOutputShape(), ShapedType::kDynamic))
+    return emitOpError("mismatch in dynamic dims in output_shape and "
+                       "static_output_shape: static_output_shape has ")
+           << llvm::count(getStaticOutputShape(), ShapedType::kDynamic)
+           << " dynamic dims while output_shape has " << getOutputShape().size()
+           << " values";
+
   return success();
 }
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(
-      context);
+  results.add<
+      ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
+      ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(context);
 }
 
 /// Compute the layout map after collapsing a given source MemRef type with the
 /// specified reassociation indices.
 ///
 /// Note: All collapsed dims in a reassociation group must be contiguous. It is
 /// not possible to check this by inspecting a MemRefType in the general case.
 /// If non-contiguity cannot be checked statically, the collapse is assumed to
if (newResultType == op.getResultType()) {
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
-              CollapseShapeOpMemRefCastFolder>(context);
+  results.add<
+      ComposeReassociativeReshapeOps<CollapseShapeOp, ReshapeOpKind::kCollapse>,
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

if (encSrc) {
       op->setOperand(0, convert);
       return success();
     }
     if (encDst) {
       RankedTensorType rtp =
           op.getResult().getType().template cast<RankedTensorType>();
       auto denseTp =
           RankedTensorType::get(rtp.getShape(), rtp.getElementType());
-      auto reshape = rewriter.create<ReshapeOp>(loc, denseTp, op.getSrc(),
-                                                op.getReassociation());
+      ReshapeOp reshape;
+      if constexpr (std::is_same<ReshapeOp, tensor::ExpandShapeOp>::value) {
+        reshape = rewriter.create<ReshapeOp>(
+            loc, denseTp, op.getSrc(), op.getReassociation(),
+            op.getOutputShape(), op.getStaticOutputShape());
+      } else {
+        reshape = rewriter.create<ReshapeOp>(loc, denseTp, op.getSrc(),
+      }
       Value convert = rewriter.create<ConvertOp>(loc, rtp, reshape);
       rewriter.replaceOp(op, convert);
       return success();
     }
     return failure();
   }
 };
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 
 LogicalResult ExpandShapeOp::verify() {
   auto srcType = getSrcType();
   auto resultType = getResultType();
   if (srcType.getRank() >= resultType.getRank())
     return emitOpError("expected rank expansion, but found source rank ")
            << srcType.getRank() << " >= result rank " << resultType.getRank();
 
+  if ((int64_t)getStaticOutputShape().size() != resultType.getRank())
+    return emitOpError("expected number of static shape dims to be equal to "
+                       "the output rank (")
+           << resultType.getRank() << ") but found "
+           << getStaticOutputShape().size() << " inputs instead";
+
+  if ((int64_t)getOutputShape().size() !=
+      llvm::count(getStaticOutputShape(), ShapedType::kDynamic))
+    return emitOpError("mismatch in dynamic dims in output_shape and "
+                       "static_output_shape: static_output_shape has ")
+           << llvm::count(getStaticOutputShape(), ShapedType::kDynamic)
+           << " dynamic dims while output_shape has " << getOutputShape().size()
+           << " values";
+
   return verifyTensorReshapeOp(*this, getResultType(), getSrcType());
 }
 
 LogicalResult CollapseShapeOp::verify() {
   auto srcType = getSrcType();
   auto resultType = getResultType();
   if (srcType.getRank() <= resultType.getRank())
     return emitOpError("expected rank reduction, but found source rank ")
LogicalResult matchAndRewrite(DimOp dimOp,
     rewriter.replaceOpWithNewOp<AffineApplyOp>(dimOp, product, srcDimSizes);
     return success();
   }
 };
 } // namespace
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
-              FoldReshapeWithConstant<ExpandShapeOp>,
-              FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
-              FoldDimOfCollapseShape>(context);
+  results.add<
+      ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results
-      .add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-           ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
-           FoldReshapeWithConstant<CollapseShapeOp>,
-           FoldReshapeWithFromElements<CollapseShapeOp>, FoldCollapseOfCastOp>(
-          context);
+  results.add<
+      ComposeReassociativeReshapeOps<CollapseShapeOp, ReshapeOpKind::kCollapse>,
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);
 }

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 
 #include <numeric>
 
 using namespace mlir;
 
+void mlir::inferExpandShapeOutputShape(
+    RankedTensorType expandedType, ArrayRef<ReassociationIndices> reassociation,
+    ArrayRef<OpFoldResult> inputShape, SmallVectorImpl<Value> &outputShape,
+    SmallVector<int64_t> &staticOutputShape) {
+  outputShape.resize(expandedType.getRank());
+  staticOutputShape.resize(expandedType.getRank(), ShapedType::kDynamic);
+
+  for (const auto &it : llvm::enumerate(reassociation)) {
+    ReassociationIndices indexGroup = it.value();
+    int64_t inputIndex = it.index();
+    bool foundDynamic = false;
+    for (int64_t index : indexGroup) {
+      if (ShapedType::isDynamic(expandedType.getDimSize(index))) {
+        assert(!foundDynamic &&
+               "Cannot infer expanded shape with multiple dynamic dims in the "
+               "same reassociation group!");
+        foundDynamic = true;
+        // Call get<Value>() under the assumption that we're not casting
+        // dynamism.
+        outputShape[index] = inputShape[inputIndex].get<Value>();
+      }
+    }
+
+    for (int64_t index : indexGroup) {
+      int64_t outputDimSize = expandedType.getDimSize(index);
+      if (ShapedType::isDynamic(outputDimSize))
+        continue;
+      // This restriction can be lifted -- we just need to change the output
+      // shape computation above to insert an integer division operation.
+      assert((!foundDynamic || outputDimSize == 1) &&
+             "Non-unit dimensions not supported in reassociation groups with "
+             "dynamic dims!");
+      staticOutputShape[index] = outputDimSize;
+    }
+  }
+
+  assert(llvm::count(staticOutputShape, ShapedType::kDynamic) ==
+             (outputShape.size() - llvm::count(outputShape, Value{})) &&
+         "Missing output shape entries!");
+  llvm::erase_value(outputShape, Value{});
+}
+
+void mlir::inferExpandShapeOutputShape(
+    RankedTensorType expandedType, ArrayRef<ReassociationExprs> reassociation,
+    ArrayRef<OpFoldResult> inputShape, SmallVectorImpl<Value> &outputShape,
+    SmallVector<int64_t> &staticOutputShape) {
+  inferExpandShapeOutputShape(expandedType,
+                              convertReassociationMapsToIndices(reassociation),
+                              inputShape, outputShape, staticOutputShape);
+}
+
 Optional<SmallVector<ReassociationIndices>>
 mlir::getReassociationIndicesForReshape(ShapedType sourceType,
                                         ShapedType targetType) {
   if (sourceType.getRank() > targetType.getRank())
     return getReassociationIndicesForCollapse(sourceType.getShape(),
                                               targetType.getShape());
   if (sourceType.getRank() < targetType.getRank())
     return getReassociationIndicesForCollapse(targetType.getShape(),
SmallVector<Attribute, 4> reassociationAttr =
       llvm::to_vector<4>(llvm::map_range(
           reassociation, [&](const ReassociationIndices &indices) -> Attribute {
             return b.getI64ArrayAttr(indices).cast<Attribute>();
           }));
   return b.getArrayAttr(reassociationAttr);
 }
 
 SmallVector<ReassociationIndices, 2> mlir::convertReassociationMapsToIndices(
-    OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs) {
+    ArrayRef<ReassociationExprs> reassociationExprs) {
   SmallVector<ReassociationIndices, 2> reassociationIndices;
   for (const auto &exprs : reassociationExprs) {
     ReassociationIndices indices;
     indices.reserve(exprs.size());
     for (const auto &expr : exprs)
       indices.push_back(expr.cast<AffineDimExpr>().getPosition());
     reassociationIndices.push_back(indices);
   }

 }
 
 LogicalResult mlir::reshapeLikeShapesAreCompatible(
     function_ref<LogicalResult(const Twine &)> emitError,
     ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
     ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape) {
   unsigned expandedDimStart = 0;
   for (const auto &map : llvm::enumerate(reassociationMaps)) {
-    std::optional<int64_t> dynamicShape;
+    bool foundDynamicShape = false;
     int64_t linearizedStaticShape = 1;
     for (const auto &dim : llvm::enumerate(
              expandedShape.slice(expandedDimStart, map.value().size()))) {
-      if (ShapedType::isDynamic(dim.value())) {
-        if (isExpandingReshape && dynamicShape) {
-          return emitError("invalid to have a single dimension (" +
-                           Twine(map.index()) +
-                           ") expanded into multiple dynamic dims (" +
-                           Twine(expandedDimStart + dynamicShape.value()) +
-                           "," + Twine(expandedDimStart + dim.index()) + ")");
-        }
-        dynamicShape = dim.index();
-      } else {
+      if (ShapedType::isDynamic(dim.value()))
+        foundDynamicShape = true;
+      else
         linearizedStaticShape *= dim.value();
-      }
-    }
-    if (dynamicShape) {
+    if (foundDynamicShape) {
       if (!ShapedType::isDynamic(collapsedShape[map.index()])) {
         return emitError(
             "expected dimension " + Twine(map.index()) +
             " of collapsed type to be dynamic since one or more of the "
             "corresponding dimensions in the expanded type is dynamic");
       }
     } else {
       if (collapsedShape[map.index()] != linearizedStaticShape) {

mlir/test/Dialect/Tensor/ops.mlir

func.func @insert_slice(
   return
 }
 
 // -----
 
 func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>)
     -> (tensor<f32>, tensor<1x1xf32>) {
   %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>
-  %1 = tensor.expand_shape %0 [] : tensor<f32> into tensor<1x1xf32>
+  %1 = tensor.expand_shape %0 [] [1, 1] : tensor<f32> into tensor<1x1xf32>
   return %0, %1 : tensor<f32>, tensor<1x1xf32>
 }
 // CHECK-LABEL: func @tensor_reshape_zero_dim
 //       CHECK:   tensor.collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
-//       CHECK:   tensor.expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
+//       CHECK:   tensor.expand_shape %{{.*}} [] [1, 1] : tensor<f32> into tensor<1x1xf32>
+
+// -----
+
+func.func @tensor_expand_shape_dynamic_dim(%arg0 : tensor<?x?xf32>, %sz0 : index, %sz1 : index, %sz2 : index)
+    -> (tensor<5x?x?x?xf32>) {
+  %1 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] [5, %sz0, %sz1, %sz2] : tensor<?x?xf32> into tensor<5x?x?x?xf32>
+  return %1 : tensor<5x?x?x?xf32>
+}
+
+// CHECK-LABEL:  func.func @tensor_expand_shape_dynamic_dim(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<5x?x?x?xf32> {
+//       CHECK:    %expanded = tensor.expand_shape %arg0 {{\[\[}}0, 1], [2, 3{{\]\]}} [5, %arg1, %arg2, %arg3] : tensor<?x?xf32> into tensor<5x?x?x?xf32>
+//       CHECK:    return %expanded : tensor<5x?x?x?xf32>
+//       CHECK:  }
+
 
 // -----
 
 func.func @legal_collapsing_reshape_dynamic_tensor
   (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
 {
   %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
     tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>