[mlir][MemRef] Simplify extract_strided_metadata(expand_shape)

[qcolombet]

Summary

Add a pattern to the pass that simplifies extract_strided_metadata(other_op(memref)).

The new pattern gets rid of the expand_shape operation while materializing its effects on the sizes, and the strides of the base object.

In other words, this simplification replaces:

baseBuffer, offset, sizes, strides =
    extract_strided_metadata(expand_shape(memref))

With

baseBuffer, baseOffset, baseSizes, baseStrides =
    extract_strided_metadata(memref)
sizes#reassIdx =
  baseSizes#reassDim / product(expandShapeSizes#j,
                          for j in group excluding
                                 reassIdx)
strides#reassIdx =
baseStrides#reassDim * product(expandShapeSizes#j,
                           for j in
                                 reassIdx+1..
                                 reassIdx+group.size-1)

Where reassIdx is a reassociation index for the group at reassDim and expandShapeSizes#j is either:

• The constant size at dimension j, derived directly from the result type of the expand_shape op, or
• An affine expression: baseSizes#reassDim / product of all constant sizes in expandShapeSizes.

[chelini]

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

316

369

repeated the

[chelini]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

211

231

I would mention explicitly that this pattern does not affect base buffer and offset, as mentioned in the test cases.

[qcolombet]

Thanks for the quick review @chelini !

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

231

Good point!
It is not that clear in the high level comment of that pattern/commit message too since I changed the name of the offset value (from offset to baseOffset).
I'll fix that and also call it out.

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

316

Good catch x)

398

BTW something I haven't done yet is making sure this gets simplified into baseSizes#1 / 10 since at this point we know that baseSizes#1 is divisible by both 10 and 3.

[qcolombet]

• Fix typos
• Add a few comments
• Make sure the constant propagation happens for dynSize.floorDiv(sizes) * stride happens, since we know that dynSize is a multiple of sizes

[qcolombet]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

211

I'm actually with you @chelini on that one, but following the LLVM coding standard, I think this is one of the mandatory auto. (I.e., clang-tidy, I think, would replace the type by auto here.)

More specifically this sentence "[...]Don’t “almost always” use auto, but do use auto with initializers like cast<Foo>(...)[...]" in the coding guide.
https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

I'm happy to change it though.
Just let me know.

[qcolombet]

• Add the handling of 0-D
• Add a test case for 0-D

[qcolombet]

Fix typo in comment

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

194

nit: affine.apply

268

Can anything be refactored and reused from MemRefOps.cpp::namespace saturated_arith { ?

[qcolombet]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

268

Thanks for the pointer, I didn't know about these utilities.

At first glance, it doesn't look like it would fit here.

In particular, here we don't saturate when we see a dynamic stride/size, we replace it by 1 and patch up the expression with the proper dynamic value later.

The suffix product that we do here is similar to what computeExpandedLayoutMap does, except in our case we don't want to saturate.

I'm guessing the suffix product is used in a lot of places given that's how strides are computed, maybe we can refactor on that front. Though how we deal with dynamic dimensions is where things diverge and I don't know how to reconcile that at the moment.

Put differently, I don't know yet what the refactoring should look like.

Do you want me to spend more time looking into this or should we move on for now and wait for more use cases to show up?

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

268

I was hoping dome refactoring into simple utils would be possible because the code below feels like it's been rewritten a bunch of times in different ways.

For instance, I find the logic of the hasDynamicSize + productOfKnownStaticSize to be quite tricky to follow with further uses of hasDynamicSize.
I would hope that we can express this much more like functional-style applications without all these loops and levels of nesting.

For instance, l269 - 300 should look like:

SmallVector<int64_t> suffixProductOfSizes = 
  suffixProduct(expandShapeType.getShape(), begin, end, [](int64_t val, int64_t acc){ return saturated_arith(...); })

I made similar usability comments in https://reviews.llvm.org/D128986 but not exactly the same.

TL;DR I'd really prefer we avoid complex loopy logic for things that can be functional-style and should be reusable.

360

I'm sorry this is really hard to follow and I doubt this code will be maintainable..

[qcolombet]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

268

Okay, let me take a closer look.

Again, unless I missed something, the saturated_arith constructs are not good in that case because IIUC it saturates on the first dynamic value, whereas I need to get the full product without the dynamic dimension.

E.g., for 2x?x3, I need 6, not saturated, because this number will let us compute the actual value of ? from the original dimension for that group.

I feel that saturated_arith is well suited to construct the type of a shape (e.g., a stride is dynamic as soon as one of the dimension is dynamic), but not so much to compute actual values.

360

What do you think would make the code more understandable / maintainable?

Essentially for every reassociation group we have at most one dynamic size.
What this code does is the computation of this dynamic size as: origSize / product(all other sizes in that group).

E.g., for ?x4 -> [0,1,2][3] 2x?x3x4
This will compute the ? of the resulting shape as an affine.apply of ()[s0] -> (s0 / (2 * 3)), where s0 is the original dynamic dimension for that group.

The productOfKnownStaticSize holds the product of the static sizes for the related group (2 * 3 in the example), then dynSizeInput hold the original size for that group.

Maybe I could split the loop between the dimensions that come before the dynamic dimension and the ones that come after. That way I would not have to "propagate" a fake dynamic size expr to keep the final strideExpr stable between the dynamic and non-dynamic case, which should help readability.

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

268

Yes, I see that the saturated_arith itself is no use when you want to construct the AffineExpr, I just wanted to give you a feel of an API that avoids deep nesting and that will be easier to maintain longer term.

360

Essentially for every reassociation group we have at most one dynamic size.

No need to worry about this immediately but this will change (e.g. https://github.com/llvm/llvm-project/issues/51564)

Yes, less deeply nested fusion would def. be a good start, then we can iterate.

[nicolasvasilache]

Surfacing a general comment that appeared ofline:

I think we can always consider the general case with all ? and let the folders do their job on OpFoldResult:

sizes<?>, strides<?> expand to sizes<?x..x?>, strides<?x..x?> by expansionFactors(?x...x?)

If we reduce to this, this is all about determining SmallVector<OpFoldResult> expansionFactors(...).
In the most general case (NYI), these will be SSA values in the op (see e.g. https://github.com/llvm/llvm-project/issues/51564).

But even in the absence of the general case, I believe we greatly simplify the problem by thinking in those terms.

I'd add a helper SmallVector<OpFoldResult> ExpandShapeOp:::buildExpansionFactors(b, loc, groupId) to the op itself that constructs expansionFactors from the source memref and the result type (actually may need to be in some utils because of Affine cyclic dependencies .. sigh).

We'd likely want a counterpart for CollapseShape when we get there.

I think these suggestions will make the impl. both future proof and readable.

Still I have not tried so I may be missing something.

[qcolombet]

• Rework the code to make it more approachable:
  • Compute all the static sizes and strides, first
  • Adjust the strides impacted by the dynamic size, second

Next: Will look into the buildExpansionFactors idea.
I wanted to share this version since it was what was driving our offline conversation with Nicolas.

[qcolombet]

• Move the sizes and strides computations in their own helper functions

The buildExpansionFactors didn't pan out as well as I would have hoped because the expressions are too different when you factor dynamic size and/or strides in the mix.

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

170

nit: typo

177

I would make this a first-class citizen of expand: SmallVector<OpFoldResult> memref::ExpandShapeOp::buildExpandedSizes(...).

Then as the op semantics evolve to take the result shape operands, we can update this.

178

ArrayRef ?

226

same for making this a first-class citizen of expand

227

ArrayRef ?

[qcolombet]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

177

That's not possible because the MemRef dialect would depend on Affine dialect if we were to do that.

[qcolombet]

• Use ArrayRef instead of SmallVectorImpl
• Fix typo

[qcolombet]

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

177

For the record, my plan was to move that into a common place (e.g., ReshapeOpsUtils), when we think it looks good.

178

Good point!

[nicolasvasilache]

Thanks for splitting it up this way, let's go with this for now

I will try to simplify some of the impl (at a tradeoff of adding a bit more IR that folds) as a separate PR but this is a great step towards removing logic from LLVM while also expanding the cases we can support.

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

177

Ah yes, I keep on forgetting this unfortunate layering .. could you please add a TODO?

[chelini]

Please go ahead with Nicolas. He has better ideas on how to bring this up.

[qcolombet]

• Add a TODO for moving the utility functions to ExpandShapeOp

[qcolombet]

Thanks again @chelini and @nicolasvasilache for the reviews!

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

 /// baseBuffer, baseOffset, baseSizes, baseStrides =
 ///     extract_strided_metadata(memref)
 /// strides#i = baseStrides#i * subSizes#i
 /// offset = baseOffset + sum(subOffset#i * strides#i)
 /// sizes = subSizes
 /// \endverbatim
 ///
 /// In other words, get rid of the subview in that expression and canonicalize
-/// on its effects on the offset, the sizes, and the strides using affine apply.
+/// on its effects on the offset, the sizes, and the strides using affine.apply.
 struct ExtractStridedMetadataOpSubviewFolder
     : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
 public:
   using OpRewritePattern<memref::ExtractStridedMetadataOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
                                 PatternRewriter &rewriter) const override {
     auto subview = op.getSource().getDefiningOp<memref::SubViewOp>();
LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
     }
     assert(insertedDims == subRank &&
            "Should have populated all the values at this point");
 
     rewriter.replaceOp(op, results);
     return success();
   }
 };
+
+/// Compute the expanded sizes of the given \p expandShape for the
+/// \p groupId-th reassociation group.
+/// \p origSizes hold the sizes of the source shape as values.
+/// This is used to compute the new sizes in cases of dynamic shapes.
+///
+/// \post result.size() == expandShape.getReassociationIndices()[groupId].size()
+///
+/// TODO: Move this utility function directly within ExpandShapeOp. For now,
+/// this is not possible because this function uses the Affine dialect and the
+/// MemRef dialect cannot depend on the Affine dialect.
+static SmallVector<OpFoldResult>
+getExpandedSizes(memref::ExpandShapeOp expandShape, OpBuilder &builder,
+                 ArrayRef<OpFoldResult> origSizes, unsigned groupId) {
+  SmallVector<int64_t, 2> reassocGroup =
+      expandShape.getReassociationIndices()[groupId];
+  assert(!reassocGroup.empty() &&
+         "Reassociation group should have at least one dimension");
+
+  unsigned groupSize = reassocGroup.size();
+  SmallVector<OpFoldResult> expandedSizes(groupSize);
+
+  uint64_t productOfAllStaticSizes = 1;
+  Optional<unsigned> dynSizeIdx;
+  MemRefType expandShapeType = expandShape.getResultType();
+
+  // Fill up all the statically known sizes.
+  for (unsigned i = 0; i < groupSize; ++i) {
+    uint64_t dimSize = expandShapeType.getDimSize(reassocGroup[i]);
+    if (ShapedType::isDynamic(dimSize)) {
+      assert(!dynSizeIdx && "There must be at most one dynamic size per group");
+      dynSizeIdx = i;
+      continue;
+    }
+    productOfAllStaticSizes *= dimSize;
+    expandedSizes[i] = builder.getIndexAttr(dimSize);
+  }
+
+  // Compute the dynamic size using the original size and all the other known
+  // static sizes:
+  // expandSize = origSize / productOfAllStaticSizes.
+  if (dynSizeIdx) {
+    AffineExpr s0 = builder.getAffineSymbolExpr(0);
+    expandedSizes[*dynSizeIdx] = makeComposedFoldedAffineApply(
+        builder, expandShape.getLoc(), s0.floorDiv(productOfAllStaticSizes),
+        origSizes[groupId]);
+  }
+
+  return expandedSizes;
+}
+
+/// Compute the expanded strides of the given \p expandShape for the
+/// \p groupId-th reassociation group.
+/// \p origStrides and \p origSizes hold respectively the strides and sizes
+/// of the source shape as values.
+/// This is used to compute the strides in cases of dynamic shapes and/or
+/// dynamic stride for this reassociation group.
+///
+/// \post result.size() == expandShape.getReassociationIndices()[groupId].size()
+///
+/// TODO: Move this utility function directly within ExpandShapeOp. For now,
+/// this is not possible because this function uses the Affine dialect and the
+/// MemRef dialect cannot depend on the Affine dialect.
+SmallVector<OpFoldResult> getExpandedStrides(memref::ExpandShapeOp expandShape,
+                                             OpBuilder &builder,
+                                             ArrayRef<OpFoldResult> origSizes,
+                                             ArrayRef<OpFoldResult> origStrides,
+                                             unsigned groupId) {
+  SmallVector<int64_t, 2> reassocGroup =
+      expandShape.getReassociationIndices()[groupId];
+  assert(!reassocGroup.empty() &&
+         "Reassociation group should have at least one dimension");
+
+  unsigned groupSize = reassocGroup.size();
+  MemRefType expandShapeType = expandShape.getResultType();
+
+  Optional<int64_t> dynSizeIdx;
+
+  // Fill up the expanded strides, with the information we can deduce from the
+  // resulting shape.
+  uint64_t currentStride = 1;
+  SmallVector<OpFoldResult> expandedStrides(groupSize);
+  for (int i = groupSize - 1; i >= 0; --i) {
+    expandedStrides[i] = builder.getIndexAttr(currentStride);
+    uint64_t dimSize = expandShapeType.getDimSize(reassocGroup[i]);
+    if (ShapedType::isDynamic(dimSize)) {
+      assert(!dynSizeIdx && "There must be at most one dynamic size per group");
+      dynSizeIdx = i;
+      continue;
+    }
+
+    currentStride *= dimSize;
+  }
+
+  // Collect the statically known information about the original stride.
+  Value source = expandShape.getSrc();
+  auto sourceType = source.getType().cast<MemRefType>();
+  SmallVector<int64_t> strides;
+  int64_t offset;
+  bool hasKnownStridesAndOffset =
+      succeeded(getStridesAndOffset(sourceType, strides, offset));
+  (void)hasKnownStridesAndOffset;
+  assert(hasKnownStridesAndOffset &&
+         "getStridesAndOffset must work on valid expand_shape");
+
+  OpFoldResult origStride =
+      ShapedType::isDynamicStrideOrOffset(strides[groupId])
+          ? origStrides[groupId]
+          : builder.getIndexAttr(strides[groupId]);
+
+  // Apply the original stride to all the strides.
+  int64_t doneStrideIdx = 0;
+  // If we saw a dynamic dimension, we need to fix-up all the strides up to
+  // that dimension with the dynamic size.
+  if (dynSizeIdx) {
+    int64_t productOfAllStaticSizes = currentStride;
+    assert(ShapedType::isDynamic(sourceType.getDimSize(groupId)) &&
+           "We shouldn't be able to change dynamicity");
+    OpFoldResult origSize = origSizes[groupId];
+
+    AffineExpr s0 = builder.getAffineSymbolExpr(0);
+    AffineExpr s1 = builder.getAffineSymbolExpr(1);
+    for (; doneStrideIdx < *dynSizeIdx; ++doneStrideIdx) {
+      int64_t baseExpandedStride = expandedStrides[doneStrideIdx]
+                                       .get<Attribute>()
+                                       .cast<IntegerAttr>()
+                                       .getInt();
+      expandedStrides[doneStrideIdx] = makeComposedFoldedAffineApply(
+          builder, expandShape.getLoc(),
+          (s0 * baseExpandedStride).floorDiv(productOfAllStaticSizes) * s1,
+          {origSize, origStride});
+    }
+  }
+
+  // Now apply the origStride to the remaining dimensions.
+  AffineExpr s0 = builder.getAffineSymbolExpr(0);
+  for (; doneStrideIdx < groupSize; ++doneStrideIdx) {
+    int64_t baseExpandedStride = expandedStrides[doneStrideIdx]
+                                     .get<Attribute>()
+                                     .cast<IntegerAttr>()
+                                     .getInt();
+    expandedStrides[doneStrideIdx] = makeComposedFoldedAffineApply(
+        builder, expandShape.getLoc(), s0 * baseExpandedStride, {origStride});
+  }
+
+  return expandedStrides;
+}
+
+/// Replace `baseBuffer, offset, sizes, strides =
+///              extract_strided_metadata(expand_shape(memref))`
+/// With
+///
+/// \verbatim
+/// baseBuffer, offset, baseSizes, baseStrides =
+///     extract_strided_metadata(memref)
+/// sizes#reassIdx =
+///     baseSizes#reassDim / product(expandShapeSizes#j,
+///                                  for j in group excluding reassIdx)
+/// strides#reassIdx =
+///     baseStrides#reassDim * product(expandShapeSizes#j, for j in
+///                                    reassIdx+1..reassIdx+group.size-1)
+/// \endverbatim
+///
+/// Where reassIdx is a reassociation index for the group at reassDim
+/// and expandShapeSizes#j is either:
+/// - The constant size at dimension j, derived directly from the result type of
+///   the expand_shape op, or
+/// - An affine expression: baseSizes#reassDim / product of all constant sizes
+///   in expandShapeSizes. (Remember expandShapeSizes has at most one dynamic
+///   element.)
+///
+/// Notice that `baseBuffer` and `offset` are unchanged.
+///
+/// In other words, get rid of the expand_shape in that expression and
+/// materialize its effects on the sizes and the strides using affine apply.
+struct ExtractStridedMetadataOpExpandShapeFolder
+    : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
+public:
+  using OpRewritePattern<memref::ExtractStridedMetadataOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
+                                PatternRewriter &rewriter) const override {
+    auto expandShape = op.getSource().getDefiningOp<memref::ExpandShapeOp>();
+    if (!expandShape)
+      return failure();
+
+    // Build a plain extract_strided_metadata(memref) from
+    // extract_strided_metadata(expand_shape(memref)).
+    Location origLoc = op.getLoc();
+    IndexType indexType = rewriter.getIndexType();
+    Value source = expandShape.getSrc();
+    auto sourceType = source.getType().cast<MemRefType>();
+    unsigned sourceRank = sourceType.getRank();
+    SmallVector<Type> sizeStrideTypes(sourceRank, indexType);
+
+    auto newExtractStridedMetadata =
+        rewriter.create<memref::ExtractStridedMetadataOp>(
+            origLoc, op.getBaseBuffer().getType(), indexType, sizeStrideTypes,
+            sizeStrideTypes, source);
+
+    // Collect statically known information.
+    SmallVector<int64_t> strides;
+    int64_t offset;
+    bool hasKnownStridesAndOffset =
+        succeeded(getStridesAndOffset(sourceType, strides, offset));
+    (void)hasKnownStridesAndOffset;
+    assert(hasKnownStridesAndOffset &&
+           "getStridesAndOffset must work on valid expand_shape");
+    MemRefType expandShapeType = expandShape.getResultType();
+    unsigned expandShapeRank = expandShapeType.getRank();
+
+    // The result value will start with the base_buffer and offset.
+    unsigned baseIdxInResult = 2;
+    SmallVector<OpFoldResult> results(baseIdxInResult + expandShapeRank * 2);
+    results[0] = newExtractStridedMetadata.getBaseBuffer();
+    results[1] = ShapedType::isDynamicStrideOrOffset(offset)
+                     ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
+                     : rewriter.getIndexAttr(offset);
+
+    // Get the special case of 0-D out of the way.
+    if (sourceRank == 0) {
+      Value constantOne = getValueOrCreateConstantIndexOp(
+          rewriter, origLoc, rewriter.getIndexAttr(1));
+      SmallVector<Value> resultValues(baseIdxInResult + expandShapeRank * 2,
+                                      constantOne);
+      for (unsigned i = 0; i < baseIdxInResult; ++i)
+        resultValues[i] =
+            getValueOrCreateConstantIndexOp(rewriter, origLoc, results[i]);
+      rewriter.replaceOp(op, resultValues);
+      return success();
+    }
+
+    // Compute the expanded strides and sizes from the base strides and sizes.
+    SmallVector<OpFoldResult> origSizes =
+        getAsOpFoldResult(newExtractStridedMetadata.getSizes());
+    SmallVector<OpFoldResult> origStrides =
+        getAsOpFoldResult(newExtractStridedMetadata.getStrides());
+    unsigned idx = 0, endIdx = expandShape.getReassociationIndices().size();
+    for (; idx != endIdx; ++idx) {
+      SmallVector<OpFoldResult> expandedSizes =
+          getExpandedSizes(expandShape, rewriter, origSizes, /*groupId=*/idx);
+      SmallVector<OpFoldResult> expandedStrides = getExpandedStrides(
+          expandShape, rewriter, origSizes, origStrides, /*groupId=*/idx);
+
+      unsigned groupSize = expandShape.getReassociationIndices()[idx].size();
+      const unsigned sizeStartIdx = baseIdxInResult;
+      const unsigned strideStartIdx = sizeStartIdx + expandShapeRank;
+      for (unsigned i = 0; i < groupSize; ++i) {
+        results[sizeStartIdx + i] = expandedSizes[i];
+        results[strideStartIdx + i] = expandedStrides[i];
+      }
+      baseIdxInResult += groupSize;
+    }
+    assert(idx == sourceRank &&
+           "We should have visited all the input dimensions");
+    assert(baseIdxInResult == expandShapeRank + 2 &&
+           "We should have populated all the values");
+    rewriter.replaceOp(
+        op, getValueOrCreateConstantIndexOp(rewriter, origLoc, results));
+    return success();
+  }
+};
 } // namespace
 
 void memref::populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<ExtractStridedMetadataOpSubviewFolder>(patterns.getContext());
+  patterns.add<ExtractStridedMetadataOpSubviewFolder,
+               ExtractStridedMetadataOpExpandShapeFolder>(
+      patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

func.func @extract_strided_metadata_of_subview_all_dynamic(
 
   %base_buffer, %offset, %sizes:3, %strides:3 = memref.extract_strided_metadata %subview :
     memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
     -> memref<f32>, index, index, index, index, index, index, index
 
   return %base_buffer, %offset, %sizes#0, %sizes#1, %sizes#2, %strides#0, %strides#1, %strides#2 :
     memref<f32>, index, index, index, index, index, index, index
 }
+
+// -----
+
+// Check that we properly simplify extract_strided_metadata of expand_shape
+// into:
+// baseBuffer, baseOffset, baseSizes, baseStrides =
+//     extract_strided_metadata(memref)
+// sizes#reassIdx =
+//     baseSizes#reassDim / product(expandShapeSizes#j,
+//                                  for j in group excluding reassIdx)
+// strides#reassIdx =
+//     baseStrides#reassDim * product(expandShapeSizes#j, for j in
+//                                    reassIdx+1..reassIdx+group.size)
+//
+// Here we have:
+// For the group applying to dim0:
+// size 0 = 3
+// size 1 = 5
+// size 2 = 2
+// stride 0 = baseStrides#0 * 5 * 2
+//          = 4 * 5 * 2
+//          = 40
+// stride 1 = baseStrides#0 * 2
+//          = 4 * 2
+//          = 8
+// stride 2 = baseStrides#0
+//          = 4
+//
+// For the group applying to dim1:
+// size 3 = 2
+// size 4 = 2
+// stride 3 = baseStrides#1 * 2
+//          = 1 * 2
+//          = 2
+// stride 4 = baseStrides#1
+//          = 1
+//
+// Base and offset are unchanged.
+//
+// CHECK-LABEL: func @extract_strided_metadata_of_expand_shape_all_static
+//  CHECK-SAME: (%[[ARG:.*]]: memref<30x4xi16>)
+//
+//   CHECK-DAG: %[[C40:.*]] = arith.constant 40 : index
+//   CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+//   CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
+//   CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+//   CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//
+//   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]] : memref<30x4xi16> -> memref<i16>, index, index, index, index, index
+//
+//   CHECK: return %[[BASE]], %[[C0]], %[[C3]], %[[C5]], %[[C2]], %[[C2]], %[[C2]], %[[C40]], %[[C8]], %[[C4]], %[[C2]], %[[C1]] : memref<i16>, index, index, index, index, index, index, index, index, index, index, index
+func.func @extract_strided_metadata_of_expand_shape_all_static(
+    %arg : memref<30x4xi16>)
+    -> (memref<i16>, index,
+       index, index, index, index, index,
+       index, index, index, index, index) {
+
+  %expand_shape = memref.expand_shape %arg[[0, 1, 2], [3, 4]] :
+    memref<30x4xi16> into memref<3x5x2x2x2xi16>
+
+  %base, %offset, %sizes:5, %strides:5 = memref.extract_strided_metadata %expand_shape :
+    memref<3x5x2x2x2xi16>
+    -> memref<i16>, index,
+       index, index, index, index, index,
+       index, index, index, index, index
+
+  return %base, %offset,
+    %sizes#0, %sizes#1, %sizes#2, %sizes#3, %sizes#4,
+    %strides#0, %strides#1, %strides#2, %strides#3, %strides#4 :
+      memref<i16>, index,
+      index, index, index, index, index,
+      index, index, index, index, index
+}
+
+// -----
+
+// Check that we properly simplify extract_strided_metadata of expand_shape
+// when dynamic sizes, strides, and offsets are involved.
+// See extract_strided_metadata_of_expand_shape_all_static for an explanation
+// of the expansion.
+//
+// One of the important characteristic of this test is that the dynamic
+// dimensions produced by the expand_shape appear both in the first dimension
+// (for group 1) and the non-first dimension (second dimension for group 2.)
+// The idea is to make sure that:
+// 1. We properly account for dynamic shapes even when the strides are not
+//    affected by them. (When the dynamic dimension is the first one.)
+// 2. We properly compute the strides affected by dynamic shapes. (When the
+//    dynamic dimension is not the first one.)
+//
+// Here we have:
+// For the group applying to dim0:
+// size 0 = baseSizes#0 / (all static sizes in that group)
+//        = baseSizes#0 / (7 * 8 * 9)
+//        = baseSizes#0 / 504
+// size 1 = 7
+// size 2 = 8
+// size 3 = 9
+// stride 0 = baseStrides#0 * 7 * 8 * 9
+//          = baseStrides#0 * 504
+// stride 1 = baseStrides#0 * 8 * 9
+//          = baseStrides#0 * 72
+// stride 2 = baseStrides#0 * 9
+// stride 3 = baseStrides#0
+//
+// For the group applying to dim1:
+// size 4 = 10
+// size 5 = 2
+// size 6 = baseSizes#1 / (all static sizes in that group)
+//        = baseSizes#1 / (10 * 2 * 3)
+//        = baseSizes#1 / 60
+// size 7 = 3
+// stride 4 = baseStrides#1 * size 5 * size 6 * size 7
+//          = baseStrides#1 * 2 * (baseSizes#1 / 60) * 3
+//          = baseStrides#1 * (baseSizes#1 / 60) * 6
+//          and since we know that baseSizes#1 is a multiple of 60:
+//          = baseStrides#1 * (baseSizes#1 / 10)
+// stride 5 = baseStrides#1 * size 6 * size 7
+//          = baseStrides#1 * (baseSizes#1 / 60) * 3
+//          = baseStrides#1 * (baseSizes#1 / 20)
+// stride 6 = baseStrides#1 * size 7
+//          = baseStrides#1 * 3
+// stride 7 = baseStrides#1
+//
+// Base and offset are unchanged.
+//
+//   CHECK-DAG: #[[$DIM0_SIZE_MAP:.*]] = affine_map<()[s0] -> (s0 floordiv 504)>
+//   CHECK-DAG: #[[$DIM6_SIZE_MAP:.*]] = affine_map<()[s0] -> (s0 floordiv 60)>
+//
+//   CHECK-DAG: #[[$DIM0_STRIDE_MAP:.*]] = affine_map<()[s0] -> (s0 * 504)>
+//   CHECK-DAG: #[[$DIM1_STRIDE_MAP:.*]] = affine_map<()[s0] -> (s0 * 72)>
+//   CHECK-DAG: #[[$DIM2_STRIDE_MAP:.*]] = affine_map<()[s0] -> (s0 * 9)>
+//   CHECK-DAG: #[[$DIM4_STRIDE_MAP:.*]] = affine_map<()[s0, s1] -> ((s0 floordiv 10) * s1)>
+//   CHECK-DAG: #[[$DIM5_STRIDE_MAP:.*]] = affine_map<()[s0, s1] -> ((s0 floordiv 20) * s1)>
+//   CHECK-DAG: #[[$DIM6_STRIDE_MAP:.*]] = affine_map<()[s0] -> (s0 * 3)>
+// CHECK-LABEL: func @extract_strided_metadata_of_expand_shape_all_dynamic
+//  CHECK-SAME: (%[[ARG:.*]]: memref<?x?xf32,
+//
+//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//   CHECK-DAG: %[[C9:.*]] = arith.constant 9 : index
+//   CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+//   CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index
+//   CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//
+//   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]] : memref<?x?xf32, strided<[?, ?], offset: ?>> -> memref<f32>, index, index, index, index, index
+//
+//   CHECK-DAG: %[[DYN_SIZE0:.*]] = affine.apply #[[$DIM0_SIZE_MAP]]()[%[[SIZES]]#0]
+//   CHECK-DAG: %[[DYN_SIZE6:.*]] = affine.apply #[[$DIM6_SIZE_MAP]]()[%[[SIZES]]#1]
+//   CHECK-DAG: %[[DYN_STRIDE0:.*]] = affine.apply #[[$DIM0_STRIDE_MAP]]()[%[[STRIDES]]#0]
+//   CHECK-DAG: %[[DYN_STRIDE1:.*]] = affine.apply #[[$DIM1_STRIDE_MAP]]()[%[[STRIDES]]#0]
+//   CHECK-DAG: %[[DYN_STRIDE2:.*]] = affine.apply #[[$DIM2_STRIDE_MAP]]()[%[[STRIDES]]#0]
+//   CHECK-DAG: %[[DYN_STRIDE4:.*]] = affine.apply #[[$DIM4_STRIDE_MAP]]()[%[[SIZES]]#1, %[[STRIDES]]#1]
+//   CHECK-DAG: %[[DYN_STRIDE5:.*]] = affine.apply #[[$DIM5_STRIDE_MAP]]()[%[[SIZES]]#1, %[[STRIDES]]#1]
+//   CHECK-DAG: %[[DYN_STRIDE6:.*]] = affine.apply #[[$DIM6_STRIDE_MAP]]()[%[[STRIDES]]#1]
+
+//   CHECK: return %[[BASE]], %[[OFFSET]], %[[DYN_SIZE0]], %[[C7]], %[[C8]], %[[C9]], %[[C10]], %[[C2]], %[[DYN_SIZE6]], %[[C3]], %[[DYN_STRIDE0]], %[[DYN_STRIDE1]], %[[DYN_STRIDE2]], %[[STRIDES]]#0, %[[DYN_STRIDE4]], %[[DYN_STRIDE5]], %[[DYN_STRIDE6]], %[[STRIDES]]#1 : memref<f32>, index, index, index, index, index, index, index, index, index, index, index, index, index
+func.func @extract_strided_metadata_of_expand_shape_all_dynamic(
+    %base: memref<?x?xf32, strided<[?,?], offset:?>>,
+    %offset0: index, %offset1: index, %offset2: index,
+    %size0: index, %size1: index, %size2: index,
+    %stride0: index, %stride1: index, %stride2: index)
+    -> (memref<f32>, index,
+       index, index, index, index, index, index, index, index,
+       index, index, index, index, index, index, index, index) {
+
+  %subview = memref.expand_shape %base[[0, 1, 2, 3],[4, 5, 6, 7]] :
+    memref<?x?xf32, strided<[?,?], offset: ?>> into
+      memref<?x7x8x9x10x2x?x3xf32, strided<[?, ?, ?, ?, ?, ?, ?, ?], offset: ?>>
+
+  %base_buffer, %offset, %sizes:8, %strides:8 = memref.extract_strided_metadata %subview :
+    memref<?x7x8x9x10x2x?x3xf32, strided<[?, ?, ?, ?, ?, ?, ?, ?], offset: ?>>
+    -> memref<f32>, index,
+       index, index, index, index, index, index, index, index,
+       index, index, index, index, index, index, index, index
+
+  return %base_buffer, %offset,
+    %sizes#0, %sizes#1, %sizes#2, %sizes#3, %sizes#4, %sizes#5, %sizes#6, %sizes#7,
+    %strides#0, %strides#1, %strides#2, %strides#3, %strides#4, %strides#5, %strides#6, %strides#7 :
+      memref<f32>, index,
+      index, index, index, index, index, index, index, index,
+      index, index, index, index, index, index, index, index
+}
+
+
+// -----
+
+// Check that we properly handle extract_strided_metadata of expand_shape for
+// 0-D input.
+// The 0-D case is pretty boring:
+// All expanded sizes are 1, likewise for the strides, and we keep the
+// original base and offset.
+// We have still a test for it, because since the input reassociation map
+// of the expand_shape is empty, the handling of such shape hits a corner
+// case.
+// CHECK-LABEL: func @extract_strided_metadata_of_expand_shape_all_static_0_rank
+//  CHECK-SAME: (%[[ARG:.*]]: memref<i16, strided<[], offset: ?>>)
+//
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//
+//   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]] = memref.extract_strided_metadata %[[ARG]] : memref<i16, strided<[], offset: ?>> -> memref<i16>, index
+//
+//   CHECK: return %[[BASE]], %[[OFFSET]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]] : memref<i16>, index, index, index, index, index, index, index, index, index, index, index
+func.func @extract_strided_metadata_of_expand_shape_all_static_0_rank(
+    %arg : memref<i16, strided<[], offset: ?>>)
+    -> (memref<i16>, index,
+       index, index, index, index, index,
+       index, index, index, index, index) {
+
+  %expand_shape = memref.expand_shape %arg[] :
+    memref<i16, strided<[], offset: ?>> into memref<1x1x1x1x1xi16, strided<[1,1,1,1,1], offset: ?>>
+
+  %base, %offset, %sizes:5, %strides:5 = memref.extract_strided_metadata %expand_shape :
+    memref<1x1x1x1x1xi16, strided<[1,1,1,1,1], offset: ?>>
+    -> memref<i16>, index,
+       index, index, index, index, index,
+       index, index, index, index, index
+
+  return %base, %offset,
+    %sizes#0, %sizes#1, %sizes#2, %sizes#3, %sizes#4,
+    %strides#0, %strides#1, %strides#2, %strides#3, %strides#4 :
+      memref<i16>, index,
+      index, index, index, index, index,
+      index, index, index, index, index
+}