[mlir][memref] Fix ExpandShapeOp builder

[springerm]

Summary

The result type of an ExpandShapeOp cannot be inferred from the source type and the reassociations. The result shape must be specified. Strides can be computed.

[mravishankar]

I just have a question here (looking through the code, but before I dig into more details). Is it even valid to reshape a memref with strides. The memref could be result of a subview, but if the subview is not itself of a contiguous region of memory, then the reshape is invalid (you cannot reshape any arbitrary subview). Apart from this the rest of looks fine.

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

1231

Nit: Could we make the order resultShape, src, reassociation, attrs? This is consistent with the resulttype, operands, attrs scheme, except the result type is replaced with result shape.

[springerm]

In D112511#3088944, @mravishankar wrote:

I just have a question here (looking through the code, but before I dig into more details). Is it even valid to reshape a memref with strides. The memref could be result of a subview, but if the subview is not itself of a contiguous region of memory, then the reshape is invalid (you cannot reshape any arbitrary subview). Apart from this the rest of looks fine.

I think this is a problem only for collapse_shape, where two non-contiguous regions may be collapsed. I think that's what isReshapableDimBand is checking.

In the case of expand_shape, we split a dim into multiple dims, so just adding new strides/sizes. I think this can always be done without a copy. I'm still a bit confused about strides etc. (esp. when they are dynamic), so maybe I'm missing something here.

[springerm]

address comments

[nicolasvasilache]

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

1494

We've been moving away from AffineMap from reassociations.
If we don't use m.compose(...) we should just use the list of indices.

1506

do you need the full AffineExpr or can you use some of the saturated arith stuff ?

1516

use ShapeType::isDynamic(idx), here and below plz.

[mravishankar]

In the case of expand_shape, we split a dim into multiple dims, so just adding new strides/sizes. I think this can always be done without a copy. I'm still a bit confused about strides etc. (esp. when they are dynamic), so maybe I'm missing something here.

yeah, that makes sense. I tried a few examples, and couldnt find a counter-example to this. So by a totally unscientific method, this might be fine.

Maybe some lit-tests for the expanded_shape case with non-trivial strides would help catch bugs?

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

 def IndexListArrayAttr :
   TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
 
 class MemRef_ReassociativeReshapeOp<string mnemonic, list<OpTrait> traits = []> :
     MemRef_Op<mnemonic, !listconcat(traits,
       [NoSideEffect, ViewLikeOpInterface])>,
     Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyStridedMemRef:$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);
-      build($_builder, $_state, src, reassociationMaps, attrs);
-    }]>,
-
-    // Builders for a reshape whose result type is passed explicitly. This may
-    // be either a contracting or expanding reshape.
-    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);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
-    }]>
-  ];
 
   code commonExtraClassDeclaration = [{
     SmallVector<AffineMap, 4> getReassociationMaps();
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : reassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
let description = [{
     Example:
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
     %1 = memref.expand_shape %0 [[0, 1], [2]] :
       memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
     ```
   }];
+
+  let builders = [
+    // Builders for a contracting reshape whose result type is computed from
+    // `src`, `reassociation` and `resultShape`.
+    OpBuilder<(ins "Value":$src,
+      "ArrayRef<ReassociationIndices>":$reassociation,
+      "ArrayRef<int64_t>":$resultShape,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+    OpBuilder<(ins "Value":$src,
+      "ArrayRef<ReassociationExprs>":$reassociation,
+      "ArrayRef<int64_t>":$resultShape,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+    [{
+      auto reassociationMaps =
+          convertReassociationMapsToIndices($_builder, reassociation);
+      build($_builder, $_state, src, reassociationMaps, resultShape, attrs);
+    }]>,
+
+    // Builders for a reshape whose result type is passed explicitly. This may
+    // be either a contracting or expanding reshape.
+    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);
+      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+    }]>
+  ];
+
   let extraClassDeclaration = commonExtraClassDeclaration;
 }
 
 def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape"> {
   let summary = "operation to produce a memref with a smaller rank.";
   let description = [{
     The `memref.collapse_shape` op produces a new view with a smaller rank
     whose sizes are a reassociation of the original `view`. Depending on
let description = [{
     Examples:
 
     ```mlir
     // Dimension collapse (i, j) -> i' and k -> k'
     %1 = memref.collapse_shape %0 [[0, 1], [2]] :
       memref<?x?x?xf32, stride_spec> into memref<?x?xf32, stride_spec_2>
     ```
   }];
+
+  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);
+      build($_builder, $_state, src, reassociationMaps, attrs);
+    }]>,
+
+    // Builders for a reshape whose result type is passed explicitly. This may
+    // be either a contracting or expanding reshape.
+    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);
+      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+    }]>
+  ];
+
   let extraClassDeclaration = commonExtraClassDeclaration;
 }
 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_StoreOp : MemRef_Op<"store",

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

for (auto idx = dim, e = dim + extent; idx + 1 < e; ++idx) {
     if (strides[idx] != strides[idx + 1] * sizes[idx + 1])
       return false;
   }
   return true;
 }
 
 /// Compute the MemRefType obtained by applying the `reassociation` (which is
 /// expected to be valid) to `type`.
-/// If `type` is Contiguous MemRefType, this always produce a contiguous
+/// If `type` is a contiguous MemRefType, this always produce a contiguous
+/// MemRefType.
+static MemRefType computeReshapeExpandedType(MemRefType type,
+                                             ArrayRef<AffineMap> reassociation,
+                                             ArrayRef<int64_t> resultShape) {
+  // Early-exit: if `type` is contiguous, the result must be contiguous.
+  if (canonicalizeStridedLayout(type).getLayout().isIdentity())
+    return MemRefType::Builder(type).setShape(resultShape).setLayout({});
+
+  AffineExpr offset;
+  SmallVector<AffineExpr, 4> strides;
+  auto status = getStridesAndOffset(type, strides, offset);
+  (void)status;
+  assert(succeeded(status) && "expected strided memref");
+
+  SmallVector<AffineExpr, 4> newStrides(resultShape.size(), AffineExpr());
+  newStrides.reserve(reassociation.size());
+
+  // Use the fact that reassociation is valid to simplify the logic: only use
+  // each map's rank.
+  assert(isReassociationValid(reassociation) && "invalid reassociation");
+  int64_t resultDim = resultShape.size() - 1;
+  int64_t sourceDim = strides.size() - 1;
+  for (AffineMap m : llvm::reverse(reassociation)) {
+    newStrides[resultDim] = strides[sourceDim];
+    int64_t numDynamicDims = static_cast<int64_t>(resultShape[resultDim] ==
+                                                  ShapedType::kDynamicSize);
+    --resultDim;
+
+    // If this loop has iterations, this source dim is expanded.
+    for (int64_t r = static_cast<int64_t>(m.getNumResults()) - 2; r >= 0; --r) {
+      numDynamicDims += static_cast<int64_t>(resultShape[resultDim] ==
+                                             ShapedType::kDynamicSize);
+      assert(numDynamicDims < 2 &&
+             "at most one result dim of an expansion may be dynamic");
+      if (resultShape[resultDim + 1] != ShapedType::kDynamicSize &&
+          newStrides[resultDim + 1])
+        newStrides[resultDim] =
+            newStrides[resultDim + 1] * resultShape[resultDim + 1];
+      --resultDim;
+    }
+
+    --sourceDim;
+  }
+
+  assert(resultDim == -1 && "did not process all result dims");
+  assert(sourceDim == -1 && "did not process all source dims");
+
+  // Convert back to int64_t because we don't have enough information to create
+  // new strided layouts from AffineExpr only. This corresponds to a case where
+  // copies may be necessary.
+  int64_t intOffset = ShapedType::kDynamicStrideOrOffset;
+  if (auto o = offset.dyn_cast<AffineConstantExpr>())
+    intOffset = o.getValue();
+  SmallVector<int64_t, 4> intStrides;
+  intStrides.reserve(newStrides.size());
+  for (auto stride : newStrides) {
+    if (auto cst = stride.dyn_cast_or_null<AffineConstantExpr>())
+      intStrides.push_back(cst.getValue());
+    else
+      intStrides.push_back(ShapedType::kDynamicStrideOrOffset);
+  }
+  auto layout =
+      makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());
+  return canonicalizeStridedLayout(MemRefType::Builder(type)
+                                       .setShape(resultShape)
+                                       .setLayout(AffineMapAttr::get(layout)));
+}
+
+/// Compute the MemRefType obtained by applying the `reassociation` (which is
+/// expected to be valid) to `type`.
+/// If `type` is a contiguous MemRefType, this always produce a contiguous
 /// MemRefType.
 static MemRefType
 computeReshapeCollapsedType(MemRefType type,
                             ArrayRef<AffineMap> reassociation) {
   auto sizes = type.getShape();
   AffineExpr offset;
   SmallVector<AffineExpr, 4> strides;
   auto status = getStridesAndOffset(type, strides, offset);
auto layout =
       makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());
   return canonicalizeStridedLayout(
       MemRefType::Builder(type).setShape(newSizes).setLayout(
           AffineMapAttr::get(layout)));
 }
 
 void ExpandShapeOp::build(OpBuilder &b, OperationState &result, Value src,
                           ArrayRef<ReassociationIndices> reassociation,
+                          ArrayRef<int64_t> resultShape,
                           ArrayRef<NamedAttribute> attrs) {
   auto memRefType = src.getType().cast<MemRefType>();
-  auto resultType = computeReshapeCollapsedType(
-      memRefType, getSymbolLessAffineMaps(convertReassociationIndicesToExprs(
-                      b.getContext(), reassociation)));
+  auto resultType = computeReshapeExpandedType(
+      memRefType,
+      getSymbolLessAffineMaps(
+          convertReassociationIndicesToExprs(b.getContext(), reassociation)),
+      resultShape);
   build(b, result, resultType, src, attrs);
   result.addAttribute(getReassociationAttrName(),
                       getReassociationIndicesAttribute(b, reassociation));
 }
 
 void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
                             ArrayRef<ReassociationIndices> reassociation,
                             ArrayRef<NamedAttribute> attrs) {
   auto memRefType = src.getType().cast<MemRefType>();
   auto resultType = computeReshapeCollapsedType(
       memRefType, getSymbolLessAffineMaps(convertReassociationIndicesToExprs(
                       b.getContext(), reassociation)));
   build(b, result, resultType, src, attrs);
   result.addAttribute(getReassociationAttrName(),
                       getReassociationIndicesAttribute(b, reassociation));
 }
 
-template <typename ReshapeOp,
-          bool isExpansion = std::is_same<ReshapeOp, ExpandShapeOp>::value>
-static LogicalResult verifyReshapeOp(ReshapeOp op, MemRefType expandedType,
-                                     MemRefType collapsedType) {
-  if (failed(
-          verifyReshapeLikeTypes(op, expandedType, collapsedType, isExpansion)))
+static LogicalResult verify(ExpandShapeOp op) {
+  if (failed(verifyReshapeLikeTypes(op, op.getResultType(), op.getSrcType(),
+                                    /*isExpansion=*/true)))
     return failure();
   auto maps = op.getReassociationMaps();
-  MemRefType expectedType = computeReshapeCollapsedType(expandedType, maps);
-  if (collapsedType != expectedType)
-    return op.emitOpError("expected collapsed type to be ")
-           << expectedType << ", but got " << collapsedType;
+  MemRefType expectedType = computeReshapeExpandedType(
+      op.getSrcType(), maps, op.getResultType().getShape());
+  if (op.getResultType() != expectedType)
+    return op.emitOpError("expected result type to be ")
+           << expectedType << ", but got " << op.getResultType();
   return success();
 }
 
-static LogicalResult verify(ExpandShapeOp op) {
-  return verifyReshapeOp(op, op.getResultType(), op.getSrcType());
-}
-
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<CollapseReshapeOps<ExpandShapeOp>,
               CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>>(context);
 }
 
 static LogicalResult verify(CollapseShapeOp op) {
-  return verifyReshapeOp(op, op.getSrcType(), op.getResultType());
+  if (failed(verifyReshapeLikeTypes(op, op.getSrcType(), op.getResultType(),
+                                    /*isExpansion=*/false)))
+    return failure();
+  auto maps = op.getReassociationMaps();
+  MemRefType expectedType = computeReshapeCollapsedType(op.getSrcType(), maps);
+  if (op.getResultType() != expectedType)
+    return op.emitOpError("expected result type to be ")
+           << expectedType << ", but got " << op.getResultType();
+  return success();
 }
 
 struct CollapseShapeOpMemRefCastFolder
     : public OpRewritePattern<CollapseShapeOp> {
 public:
   using OpRewritePattern<CollapseShapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CollapseShapeOp op,

mlir/test/Dialect/MemRef/invalid.mlir

func @collapse_shape_invalid_reassociation(%arg0: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected reassociation map #1 to be valid and contiguous}}
   %0 = memref.collapse_shape %arg0 [[0, 1], [1, 2]] :
     memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 
 // -----
 
 func @collapse_shape_wrong_collapsed_type(%arg0: memref<?x?x?xf32>) {
-  // expected-error @+1 {{expected collapsed type to be 'memref<?x?xf32>', but got 'memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>'}}
+  // expected-error @+1 {{expected result type to be 'memref<?x?xf32>', but got 'memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>'}}
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
     memref<?x?x?xf32> into memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>
 }
 
 // -----
 
 func @expand_shape_illegal_dynamic_memref
   (%arg0: memref<?x?x?xf32>) -> memref<?x?x?x4x?xf32> {