[mlir] Move linalg::Expand/CollapseShapeOp to memref dialect.

[pifon2a]

Summary

RFC: https://llvm.discourse.group/t/rfc-reshape-ops-restructuring/3310

[pifon2a]

Add ViewLikeInterface back.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

class Linalg_ReshapeLikeOp<string mnemonic, list<OpTrait> traits = []> :
   }];
 
   let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
 }
 
 def IndexListArrayAttr :
   TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
 
-class Linalg_ReshapeOp<string mnemonic> : Linalg_ReshapeLikeOp<mnemonic,
-    [DeclareOpInterfaceMethods<ViewLikeOpInterface>]>,
-    Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
-    Results<(outs AnyStridedMemRef:$result)> {
-  let extraClassDeclaration = commonExtraClassDeclaration # [{
-    MemRefType getSrcType() { return src().getType().cast<MemRefType>(); }
-    MemRefType getResultType() { return result().getType().cast<MemRefType>(); }
-  }];
-  let hasFolder = 1;
-  let hasCanonicalizer = 1;
-  let printer = [{ return ::print(p, *this); }];
-}
-
-def Linalg_ExpandShapeOp : Linalg_ReshapeOp<"expand_shape"> {
-  let summary = "operation to produce a memref with a higher rank.";
-  let description = [{
-    The `linalg.expand_shape` op produces a new view with a higher rank whose
-    sizes are a reassociation of the original `view`. Depending on whether or
-    not the reassociated MemRefType is contiguous, the resulting memref may
-    require explicit alloc and copies.
-
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of I64ArrayAttr attribute.
-
-    For now, it is assumed that either:
-      1. a reassociation produces and consumes contiguous MemRefType or,
-      2. the reshape op will be folded into its consumers (by changing the shape
-         of the computations).
-    All other cases are undefined behavior and a reshape op may not lower to
-    LLVM if it cannot be proven statically that it does not require alloc+copy.
-
-    The operand memref type when dimensions can be zero-ranked if the result
-    memref type is statically shaped with all dimensions being unit extent. In
-    such case the reassociation map is empty.
-
-    The verification rule is that the reassociation maps are applied to the
-    result memref with the larger rank to obtain the operand memref with the
-    smaller rank.
-
-    Example:
-
-    ```mlir
-    // Dimension expansion i -> (i', j') and (k) -> (k')
-    %1 = linalg.expand_shape %0 [[0, 1], [2]] :
-      memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
-    ```
-  }];
-}
-
-def Linalg_CollapseShapeOp : Linalg_ReshapeOp<"collapse_shape"> {
-  let summary = "operation to produce a memref with a smaller rank.";
-  let description = [{
-    The `linalg.collapse_shape` op produces a new view with a smaller rank
-    whose sizes are a reassociation of the original `view`. Depending on
-    whether or not the reassociated MemRefType is contiguous, the resulting
-    memref may require explicit alloc and copies.
-
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of I64ArrayAttr attribute.
-
-    For now, it is assumed that either:
-      1. a reassociation produces and consumes contiguous MemRefType or,
-      2. the reshape op will be folded into its consumers (by changing the shape
-         of the computations).
-    All other cases are undefined behavior and a reshape op may not lower to
-    LLVM if it cannot be proven statically that it does not require alloc+copy.
-
-    The result memref type of a reshape can be zero-ranked if the operand
-    memref type is statically shaped with all dimensions being unit extent. In
-    such case the reassociation map is empty.
-
-    The verification rule is that the reassociation maps are applied to the
-    operand memref with the larger rank to obtain the result memref with the
-    smaller rank.
-
-    Examples:
-
-    ```mlir
-    // Dimension collapse (i, j) -> i' and k -> k'
-    %1 = linalg.collapse_shape %0 [[0, 1], [2]] :
-      memref<?x?x?xf32, stride_spec> into memref<?x?xf32, stride_spec_2>
-    ```
-  }];
-}
-
 class Linalg_TensorReshapeOp<string mnemonic> : Linalg_ReshapeLikeOp<
     mnemonic,
     [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
       ["reifyReturnTypeShapesPerResultDim"]>]>,
     Arguments<(ins AnyTensor:$src,
                    IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
   let extraClassDeclaration = commonExtraClassDeclaration # [{

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 //===- MemRef.h - MemRef dialect --------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_IR_MEMREF_H_
 #define MLIR_DIALECT_MEMREF_IR_MEMREF_H_
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 

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

let description = [{
     chosen.
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // AllocaScopeOp
 //===----------------------------------------------------------------------===//
 
-def MemRef_AllocaScopeOp : MemRef_Op<"alloca_scope", 
+def MemRef_AllocaScopeOp : MemRef_Op<"alloca_scope",
       [AutomaticAllocationScope,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"AllocaScopeReturnOp">,
        RecursiveSideEffects,
        NoRegionArguments]> {
   let summary = "explicitly delimited scope for stack allocation";
   let description = [{
     The `memref.alloca_scope` operation represents an explicitly-delimited
     scope for the alloca allocations. Any `memref.alloca` operations that are
     used within this scope are going to be cleaned up automatically once
     the control-flow exits the nested region. For example:
 
     ```mlir
     memref.alloca_scope {
       %myalloca = memref.alloca(): memref<4x3xf32>
       ...
     }
     ```
 
     Here, `%myalloca` memref is valid within the explicitly delimited scope
     and is automatically deallocated at the end of the given region. Conceptually,
-    `memref.alloca_scope` is a passthrough operation with 
-    `AutomaticAllocationScope` that spans the body of the region within the operation. 
+    `memref.alloca_scope` is a passthrough operation with
+    `AutomaticAllocationScope` that spans the body of the region within the operation.
 
     `memref.alloca_scope` may also return results that are defined in the nested
     region. To return a value, one should use `memref.alloca_scope.return`
     operation:
 
     ```mlir
     %result = memref.alloca_scope {
       ...
       memref.alloca_scope.return %value
     }
     ```
 
     If `memref.alloca_scope` returns no value, the `memref.alloca_scope.return ` can
     be left out, and will be inserted implicitly.
   }];
 
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$bodyRegion);
 }
 
 //===----------------------------------------------------------------------===//
 // AllocaScopeReturnOp
 //===----------------------------------------------------------------------===//
 
-def MemRef_AllocaScopeReturnOp : MemRef_Op<"alloca_scope.return", 
+def MemRef_AllocaScopeReturnOp : MemRef_Op<"alloca_scope.return",
       [HasParent<"AllocaScopeOp">,
        NoSideEffect,
        ReturnLike,
        Terminator]> {
   let summary = "terminator for alloca_scope operation";
   let description = [{
-    `memref.alloca_scope.return` operation returns zero or more SSA values 
+    `memref.alloca_scope.return` operation returns zero or more SSA values
     from the region within `memref.alloca_scope`. If no values are returned,
     the return operation may be omitted. Otherwise, it has to be present
     to indicate which values are going to be returned. For example:
 
     ```mlir
     memref.alloca_scope.return %value
     ```
   }];
let extraClassDeclaration = [{
     Value getViewSource() { return source(); }
   }];
 
   let assemblyFormat = [{
     $source `(` $shape `)` attr-dict `:` functional-type(operands, results)
   }];
 }
 
+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])>,
+    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>(
+            llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
+              return indexAttr.cast<IntegerAttr>().getInt();
+            })));
+      return reassociationIndices;
+    };
+    MemRefType getSrcType() { return src().getType().cast<MemRefType>(); }
+    MemRefType getResultType() { return result().getType().cast<MemRefType>(); }
+    Value getViewSource() { return src(); }
+  }];
+
+  let hasFolder = 1;
+  let hasCanonicalizer = 1;
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
+}
+
+def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape"> {
+  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`. Depending on whether or
+    not the reassociated MemRefType is contiguous, the resulting memref may
+    require explicit alloc and copies.
+
+    A reassociation is defined as a continuous grouping of dimensions and is
+    represented with an array of I64ArrayAttr attribute.
+
+    For now, it is assumed that either:
+      1. a reassociation produces and consumes contiguous MemRefType or,
+      2. the reshape op will be folded into its consumers (by changing the shape
+         of the computations).
+    All other cases are undefined behavior and a reshape op may not lower to
+    LLVM if it cannot be proven statically that it does not require alloc+copy.
+
+    The operand memref type when dimensions can be zero-ranked if the result
+    memref type is statically shaped with all dimensions being unit extent. In
+    such case the reassociation map is empty.
+
+    The verification rule is that the reassociation maps are applied to the
+    result memref with the larger rank to obtain the operand memref with the
+    smaller rank.
+
+    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 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
+    whether or not the reassociated MemRefType is contiguous, the resulting
+    memref may require explicit alloc and copies.
+
+    A reassociation is defined as a continuous grouping of dimensions and is
+    represented with an array of I64ArrayAttr attribute.
+
+    For now, it is assumed that either:
+      1. a reassociation produces and consumes contiguous MemRefType or,
+      2. the reshape op will be folded into its consumers (by changing the shape
+         of the computations).
+    All other cases are undefined behavior and a reshape op may not lower to
+    LLVM if it cannot be proven statically that it does not require alloc+copy.
+
+    The result memref type of a reshape can be zero-ranked if the operand
+    memref type is statically shaped with all dimensions being unit extent. In
+    such case the reassociation map is empty.
+
+    The verification rule is that the reassociation maps are applied to the
+    operand memref with the larger rank to obtain the result memref with the
+    smaller rank.
+
+    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 extraClassDeclaration = commonExtraClassDeclaration;
+}
+
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_StoreOp : MemRef_Op<"store",
      [TypesMatchWith<"type of 'value' matches element type of 'memref'",
                      "memref", "value",
                      "$_self.cast<MemRefType>().getElementType()">,

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

-//===- RehshapeOpsUtils.h - Utilities used by reshape ops --*- C++ -*------===//
+//===- ReshapeOpsUtils.h - Utilities used by reshape ops --*- C++ -*------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines utilities and common canonicalization patterns for

 
 namespace mlir {
 
 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();
+constexpr StringRef getReassociationAttrName() { return "reassociation"; }
 
 /// 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]]
 ///
 /// is folded into
 ///
 ///   result = [[0, 1, 2], [3, 4]].
 Optional<SmallVector<ReassociationIndices>> composeReassociationIndices(
     ArrayRef<ReassociationIndices> producerReassociations,
     ArrayRef<ReassociationIndices> consumerReassociations,
     MLIRContext *context);
 
+/// Convert reassociation indices to affine expressions.
+SmallVector<SmallVector<AffineExpr, 2>, 2> convertReassociationIndicesToExprs(
+    OpBuilder &b, ArrayRef<ReassociationIndices> reassociationIndices);
+
+/// Constructs affine maps out of Array<Array<AffineExpr>>.
+SmallVector<AffineMap, 4>
+getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation);
+
+/// 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);
+
 /// Return the reassociations maps to use to reshape given the source type and
 /// the target type when possible. Return llvm::None when this computation
 /// failed.
 Optional<SmallVector<ReassociationIndices>>
 getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType);
 
 /// Return true if the reassociation specification is valid, false otherwise.
 /// When false, the `invalidIndex` integer pointer is optionally filled with the
llvm::interleaveComma(op.reassociation(), p, [&](const Attribute &attr) {
     llvm::interleaveComma(arrayAttr, p, [&](const Attribute &attr) {
       p << attr.cast<IntegerAttr>().getInt();
     });
     p << ']';
   });
 
   p << "] ";
   p.printOptionalAttrDict(op->getAttrs(),
-                          /*elidedAttrs=*/{op.getReassociationAttrName()});
+                          /*elidedAttrs=*/{getReassociationAttrName()});
   p << ": " << op.src().getType() << " into " << op.getType();
 }
 
 template <typename ReshapeOpTy, typename InverseReshapeOpTy>
 static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
                                   ArrayRef<Attribute> operands) {
   // Fold producer-consumer reshape ops that where the operand type of the
   // producer is same as the return type of the consumer.

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 //===- LinalgToLLVM.cpp - conversion from Linalg to LLVM dialect ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/SCF/SCF.h"
desc = b.create<LLVM::InsertValueOp>(desc, adaptor.max(),
                                          rewriter.getI64ArrayAttr(1));
     desc = b.create<LLVM::InsertValueOp>(desc, adaptor.step(),
                                          rewriter.getI64ArrayAttr(2));
     rewriter.replaceOp(rangeOp, desc);
     return success();
   }
 };
 
-// ReshapeOp creates a new view descriptor of the proper rank.
-// For now, the only conversion supported is for target MemRef with static sizes
-// and strides.
-template <typename ReshapeOp>
-class ReshapeOpConversion : public ConvertOpToLLVMPattern<ReshapeOp> {
-public:
-  using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
-  using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;
-
-  LogicalResult
-  matchAndRewrite(ReshapeOp reshapeOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    MemRefType dstType = reshapeOp.getResultType();
-
-    if (!dstType.hasStaticShape())
-      return failure();
-
-    int64_t offset;
-    SmallVector<int64_t, 4> strides;
-    auto res = getStridesAndOffset(dstType, strides, offset);
-    if (failed(res) || llvm::any_of(strides, [](int64_t val) {
-          return ShapedType::isDynamicStrideOrOffset(val);
-        }))
-      return failure();
-
-    ReshapeOpAdaptor adaptor(operands);
-    MemRefDescriptor baseDesc(adaptor.src());
-    Location loc = reshapeOp->getLoc();
-    auto desc =
-        MemRefDescriptor::undef(rewriter, reshapeOp->getLoc(),
-                                this->typeConverter->convertType(dstType));
-    desc.setAllocatedPtr(rewriter, loc, baseDesc.allocatedPtr(rewriter, loc));
-    desc.setAlignedPtr(rewriter, loc, baseDesc.alignedPtr(rewriter, loc));
-    desc.setOffset(rewriter, loc, baseDesc.offset(rewriter, loc));
-    for (auto en : llvm::enumerate(dstType.getShape()))
-      desc.setConstantSize(rewriter, loc, en.index(), en.value());
-    for (auto en : llvm::enumerate(strides))
-      desc.setConstantStride(rewriter, loc, en.index(), en.value());
-    rewriter.replaceOp(reshapeOp, {desc});
-    return success();
-  }
-};
 
 // YieldOp produces and LLVM::ReturnOp.
 class YieldOpConversion : public ConvertOpToLLVMPattern<linalg::YieldOp> {
 public:
   using ConvertOpToLLVMPattern<linalg::YieldOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(linalg::YieldOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, operands);
     return success();
   }
 };
 } // namespace
 
 /// Populate the given list with patterns that convert from Linalg to LLVM.
 void mlir::populateLinalgToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
-  patterns.add<RangeOpConversion, ReshapeOpConversion<ExpandShapeOp>,
-               ReshapeOpConversion<CollapseShapeOp>, YieldOpConversion>(
-      converter);
+  patterns.add<RangeOpConversion, YieldOpConversion>(converter);
 
   // Populate the type conversions for the linalg types.
   converter.addConversion(
       [&](RangeType type) { return convertRangeType(type, converter); });
 }
 
 namespace {
 struct ConvertLinalgToLLVMPass
     : public ConvertLinalgToLLVMBase<ConvertLinalgToLLVMPass> {
   void runOnOperation() override;
 };
 } // namespace
 
 void ConvertLinalgToLLVMPass::runOnOperation() {
   auto module = getOperation();
 
   // Convert to the LLVM IR dialect using the converter defined above.
   RewritePatternSet patterns(&getContext());
   LLVMTypeConverter converter(&getContext());
   populateLinalgToLLVMConversionPatterns(converter, patterns);
+  populateMemRefToLLVMConversionPatterns(converter, patterns);
 
   LLVMConversionTarget target(getContext());
   target.addIllegalOp<RangeOp>();
   target.addLegalOp<ModuleOp, LLVM::DialectCastOp>();
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertLinalgToLLVMPass() {
   return std::make_unique<ConvertLinalgToLLVMPass>();
 }

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

 };
 } // namespace
 
 void ConvertLinalgToStandardPass::runOnOperation() {
   auto module = getOperation();
   ConversionTarget target(getContext());
   target.addLegalDialect<AffineDialect, memref::MemRefDialect, scf::SCFDialect,
                          StandardOpsDialect>();
-  target.addLegalOp<ModuleOp, FuncOp, ReturnOp>();
-  target.addLegalOp<linalg::ExpandShapeOp, linalg::CollapseShapeOp,
-                    linalg::RangeOp>();
+  target.addLegalOp<ModuleOp, FuncOp, ReturnOp, linalg::RangeOp>();
   RewritePatternSet patterns(&getContext());
   populateLinalgToStandardConversionPatterns(patterns);
   if (failed(applyFullConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertLinalgToStandardPass() {
   return std::make_unique<ConvertLinalgToStandardPass>();
 }

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
     // Reset position to beginning of new remainder block.
     rewriter.setInsertionPointToStart(remainder);
 
     *descriptor = targetDesc;
     return success();
   }
 };
 
+// ReshapeOp creates a new view descriptor of the proper rank.
+// For now, the only conversion supported is for target MemRef with static sizes
+// and strides.
+template <typename ReshapeOp>
+class ReassociatingReshapeOpConversion
+    : public ConvertOpToLLVMPattern<ReshapeOp> {
+public:
+  using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
+  using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;
+
+  LogicalResult
+  matchAndRewrite(ReshapeOp reshapeOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType dstType = reshapeOp.getResultType();
+
+    if (!dstType.hasStaticShape())
+      return failure();
+
+    int64_t offset;
+    SmallVector<int64_t, 4> strides;
+    auto res = getStridesAndOffset(dstType, strides, offset);
+    if (failed(res) || llvm::any_of(strides, [](int64_t val) {
+          return ShapedType::isDynamicStrideOrOffset(val);
+        }))
+      return failure();
+
+    ReshapeOpAdaptor adaptor(operands);
+    MemRefDescriptor baseDesc(adaptor.src());
+    Location loc = reshapeOp->getLoc();
+    auto desc =
+        MemRefDescriptor::undef(rewriter, reshapeOp->getLoc(),
+                                this->typeConverter->convertType(dstType));
+    desc.setAllocatedPtr(rewriter, loc, baseDesc.allocatedPtr(rewriter, loc));
+    desc.setAlignedPtr(rewriter, loc, baseDesc.alignedPtr(rewriter, loc));
+    desc.setOffset(rewriter, loc, baseDesc.offset(rewriter, loc));
+    for (auto en : llvm::enumerate(dstType.getShape()))
+      desc.setConstantSize(rewriter, loc, en.index(), en.value());
+    for (auto en : llvm::enumerate(strides))
+      desc.setConstantStride(rewriter, loc, en.index(), en.value());
+    rewriter.replaceOp(reshapeOp, {desc});
+    return success();
+  }
+};
 /// Conversion pattern that transforms a subview op into:
 ///   1. An `llvm.mlir.undef` operation to create a memref descriptor
 ///   2. Updates to the descriptor to introduce the data ptr, offset, size
 ///      and stride.
 /// The subview op is replaced by the descriptor.
 struct SubViewOpLowering : public ConvertOpToLLVMPattern<memref::SubViewOp> {
   using ConvertOpToLLVMPattern<memref::SubViewOp>::ConvertOpToLLVMPattern;
 
patterns.add<
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefCopyOpLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
       PrefetchOpLowering,
+      ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
+      ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
       StoreOpLowering,
       SubViewOpLowering,
       TransposeOpLowering,
       ViewOpLowering>(converter);
   // clang-format on
   auto allocLowering = converter.getOptions().allocLowering;
   if (allocLowering == LowerToLLVMOptions::AllocLowering::AlignedAlloc)
     patterns.add<AlignedAllocOpLowering, DeallocOpLowering>(converter);

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

 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
if (getResultType().hasStaticShape() && getResultType() == getSourceType())
     return source();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, linalg::ExpandShapeOp op) {
-  ::mlir::printReshapeOp<linalg::ExpandShapeOp>(p, op);
-}
-
-static void print(OpAsmPrinter &p, linalg::CollapseShapeOp op) {
-  ::mlir::printReshapeOp<linalg::CollapseShapeOp>(p, op);
-}
-
 static void print(OpAsmPrinter &p, linalg::TensorExpandShapeOp op) {
   ::mlir::printReshapeOp<linalg::TensorExpandShapeOp>(p, op);
 }
 
 static void print(OpAsmPrinter &p, linalg::TensorCollapseShapeOp op) {
   ::mlir::printReshapeOp<linalg::TensorCollapseShapeOp>(p, op);
 }
 
for (const auto &indices : reassociationIndices) {
     reassociationMap.reserve(indices.size());
     for (int64_t index : indices)
       reassociationMap.push_back(b.getAffineDimExpr(index));
     reassociationMaps.push_back(std::move(reassociationMap));
   }
   return reassociationMaps;
 }
 
-SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
-  return getSymbolLessAffineMaps(getReassociationExprs());
-}
-SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {
-  OpBuilder b(this->getContext());
-  return convertReassociationIndicesToExprs(b, getReassociationIndices());
-}
-SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() {
-  return getSymbolLessAffineMaps(getReassociationExprs());
-}
-SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() {
-  OpBuilder b(this->getContext());
-  return convertReassociationIndicesToExprs(b, getReassociationIndices());
-}
 
 SmallVector<AffineMap, 4> TensorCollapseShapeOp::getReassociationMaps() {
   return getSymbolLessAffineMaps(getReassociationExprs());
 }
 SmallVector<ReassociationExprs, 4>
 TensorCollapseShapeOp::getReassociationExprs() {
   OpBuilder b(this->getContext());
   return convertReassociationIndicesToExprs(b, getReassociationIndices());
getReassociationIndicesAttribute(OpBuilder &b,
   SmallVector<Attribute, 4> reassociationAttr =
       llvm::to_vector<4>(llvm::map_range(
           reassociation, [&](ReassociationIndices indices) -> Attribute {
             return b.getI64ArrayAttr(indices).cast<Attribute>();
           }));
   return b.getArrayAttr(reassociationAttr);
 }
 
-void mlir::linalg::ExpandShapeOp::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, reassociation)));
-  build(b, result, resultType, src, attrs);
-  result.addAttribute(getReassociationAttrName(),
-                      getReassociationIndicesAttribute(b, reassociation));
-}
-
-Value mlir::linalg::ExpandShapeOp::getViewSource() { return src(); }
-
-void mlir::linalg::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, reassociation)));
-  build(b, result, resultType, src, attrs);
-  result.addAttribute(getReassociationAttrName(),
-                      getReassociationIndicesAttribute(b, reassociation));
-}
-
-Value mlir::linalg::CollapseShapeOp::getViewSource() { return src(); }
-
-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)))
-    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;
-  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());
-}
-
-void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                                  MLIRContext *context) {
-  results.add<CollapseReshapeOps<CollapseShapeOp>,
-              CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>>(context);
-}
-
 //===----------------------------------------------------------------------===//
 // TensorReshapeOp
 //===----------------------------------------------------------------------===//
 
 /// Compute the RankedTensorType obtained by applying `reassociation` to `type`.
 static RankedTensorType
 computeTensorReshapeCollapsedType(RankedTensorType type,
                                   ArrayRef<AffineMap> reassociation) {
llvm::interleave(
       types.begin(), types.end(), [&](Type t) { appendMangledType(ss, t); },
       [&]() { ss << "_"; });
   return ss.str();
 }
 
 // TODO: Consider making all this boilerplate easy to autogenerate
 // with Tablegen. This seems a desirable property in the context of
 // OpInterfaces where a Linalg "named" op **isa** LinalgOp.
-OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
-  if (succeeded(foldMemRefCast(*this)))
-    return getResult();
-  return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
-}
-OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
-  if (succeeded(foldMemRefCast(*this)))
-    return getResult();
-  return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);
-}
 OpFoldResult TensorExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<TensorExpandShapeOp, TensorCollapseShapeOp>(*this,
                                                                    operands);
 }
 OpFoldResult TensorCollapseShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<TensorCollapseShapeOp, TensorExpandShapeOp>(*this,
                                                                    operands);
 }

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

 /// Conversion pattern that replaces `linalg.tensor_reshape` with
 /// `linalg.reshape`.
 template <typename TensorReshapeOp,
           typename Adaptor = typename TensorReshapeOp::Adaptor>
 class BufferizeTensorReshapeOp : public OpConversionPattern<TensorReshapeOp> {
 public:
   using OpConversionPattern<TensorReshapeOp>::OpConversionPattern;
   using ReshapeOp = typename std::conditional_t<
-      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value, ExpandShapeOp,
-      CollapseShapeOp>;
+      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value,
+      memref::ExpandShapeOp, memref::CollapseShapeOp>;
 
   LogicalResult
   matchAndRewrite(TensorReshapeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     Adaptor adaptor(operands, op->getAttrDictionary());
     rewriter.replaceOpWithNewOp<ReshapeOp>(op,
                                            this->getTypeConverter()
                                                ->convertType(op.getType())

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

Value maybeExpand(Value result, Type origResultType,
     if (origResultType == result.getType())
       return result;
     if (origResultType.isa<RankedTensorType>()) {
       return rewriter.create<linalg::TensorExpandShapeOp>(
           loc, origResultType, result,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     if (origResultType.isa<MemRefType>()) {
-      return rewriter.create<linalg::ExpandShapeOp>(
+      return rewriter.create<memref::ExpandShapeOp>(
           loc, origResultType, result,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     return nullptr;
   };
 
   // Return the original value if the type is unchanged, or reshape it. Return a
   // nullptr if this is an unsupported type.
   Value maybeCollapse(Value operand, Type newInputOutputType,
                       ArrayAttr reassociationMap, Location loc,
                       PatternRewriter &rewriter) const {
     auto operandType = operand.getType();
     if (operandType == newInputOutputType)
       return operand;
     if (operandType.isa<MemRefType>()) {
-      return rewriter.create<linalg::CollapseShapeOp>(
+      return rewriter.create<memref::CollapseShapeOp>(
           loc, newInputOutputType, operand,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     if (operandType.isa<RankedTensorType>()) {
       return rewriter.create<linalg::TensorCollapseShapeOp>(
           loc, newInputOutputType, operand,
           convertAffineMapArrayToExprs(reassociationMap));
     }

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

for (auto &en : llvm::enumerate(llvm::zip(
                << expectedStride << " instead of " << resultStride
                << " in dim = " << en.index();
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Reassociative reshape ops
+//===----------------------------------------------------------------------===//
+
+SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
+  return getSymbolLessAffineMaps(getReassociationExprs());
+}
+SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {
+  OpBuilder b(this->getContext());
+  return convertReassociationIndicesToExprs(b, getReassociationIndices());
+}
+
+SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() {
+  return getSymbolLessAffineMaps(getReassociationExprs());
+}
+SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() {
+  OpBuilder b(this->getContext());
+  return convertReassociationIndicesToExprs(b, getReassociationIndices());
+}
+
+static void print(OpAsmPrinter &p, ExpandShapeOp op) {
+  ::mlir::printReshapeOp<ExpandShapeOp>(p, op);
+}
+
+static void print(OpAsmPrinter &p, CollapseShapeOp op) {
+  ::mlir::printReshapeOp<CollapseShapeOp>(p, op);
+}
+
+/// Detect whether memref dims [dim, dim + extent) can be reshaped without
+/// copies.
+static bool isReshapableDimBand(unsigned dim, unsigned extent,
+                                ArrayRef<int64_t> sizes,
+                                ArrayRef<AffineExpr> strides) {
+  assert(sizes.size() == strides.size() && "mismatched ranks");
+  // off by 1 indexing to avoid out of bounds
+  //                       V
+  for (auto idx = dim, e = dim + extent; idx + 1 < e; ++idx) {
+    // Only bands of static shapes are reshapable. This is due to the fact that
+    // there is no relation between dynamic sizes and dynamic strides: we do not
+    // have enough information to know whether a "-1" size corresponds to the
+    // proper symbol in the AffineExpr of a stride.
+    if (ShapedType::isDynamic(sizes[dim + 1]))
+      return false;
+    // TODO: Refine this by passing the proper nDims and nSymbols so we can
+    // simplify on the fly and catch more reshapable cases.
+    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
+/// 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);
+  (void)status;
+  assert(succeeded(status) && "expected strided memref");
+
+  SmallVector<int64_t, 4> newSizes;
+  newSizes.reserve(reassociation.size());
+  SmallVector<AffineExpr, 4> newStrides;
+  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");
+  unsigned currentDim = 0;
+  for (AffineMap m : reassociation) {
+    unsigned dim = m.getNumResults();
+    int64_t size = 1;
+    AffineExpr stride = strides[currentDim + dim - 1];
+    if (!isReshapableDimBand(currentDim, dim, sizes, strides)) {
+      size = ShapedType::kDynamicSize;
+      stride = AffineExpr();
+    } else {
+      for (unsigned d = 0; d < dim; ++d)
+        size *= sizes[currentDim + d];
+    }
+    newSizes.push_back(size);
+    newStrides.push_back(stride);
+    currentDim += dim;
+  }
+
+  // Early-exit: if `type` is contiguous, the result must be contiguous.
+  if (canonicalizeStridedLayout(type).getAffineMaps().empty())
+    return MemRefType::Builder(type).setShape(newSizes).setAffineMaps({});
+
+  // 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(strides.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(newSizes).setAffineMaps({layout}));
+}
+
+void ExpandShapeOp::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, reassociation)));
+  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, 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)))
+    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;
+  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());
+}
+
+void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                  MLIRContext *context) {
+  results.add<CollapseReshapeOps<CollapseShapeOp>,
+              CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>>(context);
+}
+OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
+}
+OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);
+}
+
+//===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReshapeOp op) {
   Type operandType = op.source().getType();
   Type resultType = op.result().getType();
 
   Type operandElementType = operandType.cast<ShapedType>().getElementType();

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 
 #include <numeric>
 
 using namespace mlir;
 
-constexpr StringRef mlir::getReassociationAttrName() { return "reassociation"; }
-
 Optional<SmallVector<ReassociationIndices>>
 mlir::getReassociationIndicesForReshape(ShapedType sourceType,
                                         ShapedType targetType) {
   // Make the sourceType greater rank than the targetType. If they are same
   // rank, then its an unsupported reshape op.
   if (sourceType.getRank() == targetType.getRank())
     return llvm::None;
   if (sourceType.getRank() < targetType.getRank())
for (int64_t consumerIndex : consumerIndices) {
       for (int64_t producerIndex : producerReassociations[consumerIndex])
         reassociations.push_back(producerIndex);
     }
     composedIndices.push_back(std::move(reassociations));
   }
   return composedIndices;
 }
 
+SmallVector<SmallVector<AffineExpr, 2>, 2>
+mlir::convertReassociationIndicesToExprs(
+    OpBuilder &b, ArrayRef<ReassociationIndices> reassociationIndices) {
+  SmallVector<SmallVector<AffineExpr, 2>, 2> reassociationMaps;
+  for (const auto &indices : reassociationIndices) {
+    SmallVector<AffineExpr, 2> reassociationMap;
+    reassociationMap.reserve(indices.size());
+    for (int64_t index : indices)
+      reassociationMap.push_back(b.getAffineDimExpr(index));
+    reassociationMaps.push_back(std::move(reassociationMap));
+  }
+  return reassociationMaps;
+}
+
+template <typename AffineExprTy>
+unsigned getMaxPosOfType(ArrayRef<ReassociationExprs> exprArrays) {
+  unsigned pos = 0;
+  for (const auto &exprs : exprArrays) {
+    for (auto expr : exprs) {
+      expr.walk([&pos](AffineExpr e) {
+        if (auto d = e.dyn_cast<AffineExprTy>())
+          pos = std::max(pos, d.getPosition());
+      });
+    }
+  }
+  return pos;
+}
+
+ArrayAttr mlir::getReassociationIndicesAttribute(
+    OpBuilder &b, ArrayRef<ReassociationIndices> reassociation) {
+  SmallVector<Attribute, 4> reassociationAttr =
+      llvm::to_vector<4>(llvm::map_range(
+          reassociation, [&](ReassociationIndices indices) -> Attribute {
+            return b.getI64ArrayAttr(indices).cast<Attribute>();
+          }));
+  return b.getArrayAttr(reassociationAttr);
+}
+
+SmallVector<ReassociationIndices, 2> mlir::convertReassociationMapsToIndices(
+    OpBuilder &b, 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);
+  }
+  return reassociationIndices;
+}
+
+SmallVector<AffineMap, 4>
+mlir::getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation) {
+  unsigned maxDim = getMaxPosOfType<AffineDimExpr>(reassociation);
+  assert(getMaxPosOfType<AffineSymbolExpr>(reassociation) == 0 &&
+         "Expected symbol-less expressions");
+  SmallVector<AffineMap, 4> maps;
+  maps.reserve(reassociation.size());
+  for (const auto &exprs : reassociation) {
+    assert(!exprs.empty());
+    maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
+  }
+  return maps;
+}
 bool mlir::isReassociationValid(ArrayRef<AffineMap> reassociation,
                                 int *invalidIndex) {
   if (reassociation.empty())
     return true;
   unsigned nDims = reassociation[0].getNumDims();
   unsigned nextExpectedDim = 0;
   for (auto it : llvm::enumerate(reassociation)) {
     auto m = it.value();

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

func @get_gv3_memref() {
   // CHECK: llvm.insertvalue %[[DEADBEEFPTR]], {{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   // CHECK: llvm.insertvalue %[[GEP]], {{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   // CHECK: %[[OFFSET:.*]] = llvm.mlir.constant(0 : index) : i64
   // CHECK: llvm.insertvalue %[[OFFSET]], {{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   %0 = memref.get_global @gv3 : memref<f32>
   return
 }
 
+// -----
+
+func @expand_shape_static(%arg0: memref<3x4x5xf32>) -> memref<1x3x4x1x5xf32> {
+  // Reshapes that expand a contiguous tensor with some 1's.
+  %0 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]]
+      : memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+  return %0 : memref<1x3x4x1x5xf32>
+}
+// CHECK-LABEL: func @expand_shape_static
+//       CHECK:    llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(1 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(3 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(4 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(1 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(5 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 4] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(60 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(20 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(5 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(5 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.mlir.constant(1 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 4] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+
+// -----
+
+func @collapse_shape_static(%arg0: memref<1x3x4x1x5xf32>) -> memref<3x4x5xf32> {
+  %0 = memref.collapse_shape %arg0 [[0, 1], [2], [3, 4]] :
+    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+  return %0 : memref<3x4x5xf32>
+}
+// CHECK-LABEL: func @collapse_shape_static
+//       CHECK:    llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(3 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(4 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(5 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(20 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(5 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+//       CHECK:    llvm.mlir.constant(1 : index) : i64
+//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+
+// -----
+
+func @collapse_shape_fold_zero_dim(%arg0 : memref<1x1xf32>) -> memref<f32> {
+  %0 = memref.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>
+  return %0 : memref<f32>
+}
+// CHECK-LABEL: func @collapse_shape_fold_zero_dim
+//       CHECK:   llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+
+// -----
+
+func @expand_shape_zero_dim(%arg0 : memref<f32>) -> memref<1x1xf32> {
+  %0 = memref.expand_shape %arg0 [] : memref<f32> into memref<1x1xf32>
+  return %0 : memref<1x1xf32>
+}
+// CHECK-LABEL: func @expand_shape_zero_dim
+//       CHECK:   llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.mlir.constant(1 : index) : i64
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.mlir.constant(1 : index) : i64
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.mlir.constant(1 : index) : i64
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+//       CHECK:   llvm.mlir.constant(1 : index) : i64
+//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>

mlir/test/Dialect/Linalg/bufferize.mlir

 // CHECK-LABEL: func @bufferize_tensor_collapse_shape(
 // CHECK-SAME:    %[[IN:.*]]: tensor<4x5xf32>
 func @bufferize_tensor_collapse_shape(%arg0: tensor<4x5xf32>) -> tensor<20xf32> {
   %out = linalg.tensor_collapse_shape %arg0 [[0, 1]] :
      tensor<4x5xf32> into tensor<20xf32>
   return %out : tensor<20xf32>
 }
 // CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<4x5xf32>
-// CHECK: %[[RESHAPE:.*]] = linalg.collapse_shape %[[MEMREF]] {{\[}}[0, 1]]
+// CHECK: %[[RESHAPE:.*]] = memref.collapse_shape %[[MEMREF]] {{\[}}[0, 1]]
 // CHECK-SAME: : memref<4x5xf32> into memref<20xf32>
 // CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[RESHAPE]] : memref<20xf32>
 // CHECK: return %[[TENSOR]]
 
 // -----
 
 // CHECK-LABEL:   func @pad_tensor_dynamic_shape(
 // CHECK-SAME:                                   %[[IN:.*]]: tensor<4x?x2x?xf32>,

mlir/test/Dialect/Linalg/canonicalize.mlir

func @collapsing_tensor_reshapes_to_zero_dim(%arg0 : tensor<1x1x1xf32>)
   return %1 : tensor<f32>
 }
 // CHECK-LABEL: collapsing_tensor_reshapes_to_zero
 //       CHECK:   linalg.tensor_collapse_shape %{{.*}} []
 //  CHECK-SAME:     tensor<1x1x1xf32> into tensor<f32>
 
 // -----
 
-func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
-                                             -> memref<f32> {
-  %0 = linalg.collapse_shape %arg0 [[0, 1, 2]]
-      : memref<1x1x1xf32> into memref<1xf32>
-  %1 = linalg.collapse_shape %0 [] : memref<1xf32> into memref<f32>
-  return %1 : memref<f32>
-}
-// CHECK-LABEL: collapsing_memref_reshapes_to_zero
-//       CHECK:   linalg.collapse_shape %{{.*}} []
-//  CHECK-SAME:     memref<1x1x1xf32> into memref<f32>
-
-// -----
-
 func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>) -> tensor<?x6x4x?x5xf32>
 {
   %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
       : tensor<?x?xf32> into tensor<?x4x?xf32>
   %1 = linalg.tensor_expand_shape %0 [[0, 1], [2], [3, 4]]
       : tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>
   return %1 : tensor<?x6x4x?x5xf32>
 }
 // CHECK-LABEL: expanding_tensor_reshapes
 //       CHECK:   linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   linalg.tensor_expand_shape
 
 // -----
 
-func @collapsing_memref_reshapes(%arg0 : memref<?x?x?x?x?xf32>) -> memref<?x?xf32>
-{
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
-      : memref<?x?x?x?x?xf32> into memref<?x?x?xf32>
-  %1 = linalg.collapse_shape %0 [[0, 1], [2]]
-      : memref<?x?x?xf32> into memref<?x?xf32>
-  return %1 : memref<?x?xf32>
-}
-// CHECK-LABEL: collapsing_memref_reshapes
-//       CHECK:   linalg.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
-//   CHECK-NOT:   linalg.collapse_shape
-
-// -----
-
-func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>) -> memref<?x6x4x5x?xf32>
-{
-  %0 = linalg.expand_shape %arg0 [[0, 1], [2]]
-      : memref<?x?xf32> into memref<?x4x?xf32>
-  %1 = linalg.expand_shape %0 [[0, 1], [2], [3, 4]]
-      : memref<?x4x?xf32> into memref<?x6x4x5x?xf32>
-  return %1 : memref<?x6x4x5x?xf32>
-}
-// CHECK-LABEL: expanding_memref_reshapes
-//       CHECK:   linalg.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
-//   CHECK-NOT:   linalg.expand_shape
-
-// -----
-
 func @expanding_tensor_reshapes_to_zero_dim(%arg0 : tensor<f32>)
                                              -> tensor<1x1x1xf32> {
   %0 = linalg.tensor_expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
   %1 = linalg.tensor_expand_shape %0 [[0, 1, 2]]
       : tensor<1xf32> into tensor<1x1x1xf32>
   return %1 : tensor<1x1x1xf32>
 }
 // CHECK-LABEL: expanding_tensor_reshapes_to_zero
 //       CHECK:   linalg.tensor_expand_shape %{{.*}} []
 //  CHECK-SAME:     tensor<f32> into tensor<1x1x1xf32>
 
 // -----
 
-func @expanding_memref_reshapes_to_zero_dim(%arg0 : memref<f32>)
-                                             -> memref<1x1x1xf32> {
-  %0 = linalg.expand_shape %arg0 [] : memref<f32> into memref<1xf32>
-  %1 = linalg.expand_shape %0 [[0, 1, 2]]
-      : memref<1xf32> into memref<1x1x1xf32>
-  return %1 : memref<1x1x1xf32>
-}
-// CHECK-LABEL: expanding_memref_reshapes_to_zero
-//       CHECK:   linalg.expand_shape %{{.*}} []
-//  CHECK-SAME:     memref<f32> into memref<1x1x1xf32>
-
-// -----
-
 func @fold_tensor_reshape(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32>
 {
   %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
       : tensor<12x4xf32> into tensor<3x4x4xf32>
   %1 = linalg.tensor_collapse_shape %0 [[0, 1], [2]]
       : tensor<3x4x4xf32> into tensor<12x4xf32>
   return %1 : tensor<12x4xf32>
 }
%1 = linalg.tensor_collapse_shape %0 [[0, 1], [2]]
       : tensor<?x4x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 // CHECK-LABEL: @fold_tensor_reshape_dynamic
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
 
-func @fold_memref_reshape(%arg0 : memref<12x4xf32>) -> memref<12x4xf32>
-{
-  %0 = linalg.expand_shape %arg0 [[0, 1], [2]]
-      : memref<12x4xf32> into memref<3x4x4xf32>
-  %1 = linalg.collapse_shape %0 [[0, 1], [2]]
-      : memref<3x4x4xf32> into memref<12x4xf32>
-  return %1 : memref<12x4xf32>
-}
-// CHECK-LABEL: @fold_memref_reshape
-//   CHECK-NOT:   linalg.{{.*}}_shape
-
-// -----
-
-func @fold_memref_reshape_dynamic(%arg0 : memref<?x?xf32>) -> memref<?x?xf32>
-{
-  %0 = linalg.expand_shape %arg0 [[0, 1], [2]]
-      : memref<?x?xf32> into memref<?x4x?xf32>
-  %1 = linalg.collapse_shape %0 [[0, 1], [2]]
-      : memref<?x4x?xf32> into memref<?x?xf32>
-  return %1 : memref<?x?xf32>
-}
-// CHECK-LABEL: @fold_memref_reshape_dynamic
-//   CHECK-NOT:   linalg.{{.*}}_shape
-
-// -----
-
 func @reshape_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>) -> tensor<24x5x42x8xf32>
 {
   %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2, 3, 4, 5, 6]]
       : tensor<2x3x4x5x6x7x8xf32> into tensor<40320xf32>
   %1 = linalg.tensor_expand_shape %0 [[0, 1, 2, 3]]
       : tensor<40320xf32> into tensor<24x5x42x8xf32>
   return %1 : tensor<24x5x42x8xf32>
 }

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

outs(%shape : memref<?x1x?x1x?xf32>) {
          linalg.yield %arg3 : f32
        }
   return %shape : memref<?x1x?x1x?xf32>
 }
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
 //   CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
 //   CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-LABEL: func @drop_one_trip_loops
-//       CHECK: linalg.collapse_shape %{{.*}} {{\[}}[0, 1], [2]]
+//       CHECK: memref.collapse_shape %{{.*}} {{\[}}[0, 1], [2]]
 //       CHECK: linalg.generic
 //  CHECK-SAME:   indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]
 //  CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel"]
 
 // -----
 
 #accesses = [
   affine_map<(i, j, k, l, m) -> (i, k, m)>,
linalg.generic #trait
     outs(%arg0 : memref<1x1xf32>) {
        ^bb0(%arg1: f32, %arg2: f32) :
          linalg.yield %arg1 : f32
        }
   return %arg0 : memref<1x1xf32>
 }
 //       CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
 // CHECK-LABEL: func @drop_all_loops
-//       CHECK:   linalg.collapse_shape %{{.*}} []
+//       CHECK:   memref.collapse_shape %{{.*}} []
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
 //  CHECK-SAME:     iterator_types = []
 
 // -----
 
 #map0 = affine_map<(i, j) -> (i, j)>
 #access = [#map0, #map0]
func @leading_dim_1_canonicalization(%arg0: memref<1x5xf32>, %shape: memref<5xf32>) -> memref<5xf32> {
   ^bb0(%arg2: f32, %arg3: f32):     // no predecessors
     linalg.yield %arg2 : f32
   }
   return %shape : memref<5xf32>
 }
 //   CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @leading_dim_1_canonicalization
-//       CHECK:   linalg.collapse_shape %{{.*}} {{\[}}[0, 1]]
+//       CHECK:   memref.collapse_shape %{{.*}} {{\[}}[0, 1]]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[$MAP1]], #[[$MAP1]]]
 //  CHECK-SAME:     iterator_types = ["parallel"]
 
 // -----
 
 #accesses = [
   affine_map<(d0, d1) -> (0, d1)>,
   affine_map<(d0, d1) -> (d0, 0)>,
   affine_map<(d0, d1) -> (d0, d1)>
 ]
 
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel", "parallel"],
   library_call = "some_external_fn"
 }
 
 func @broadcast_test(%arg0 : memref<5xf32>, %arg1 : memref<5xf32>, %shape : memref<5x5xf32>) -> memref<5x5xf32>
 {
-  %0 = linalg.expand_shape %arg0 [[0, 1]] : memref<5xf32> into memref<1x5xf32>
-  %1 = linalg.expand_shape %arg1 [[0, 1]] : memref<5xf32> into memref<5x1xf32>
+  %0 = memref.expand_shape %arg0 [[0, 1]] : memref<5xf32> into memref<1x5xf32>
+  %1 = memref.expand_shape %arg1 [[0, 1]] : memref<5xf32> into memref<5x1xf32>
   linalg.generic #trait
      ins(%0, %1 : memref<1x5xf32>, memref<5x1xf32>)
     outs(%shape : memref<5x5xf32>) {
        ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
          %3 = addf %arg3, %arg4 : f32
          linalg.yield %3 : f32
        }
   return %shape : memref<5x5xf32>
outs(%shape : memref<?x?xf32>) {
         linalg.yield %arg2 : f32
    }
    return %shape : memref<?x?xf32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: func @broadcast_scalar
 //  CHECK-SAME:   %[[ARG0:.*]]: memref<1x1xf32>
-//       CHECK:   %[[A:.*]] = linalg.collapse_shape %[[ARG0]] []
+//       CHECK:   %[[A:.*]] = memref.collapse_shape %[[ARG0]] []
 //  CHECK-SAME:     memref<1x1xf32> into memref<f32>
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
 //  CHECK-SAME:     iterator_types = ["parallel", "parallel"]
 //  CHECK-SAME:     %[[A]]
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d2)>
 func @fold_unit_dim_memref_reshape_op(%arg0 : memref<5xf32>) -> memref<2x5xf32>
 {
   %1 = memref.alloc() : memref<1x2x5xf32>
   linalg.generic {i64, indexing_maps = [#map1, #map0],
     iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%arg0 : memref<5xf32>) outs(%1 : memref<1x2x5xf32>) {
     ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
       linalg.yield %arg1 : f32
     }
-  %3 = linalg.collapse_shape %1 [[0, 1], [2]]
+  %3 = memref.collapse_shape %1 [[0, 1], [2]]
     : memref<1x2x5xf32> into memref<2x5xf32>
   return %3 : memref<2x5xf32>
 }
 // CHECK-LABEL: func @fold_unit_dim_memref_reshape_op
 //       CHECK:   %[[ALLOC:.*]] = memref.alloc() : memref<1x2x5xf32>
-//       CHECK:   %[[OUT:.*]] = linalg.collapse_shape %[[ALLOC]]
+//       CHECK:   %[[OUT:.*]] = memref.collapse_shape %[[ALLOC]]
 //       CHECK:   linalg.generic
 //       CHECK-SAME:   outs(%[[OUT:.*]] :
-//       CHECK:   %[[RESULT:.*]] = linalg.collapse_shape %[[ALLOC]]
+//       CHECK:   %[[RESULT:.*]] = memref.collapse_shape %[[ALLOC]]
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 func @fold_unit_dim_for_init_memref(%input: memref<1x1000xf32>) -> memref<1xf32> {
   %cst = constant 0.0 : f32
   %init = memref.alloc() : memref<1xf32>
   linalg.generic {
       indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
       iterator_types = ["parallel", "reduction"]}
     ins(%input : memref<1x1000xf32>)outs(%init : memref<1xf32>) {
   ^bb0(%arg1: f32, %arg2: f32):
     %1823 = addf %arg1, %arg2 : f32
     linalg.yield %1823 : f32
   }
   return %init : memref<1xf32>
 }
 
 
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>
 
 //       CHECK: func @fold_unit_dim_for_init_memref
 //       CHECK: %[[INIT:.+]] = memref.alloc() : memref<1xf32>
-//       CHECK: %[[INPUT_RESHAPE:.+]] = linalg.collapse_shape %{{.+}} {{\[}}[0, 1]] : memref<1x1000xf32> into memref<1000xf32>
-//       CHECK: %[[INIT_RESHAPE:.+]] = linalg.collapse_shape %[[INIT]] [] : memref<1xf32> into memref<f32>
+//       CHECK: %[[INPUT_RESHAPE:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0, 1]] : memref<1x1000xf32> into memref<1000xf32>
+//       CHECK: %[[INIT_RESHAPE:.+]] = memref.collapse_shape %[[INIT]] [] : memref<1xf32> into memref<f32>
 //       CHECK: linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP2]]]
 //  CHECK-SAME:     iterator_types = ["reduction"]
 //  CHECK-SAME:   ins(%[[INPUT_RESHAPE]] : memref<1000xf32>)
 //  CHECK-SAME:   outs(%[[INIT_RESHAPE]] : memref<f32>)
 //       CHECK: return %[[INIT:.+]] : memref<1xf32>
 
 
 

mlir/test/Dialect/Linalg/invalid.mlir

 
 // -----
 
 // expected-error @+1 {{expected valid keyword}}
 !invalid_type = type !linalg<"?">
 
 // -----
 
-func @reshape(%arg0: memref<f32>) {
-  // expected-error @+1 {{expected non-zero memref ranks}}
-  %0 = linalg.expand_shape %arg0 [[0]] : memref<f32> into memref<f32>
-}
-
-// -----
-
-func @collapse_to_higher_rank(%arg0: memref<f32>) {
-  // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
-  %0 = linalg.collapse_shape %arg0 [[0]] : memref<f32> into memref<1xf32>
-}
-
-// -----
-
-func @expand_to_smaller_rank(%arg0: memref<1xf32>) {
-  // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
-  %0 = linalg.expand_shape %arg0 [[0]] : memref<1xf32> into memref<f32>
-}
-
-// -----
-
-func @reshape(%arg0: memref<?xf32>) {
-  // expected-error @+1 {{expected to collapse or expand dims}}
-  %0 = linalg.collapse_shape %arg0 [[0]] : memref<?xf32> into memref<?xf32>
-}
-
-// -----
-
-func @reshape(%arg0: memref<?x?x?xf32>) {
-  // expected-error @+1 {{expected rank of the collapsed type(2) to be the number of reassociation maps(1)}}
-  %0 = linalg.collapse_shape %arg0 [[0, 1]] :
-    memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
-}
-
-// -----
-
-func @reshape(%arg0: memref<?x?x?xf32>) {
-  // expected-error @+1 {{expected reassociation map #1 to be valid and contiguous}}
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [1, 2]] :
-    memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
-}
-
-// -----
-
-func @reshape(%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)>>'}}
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2]] :
-    memref<?x?x?xf32> into memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>
-}
-
-// -----
-
 func @pooling_rank_mismatch(%arg0: memref<?x?x?xf32>,
                             %arg1: memref<2x3xf32>,
                             %arg2: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected operand rank (2) to match the result rank of indexing_map #1 (3)}}
   linalg.pooling_max(%arg0, %arg1, %arg2) {strides = [2, 1, 2]}:
     memref<?x?x?xf32>, memref<2x3xf32>, memref<?x?x?xf32>
   return
 }

 func @matching_inits(%m: memref<?x?xf32>, %t: tensor<?x?xf32>) {
   // expected-error @+1 {{expected type of operand #2 ('tensor<?x?xf32>') to match type of corresponding result ('tensor<?xf32>')}}
   %res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>)
                       outs(%t : tensor<?x?xf32>)
                         -> tensor<?xf32>
   return
 }
 
-
 // -----
 
 func @init_tensor_err(%arg0 : index, %arg1 : index)
 {
   // expected-error @+1 {{specified type 'tensor<4x?x?x5xf32>' does not match the inferred type 'tensor<4x5x?x?xf32>'}}
   %1 = linalg.init_tensor [4, 5, %arg0, %arg1] : tensor<4x?x?x5xf32>
   return
 }
func @illegal_expanding_reshape_dynamic_tensor
   // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
   %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3, 4]]
       : tensor<?x?x?xf32> into tensor<?x?x?x4x?xf32>
   return %0 : tensor<?x?x?x4x?xf32>
 }
 
 // -----
 
-func @illegal_expanding_reshape_dynamic_memref
-  (%arg0: memref<?x?x?xf32>) -> memref<?x?x?x4x?xf32>
-{
-  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
-  %0 = linalg.expand_shape %arg0 [[0], [1], [2, 3, 4]]
-      : memref<?x?x?xf32> into memref<?x?x?x4x?xf32>
-  return %0 : memref<?x?x?x4x?xf32>
-}
-
-// -----
 
 func @illegal_expanding_reshape_static_tensor
   (%arg0: tensor<2x3x20xf32>) -> tensor<2x3x2x4x5xf32>
 {
   // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
   %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3, 4]]
       : tensor<2x3x20xf32> into tensor<2x3x2x4x5xf32>
   return %0 : tensor<2x3x2x4x5xf32>
 }
 
 // -----
 
 func @illegal_collapsing_reshape_static_tensor
   (%arg0: tensor<2x3x2x4x5xf32>) -> tensor<2x3x20xf32>
 {
   // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
   %0 = linalg.tensor_collapse_shape %arg0 [[0], [1], [2, 3, 4]]
       : tensor<2x3x2x4x5xf32> into tensor<2x3x20xf32>
   return %0 : tensor<2x3x20xf32>
 }
 
 // -----
 
-func @illegal_expanding_reshape_static_memref
-  (%arg0: memref<2x3x20xf32>) -> memref<2x3x2x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
-  %0 = linalg.expand_shape %arg0 [[0], [1], [2, 3, 4]]
-      : memref<2x3x20xf32> into memref<2x3x2x4x5xf32>
-  return %0 : memref<2x3x2x4x5xf32>
-}
-
-// -----
-
-func @illegal_collapsing_reshape_static_memref
-  (%arg0: memref<2x3x2x4x5xf32>) -> memref<2x3x20xf32>
-{
-  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
-  %0 = linalg.collapse_shape %arg0 [[0], [1], [2, 3, 4]]
-      : memref<2x3x2x4x5xf32> into memref<2x3x20xf32>
-  return %0 : memref<2x3x20xf32>
-}
-
-// -----
-
 func @illegal_expanding_reshape_mixed_tensor(%arg0 : tensor<?x?xf32>) -> tensor<?x4x5xf32>
 {
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
   %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
       : tensor<?x?xf32> into tensor<?x4x5xf32>
   return %0 : tensor<?x4x5xf32>
 }
 
func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32>
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
   %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2]]
       : tensor<?x4x5xf32> into tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 
 // -----
 
-func @illegal_expanding_reshape_mixed_memref(%arg0 : memref<?x?xf32>) -> memref<?x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
-  %0 = linalg.expand_shape %arg0 [[0, 1], [2]]
-      : memref<?x?xf32> into memref<?x4x5xf32>
-  return %0 : memref<?x4x5xf32>
-}
-
-// -----
-
-func @illegal_expanding_reshape_mixed_memref_2(%arg0 : memref<?x?xf32>) -> memref<?x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
-  %0 = linalg.expand_shape %arg0 [[0], [1, 2]]
-      : memref<?x?xf32> into memref<?x4x5xf32>
-  return %0 : memref<?x4x5xf32>
-}
-
-// -----
-
-func @illegal_collapsing_reshape_mixed_memref(%arg0 : memref<?x4x5xf32>) -> memref<?x?xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2]]
-      : memref<?x4x5xf32> into memref<?x?xf32>
-  return %0 : memref<?x?xf32>
-}
-
-// -----
-
-func @illegal_collapse_reshape_mixed_memref_2(%arg0 : memref<?x4x5xf32>) -> memref<?x?xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
-  %0 = linalg.collapse_shape %arg0 [[0], [1, 2]]
-      : memref<?x4x5xf32> into memref<?x?xf32>
-  return %0 : memref<?x?xf32>
-}
-
-// -----
-
 func @pad_result_type(%arg0: tensor<?x2x3x4xi32>, %arg1: index, %arg2: i32) -> tensor<?x?x?x8xf32> {
   // expected-error @+1 {{specified type 'tensor<?x?x?x8xf32>' does not match the inferred type 'tensor<?x?x?x9xi32>}}
   %0 = linalg.pad_tensor %arg0 low[1, %arg1, 2, 2] high[1, 2, %arg1, 3] {
   ^bb0(%arg3: index, %arg4: index):  // no predecessors
     linalg.yield %arg2 : i32
   } : tensor<?x2x3x4xi32> to tensor<?x?x?x8xf32>
   return %0 : tensor<?x?x?x8xf32>
 }
#attrs = {
   iterator_types = ["parallel"]
 }
 
 func @invalid_reverse(%A: memref<5xf32>, %B: memref<5xf32>) {
   // expected-error @+1 {{unexpected result less than 0 at expression #0 in}}
   linalg.generic #attrs ins(%A: memref<5xf32>) outs(%B: memref<5xf32>) {
     ^bb0(%a: f32, %b: f32):
     linalg.yield %a : f32
-  } 
+  }
   return
 }

mlir/test/Dialect/Linalg/llvm.mlir

 // RUN: mlir-opt %s -convert-linalg-to-llvm | FileCheck %s
 
 func @range(%arg0: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %R = linalg.range %c0:%arg0:%c1 : !linalg.range
   return
 }
 // CHECK-LABEL: func @range
 //       CHECK:   constant 0 : index
 //       CHECK:   constant 1 : index
 //       CHECK:   llvm.mlir.undef : !llvm.struct<(i64, i64, i64)>
 //       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(i64, i64, i64)>
 //       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(i64, i64, i64)>
 //       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(i64, i64, i64)>
-
-func @expand_shape_static(%arg0: memref<3x4x5xf32>) -> memref<1x3x4x1x5xf32> {
-  // Reshapes that expand a contiguous tensor with some 1's.
-  %0 = linalg.expand_shape %arg0 [[0, 1], [2], [3, 4]]
-      : memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
-  return %0 : memref<1x3x4x1x5xf32>
-}
-// CHECK-LABEL: func @expand_shape_static
-//       CHECK:    llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(1 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(3 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(4 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(1 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(5 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 4] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(60 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(20 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(5 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(5 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.mlir.constant(1 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 4] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-
-func @collapse_shape_static(%arg0: memref<1x3x4x1x5xf32>) -> memref<3x4x5xf32> {
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2], [3, 4]] :
-    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
-  return %0 : memref<3x4x5xf32>
-}
-// CHECK-LABEL: func @collapse_shape_static
-//       CHECK:    llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<5 x i64>, array<5 x i64>)>
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(3 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(4 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(5 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(20 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(5 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:    llvm.mlir.constant(1 : index) : i64
-//       CHECK:    llvm.insertvalue %{{.*}}, %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-
-func @collapse_shape_fold_zero_dim(%arg0 : memref<1x1xf32>) -> memref<f32> {
-  %0 = linalg.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>
-  return %0 : memref<f32>
-}
-// CHECK-LABEL: func @collapse_shape_fold_zero_dim
-//       CHECK:   llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-
-func @expand_shape_zero_dim(%arg0 : memref<f32>) -> memref<1x1xf32> {
-  %0 = linalg.expand_shape %arg0 [] : memref<f32> into memref<1x1xf32>
-  return %0 : memref<1x1xf32>
-}
-// CHECK-LABEL: func @expand_shape_zero_dim
-//       CHECK:   llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.extractvalue %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.mlir.constant(1 : index) : i64
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.mlir.constant(1 : index) : i64
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.mlir.constant(1 : index) : i64
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
-//       CHECK:   llvm.mlir.constant(1 : index) : i64
-//       CHECK:   llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>

mlir/test/Dialect/Linalg/roundtrip.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 // TODO: Re-enable LLVM lowering test.
 //
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // DISABLED: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
 
 // CHECK-DAG: #[[$id_2d:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
 // CHECK-DAG: #[[$id_1d:.*]] = affine_map<(d0, d1, d2) -> (d1)>
 // CHECK-DAG: #[[$permute_0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
 // CHECK-DAG: #[[$permute_1:.*]] = affine_map<(d0, d1, d2) -> (d2, d1, d0)>
 // CHECK-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
-// CHECK-DAG: #[[$strided2DOFF0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>
 // CHECK-DAG: #[[$strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
-// CHECK-DAG: #[[$strided3DOFF0:.*]] = affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2)>
 // CHECK-DAG: #[[$strided3DT:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2 * s1 + s0 + d1 * s2 + d0)>
 // CHECK-DAG: #[[$strided6D:.*]] = affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4 + d4 * s5 + d5)>
 
 func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
                   %pad_value: f32) -> tensor<6x?x?x?xf32> {
   %0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
       linalg.yield %pad_value : f32

 // CHECK-SAME:  outs(%{{.*}}: memref<?xf32, #[[$strided1D]]>)
 // CHECK: linalg.dot
 // CHECK-SAME:   ins(%{{.*}}, %{{.*}}: memref<?xf32, #[[$strided1D]]>,
 // CHECK-SAME:                         memref<?xf32, #[[$strided1D]]>)
 // CHECK-SAME:  outs(%{{.*}}: memref<f32>)
 
 // -----
 
-
 func @fill_view(%arg0: memref<?xf32, offset: ?, strides: [1]>, %arg1: f32) {
   linalg.fill(%arg1, %arg0) : f32, memref<?xf32, offset: ?, strides: [1]>
   return
 }
 // CHECK-LABEL: func @fill_view(
 //       CHECK:  %{{.*}}: memref<?xf32, #[[$strided1D]]>, %{{.*}}: f32) {
 //       CHECK:   linalg.fill(%{{.*}}, %{{.*}}) : f32, memref<?xf32, #[[$strided1D]]>
 

 //       CHECK:  ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32):
 //       CHECK:    %{{.*}} = linalg.index 0 : index
 //       CHECK:    %{{.*}} = linalg.index 1 : index
 //       CHECK:    %{{.*}} = linalg.index 2 : index
 //       CHECK:    linalg.yield %{{.*}} : f32
 
 // -----
 
-func @reshape_static(%arg0: memref<3x4x5xf32>, %arg1: tensor<3x4x5xf32>,
-                     %arg2: tensor<3x?x5xf32>) {
-  // Reshapes that collapse and expand back a contiguous buffer.
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2]] :
-    memref<3x4x5xf32> into memref<12x5xf32>
-  %r0 = linalg.expand_shape %0 [[0, 1], [2]] :
-    memref<12x5xf32> into memref<3x4x5xf32>
-  %1 = linalg.collapse_shape %arg0 [[0], [1, 2]] :
-    memref<3x4x5xf32> into memref<3x20xf32>
-  %r1 = linalg.expand_shape %1 [[0], [1, 2]] :
-    memref<3x20xf32> into memref<3x4x5xf32>
-  %2 = linalg.collapse_shape %arg0 [[0, 1, 2]] :
-    memref<3x4x5xf32> into memref<60xf32>
-  %r2 = linalg.expand_shape %2 [[0, 1, 2]] :
-    memref<60xf32> into memref<3x4x5xf32>
-  // Reshapes that expand and collapse back a contiguous buffer with some 1's.
-  %3 = linalg.expand_shape %arg0 [[0, 1], [2], [3, 4]] :
-    memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
-  %r3 = linalg.collapse_shape %3 [[0, 1], [2], [3, 4]] :
-    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
-  // Reshapes on tensors.
-  %t0 = linalg.tensor_expand_shape %arg1 [[0, 1], [2], [3, 4]] :
-    tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
-  %rt0 = linalg.tensor_collapse_shape %t0 [[0, 1], [2], [3, 4]] :
-    tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
-  %t1 = linalg.tensor_expand_shape %arg2 [[0, 1], [2], [3, 4]] :
-    tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
-  %rt1 = linalg.tensor_collapse_shape %t1 [[0], [1, 2], [3, 4]] :
-    tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
-  return
-}
-// CHECK-LABEL: func @reshape_static
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0], [1, 2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0], [1, 2]]
-//  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1, 2]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1, 2]]
-//  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
-//  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
-//  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
-//
-//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
-//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
-//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
-//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
-
-// -----
-
-func @reshape_dynamic(%arg0: memref<?x?x?xf32>,
-                      %arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
-                      %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>) {
-  %0 = linalg.collapse_shape %arg0 [[0, 1], [2]] :
-    memref<?x?x?xf32> into memref<?x?xf32>
-  %r0 = linalg.expand_shape %0 [[0, 1], [2]] :
-    memref<?x?xf32> into memref<?x4x?xf32>
-  %1 = linalg.collapse_shape %arg1 [[0, 1], [2]] :
-    memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]> into
-    memref<?x?xf32, offset : 0, strides : [?, 1]>
-  %r1 = linalg.expand_shape %1 [[0, 1], [2]] :
-    memref<?x?xf32, offset : 0, strides : [?, 1]> into
-    memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
-  %2 = linalg.collapse_shape %arg2 [[0, 1], [2]] :
-    memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
-    memref<?x?xf32, offset : ?, strides : [?, 1]>
-  %r2 = linalg.expand_shape %2 [[0, 1], [2]] :
-    memref<?x?xf32, offset : ?, strides : [?, 1]> into
-    memref<?x4x?xf32, offset : ?, strides : [?, ?, 1]>
-  return
-}
-// CHECK-LABEL: func @reshape
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32> into memref<?x4x?xf32>
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3DOFF0]]> into memref<?x?xf32, #[[$strided2DOFF0]]>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x4x?xf32, #[[$strided3DOFF0]]>
-//       CHECK:   linalg.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
-//       CHECK:   linalg.expand_shape {{.*}} {{\[}}[0, 1], [2]]
-//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
 
 func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,
                 %ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>)
   -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
 {
   linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>)
                      outs(%c3: memref<?x?x?xf32>)
   linalg.batch_matmul ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>) -> (tensor<f32>, tensor<1x1xf32>)
   return %0, %1 : tensor<f32>, tensor<1x1xf32>
 }
 // CHECK-LABEL: func @tensor_reshape_zero_dim
 //       CHECK:   linalg.tensor_collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
 //       CHECK:   linalg.tensor_expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
 
 // -----
 
-func @memref_reshape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>) -> (memref<f32>, memref<1x1xf32>)
-{
-  %0 = linalg.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>
-  %1 = linalg.expand_shape %0 [] : memref<f32> into memref<1x1xf32>
-  return %0, %1 : memref<f32>, memref<1x1xf32>
-}
-// CHECK-LABEL: func @memref_reshape_zero_dim
-//       CHECK:   linalg.collapse_shape %{{.*}} [] : memref<1x1xf32> into memref<f32>
-//       CHECK:   linalg.expand_shape %{{.*}} [] : memref<f32> into memref<1x1xf32>
-
-// -----
 
 func @init_tensor(%arg0 : index, %arg1 : index)
 {
   %0 = linalg.init_tensor [3, 42] : tensor<3x42xf32>
   %1 = linalg.init_tensor [4, %arg0, %arg1, 5] : tensor<4x?x?x5xf32>
   return
 }
 // CHECK-LABEL: func @init_tensor
func @legal_collapsing_reshape_dynamic_tensor
   return %0 : tensor<?x?x?xf32>
 }
 //      CHECK: func @legal_collapsing_reshape_dynamic_tensor
 //      CHECK:   linalg.tensor_collapse_shape
 // CHECK-SAME:    [0], [1], [2, 3, 4]
 
 // -----
 
-func @legal_collapsing_reshape_dynamic_memref
-  (%arg0: memref<?x?x?x4x?xf32>) -> memref<?x?x?xf32>
-{
-  %0 = linalg.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
-    memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
-  return %0 : memref<?x?x?xf32>
-}
-//      CHECK: func @legal_collapsing_reshape_dynamic_memref
-//      CHECK:   linalg.collapse_shape
-// CHECK-SAME:    [0], [1], [2, 3, 4]
-
-// -----
-
 func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
   %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
   %1 = linalg.fill(%arg2, %0) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 // CHECK: %{{.+}} = linalg.fill(%{{.+}}, %{{.+}}) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
 
 // -----

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK-LABEL: func @allocator
 // CHECK:   %[[alloc:.+]] = memref.alloc
 // CHECK:   memref.store %[[alloc:.+]], %arg0
 func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
   %0 = memref.alloc(%arg1) : memref<?xi32>
   memref.store %0, %arg0[] : memref<memref<?xi32>>
   return 
 }
+
+// -----
+
+func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
+                                             -> memref<f32> {
+  %0 = memref.collapse_shape %arg0 [[0, 1, 2]]
+      : memref<1x1x1xf32> into memref<1xf32>
+  %1 = memref.collapse_shape %0 [] : memref<1xf32> into memref<f32>
+  return %1 : memref<f32>
+}
+// CHECK-LABEL: collapsing_memref_reshapes_to_zero
+//       CHECK:   memref.collapse_shape %{{.*}} []
+//  CHECK-SAME:     memref<1x1x1xf32> into memref<f32>
+
+// -----
+
+func @collapsing_memref_reshapes(%arg0 : memref<?x?x?x?x?xf32>)
+    -> memref<?x?xf32> {
+  %0 = memref.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
+      : memref<?x?x?x?x?xf32> into memref<?x?x?xf32>
+  %1 = memref.collapse_shape %0 [[0, 1], [2]]
+      : memref<?x?x?xf32> into memref<?x?xf32>
+  return %1 : memref<?x?xf32>
+}
+// CHECK-LABEL: collapsing_memref_reshapes
+//       CHECK:   memref.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
+//   CHECK-NOT:   memref.collapse_shape
+
+// -----
+
+func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>)
+    -> memref<?x6x4x5x?xf32> {
+  %0 = memref.expand_shape %arg0 [[0, 1], [2]]
+      : memref<?x?xf32> into memref<?x4x?xf32>
+  %1 = memref.expand_shape %0 [[0, 1], [2], [3, 4]]
+      : memref<?x4x?xf32> into memref<?x6x4x5x?xf32>
+  return %1 : memref<?x6x4x5x?xf32>
+}
+// CHECK-LABEL: expanding_memref_reshapes
+//       CHECK:   memref.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
+//   CHECK-NOT:   memref.expand_shape
+
+// -----
+
+func @expanding_memref_reshapes_to_zero_dim(%arg0 : memref<f32>)
+                                             -> memref<1x1x1xf32> {
+  %0 = memref.expand_shape %arg0 [] : memref<f32> into memref<1xf32>
+  %1 = memref.expand_shape %0 [[0, 1, 2]]
+      : memref<1xf32> into memref<1x1x1xf32>
+  return %1 : memref<1x1x1xf32>
+}
+// CHECK-LABEL: expanding_memref_reshapes_to_zero
+//       CHECK:   memref.expand_shape %{{.*}} []
+//  CHECK-SAME:     memref<f32> into memref<1x1x1xf32>
+
+// -----
+
+func @fold_memref_reshape(%arg0 : memref<12x4xf32>) -> memref<12x4xf32> {
+  %0 = memref.expand_shape %arg0 [[0, 1], [2]]
+      : memref<12x4xf32> into memref<3x4x4xf32>
+  %1 = memref.collapse_shape %0 [[0, 1], [2]]
+      : memref<3x4x4xf32> into memref<12x4xf32>
+  return %1 : memref<12x4xf32>
+}
+// CHECK-LABEL: @fold_memref_reshape
+//   CHECK-NOT:   linalg.{{.*}}_shape
+
+// -----
+
+func @fold_memref_reshape_dynamic(%arg0 : memref<?x?xf32>) -> memref<?x?xf32> {
+  %0 = memref.expand_shape %arg0 [[0, 1], [2]]
+      : memref<?x?xf32> into memref<?x4x?xf32>
+  %1 = memref.collapse_shape %0 [[0, 1], [2]]
+      : memref<?x4x?xf32> into memref<?x?xf32>
+  return %1 : memref<?x?xf32>
+}
+// CHECK-LABEL: @fold_memref_reshape_dynamic
+//   CHECK-NOT:   linalg.{{.*}}_shape

mlir/test/Dialect/MemRef/invalid.mlir

 
 // -----
 
 func @copy_different_eltype(%arg0: memref<2xf32>, %arg1: memref<2xf16>) {
   // expected-error @+1 {{op requires the same element type for all operands}}
   memref.copy %arg0, %arg1 : memref<2xf32> to memref<2xf16>
   return
 }
+
+// -----
+
+func @expand_shape(%arg0: memref<f32>) {
+  // expected-error @+1 {{expected non-zero memref ranks}}
+  %0 = memref.expand_shape %arg0 [[0]] : memref<f32> into memref<f32>
+}
+
+// -----
+
+func @collapse_shape_to_higher_rank(%arg0: memref<f32>) {
+  // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
+  %0 = memref.collapse_shape %arg0 [[0]] : memref<f32> into memref<1xf32>
+}
+
+// -----
+
+func @expand_shape_to_smaller_rank(%arg0: memref<1xf32>) {
+  // expected-error @+1 {{expected the type 'memref<f32>' to have higher rank than the type = 'memref<1xf32>'}}
+  %0 = memref.expand_shape %arg0 [[0]] : memref<1xf32> into memref<f32>
+}
+
+// -----
+
+func @collapse_shape(%arg0: memref<?xf32>) {
+  // expected-error @+1 {{expected to collapse or expand dims}}
+  %0 = memref.collapse_shape %arg0 [[0]] : memref<?xf32> into memref<?xf32>
+}
+
+// -----
+
+func @collapse_shape_mismatch_indices_num(%arg0: memref<?x?x?xf32>) {
+  // expected-error @+1 {{expected rank of the collapsed type(2) to be the number of reassociation maps(1)}}
+  %0 = memref.collapse_shape %arg0 [[0, 1]] :
+    memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
+}
+
+// -----
+
+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)>>'}}
+  %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> {
+  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
+  %0 = memref.expand_shape %arg0 [[0], [1], [2, 3, 4]]
+      : memref<?x?x?xf32> into memref<?x?x?x4x?xf32>
+  return %0 : memref<?x?x?x4x?xf32>
+}
+
+// -----
+
+func @expand_shape_illegal_static_memref
+  (%arg0: memref<2x3x20xf32>) -> memref<2x3x2x4x5xf32> {
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = memref.expand_shape %arg0 [[0], [1], [2, 3, 4]]
+      : memref<2x3x20xf32> into memref<2x3x2x4x5xf32>
+  return %0 : memref<2x3x2x4x5xf32>
+}
+
+// -----
+
+func @collapse_shape_illegal_static_memref
+  (%arg0: memref<2x3x2x4x5xf32>) -> memref<2x3x20xf32> {
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = memref.collapse_shape %arg0 [[0], [1], [2, 3, 4]]
+      : memref<2x3x2x4x5xf32> into memref<2x3x20xf32>
+  return %0 : memref<2x3x20xf32>
+}
+
+// -----
+
+func @expand_shape_illegal_mixed_memref(%arg0 : memref<?x?xf32>)
+    -> memref<?x4x5xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = memref.expand_shape %arg0 [[0, 1], [2]]
+      : memref<?x?xf32> into memref<?x4x5xf32>
+  return %0 : memref<?x4x5xf32>
+}
+
+// -----
+
+func @expand_shape_illegal_mixed_memref_2(%arg0 : memref<?x?xf32>)
+    -> memref<?x4x5xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = memref.expand_shape %arg0 [[0], [1, 2]]
+      : memref<?x?xf32> into memref<?x4x5xf32>
+  return %0 : memref<?x4x5xf32>
+}
+
+// -----
+
+func @collapse_shape_illegal_mixed_memref(%arg0 : memref<?x4x5xf32>)
+    -> memref<?x?xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = memref.collapse_shape %arg0 [[0, 1], [2]]
+      : memref<?x4x5xf32> into memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}
+
+// -----
+
+func @collapse_shape_illegal_mixed_memref_2(%arg0 : memref<?x4x5xf32>)
+    -> memref<?x?xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = memref.collapse_shape %arg0 [[0], [1, 2]]
+      : memref<?x4x5xf32> into memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}

mlir/test/Dialect/MemRef/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
+// CHECK-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+// CHECK-DAG: #[[$strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
+// CHECK-DAG: #[[$strided2DOFF0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>
+// CHECK-DAG: #[[$strided3DOFF0:.*]] = affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2)>
+
 // CHECK-LABEL: test_buffer_cast
 func @test_buffer_cast(%arg0: tensor<?xi64>, %arg1: tensor<*xi64>) -> (memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>) {
   %0 = memref.buffer_cast %arg0 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>
   %1 = memref.buffer_cast %arg1 : memref<*xi64, 1>
   return %0, %1 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>
 }
 
 // CHECK-LABEL: func @memref_reinterpret_cast

 
 // CHECK-LABEL: func @memref_alloca_scope
 func @memref_alloca_scope() {
   memref.alloca_scope {
     memref.alloca_scope.return
   }
   return
 }
+
+func @expand_collapse_shape_static(%arg0: memref<3x4x5xf32>,
+                                   %arg1: tensor<3x4x5xf32>,
+                                   %arg2: tensor<3x?x5xf32>) {
+  // Reshapes that collapse and expand back a contiguous buffer.
+  %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
+    memref<3x4x5xf32> into memref<12x5xf32>
+  %r0 = memref.expand_shape %0 [[0, 1], [2]] :
+    memref<12x5xf32> into memref<3x4x5xf32>
+  %1 = memref.collapse_shape %arg0 [[0], [1, 2]] :
+    memref<3x4x5xf32> into memref<3x20xf32>
+  %r1 = memref.expand_shape %1 [[0], [1, 2]] :
+    memref<3x20xf32> into memref<3x4x5xf32>
+  %2 = memref.collapse_shape %arg0 [[0, 1, 2]] :
+    memref<3x4x5xf32> into memref<60xf32>
+  %r2 = memref.expand_shape %2 [[0, 1, 2]] :
+    memref<60xf32> into memref<3x4x5xf32>
+  // Reshapes that expand and collapse back a contiguous buffer with some 1's.
+  %3 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]] :
+    memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+  %r3 = memref.collapse_shape %3 [[0, 1], [2], [3, 4]] :
+    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+  // Reshapes on tensors.
+  %t0 = linalg.tensor_expand_shape %arg1 [[0, 1], [2], [3, 4]] :
+    tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+  %rt0 = linalg.tensor_collapse_shape %t0 [[0, 1], [2], [3, 4]] :
+    tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+  %t1 = linalg.tensor_expand_shape %arg2 [[0, 1], [2], [3, 4]] :
+    tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+  %rt1 = linalg.tensor_collapse_shape %t1 [[0], [1, 2], [3, 4]] :
+    tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
+  return
+}
+// CHECK-LABEL: func @expand_collapse_shape_static
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0], [1, 2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
+//  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1, 2]]
+//  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
+//  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+//
+//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
+
+
+func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
+         %arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
+         %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>) {
+  %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
+    memref<?x?x?xf32> into memref<?x?xf32>
+  %r0 = memref.expand_shape %0 [[0, 1], [2]] :
+    memref<?x?xf32> into memref<?x4x?xf32>
+  %1 = memref.collapse_shape %arg1 [[0, 1], [2]] :
+    memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]> into
+    memref<?x?xf32, offset : 0, strides : [?, 1]>
+  %r1 = memref.expand_shape %1 [[0, 1], [2]] :
+    memref<?x?xf32, offset : 0, strides : [?, 1]> into
+    memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
+  %2 = memref.collapse_shape %arg2 [[0, 1], [2]] :
+    memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
+    memref<?x?xf32, offset : ?, strides : [?, 1]>
+  %r2 = memref.expand_shape %2 [[0, 1], [2]] :
+    memref<?x?xf32, offset : ?, strides : [?, 1]> into
+    memref<?x4x?xf32, offset : ?, strides : [?, ?, 1]>
+  return
+}
+// CHECK-LABEL: func @expand_collapse_shape_dynamic
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32> into memref<?x4x?xf32>
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3DOFF0]]> into memref<?x?xf32, #[[$strided2DOFF0]]>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2DOFF0]]> into memref<?x4x?xf32, #[[$strided3DOFF0]]>
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
+//       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
+//  CHECK-SAME:     memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
+
+func @expand_collapse_shape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>)
+    -> (memref<f32>, memref<1x1xf32>) {
+  %0 = memref.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>
+  %1 = memref.expand_shape %0 [] : memref<f32> into memref<1x1xf32>
+  return %0, %1 : memref<f32>, memref<1x1xf32>
+}
+// CHECK-LABEL: func @expand_collapse_shape_zero_dim
+//       CHECK:   memref.collapse_shape %{{.*}} [] : memref<1x1xf32> into memref<f32>
+//       CHECK:   memref.expand_shape %{{.*}} [] : memref<f32> into memref<1x1xf32>
+
+func @collapse_shape_to_dynamic
+  (%arg0: memref<?x?x?x4x?xf32>) -> memref<?x?x?xf32> {
+  %0 = memref.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
+    memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
+  return %0 : memref<?x?x?xf32>
+}
+//      CHECK: func @collapse_shape_to_dynamic
+//      CHECK:   memref.collapse_shape
+// CHECK-SAME:    [0], [1], [2, 3, 4]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":AffineToStandard",
         ":Analysis",
         ":ConversionPassIncGen",
         ":IR",
         ":LLVMCommonConversion",
         ":LLVMDialect",
         ":LinalgOps",
         ":LinalgTransforms",
+        ":MemRefToLLVM",
         ":Pass",
         ":SCFDialect",
         ":SCFToStandard",
         ":StandardOps",
         ":Support",
         ":Transforms",
         ":VectorToLLVM",
         ":VectorToSCF",