[mlir][Linalg] Add a linalg.reshape op

[nicolasvasilache]

Summary

This diff adds a new operation to linalg to allow reshaping of an
existing view into a new view in the same buffer at the same offset.

More specifically:
The linalg.reshape op produces a new view 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 continous grouping of dimensions and is
represented with a affine map array attribute. In the future, non-continous
groupings may be allowed (i.e. permutations, reindexings etc).

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.

A reshape may either collapse or expand dimensions, depending on the
relationship between source and target memref ranks. The verification rule
is that the reassociation maps are applied to the memref with the larger
rank to obtain the memref with the smaller rank. In the case of a dimension
expansion, the reassociation maps can be interpreted as inverse maps.

Examples:

mlir
   // Dimension collapse (i, j) -> i' and k -> k'
   %1 = linalg.reshape %0 [(i, j, k) -> (i, j), 
                           (i, j, k) -> (k)] :
     memref<?x?x?xf32, stride_spec> into memref<?x?xf32, stride_spec_2>

mlir
   // Dimension expansion i -> (i', j') and (k) -> (k')
   %1 = linalg.reshape %0 [(i, j, k) -> (i, j), 
                           (i, j, k) -> (k)] :
     memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>

The relevant invalid and roundtripping tests are added.

[merge_guards_bot]

Unit tests: pass. 61176 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

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

109

Do not add more uses of -> or *, this is just giving me more work.

mlir/include/mlir/IR/AffineExpr.h

248

return (!*this || !isa<U>()) ? U(nullptr) : U(expr);

mlir/include/mlir/IR/Attributes.h

1444 ↗

(On Diff #236081)

Seems like it would be better to have something more general on ArrayAttr that can unwrap internal attributes. Preferably by returning a range instead of creating a vector.

mlir/include/mlir/IR/Builders.h

381 ↗

(On Diff #236081)

I don't see a need for this.

mlir/lib/IR/StandardTypes.cpp

15

These paths make no sense.

589

This needs a comment, specifically I don't see why we need to do all of this simplification here.

594

Why by-reference?

746

These names seem way too general for what they are used for. Why are these free functions and not methods of MemRefType?

[kiszk]

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

72

nit: a -> an

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

384

maybe enough -> have enough?

[nicolasvasilache]

Rebase and address comments.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[nicolasvasilache]

Address review comments.

mlir/include/mlir/IR/Attributes.h

1444 ↗

(On Diff #236081)

Sorry I have to punt on this one, I looked at the iterator_facade stuff but this exceeds my immediate C++ skills.

mlir/lib/IR/StandardTypes.cpp

589

This was a bad rebase, should be good now.

594

This was a bad rebase, should be good now.

746

It's unnecessary, killing.

[nicolasvasilache]

Update commit message.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[rriddle]

mlir/include/mlir/IR/StandardTypes.h

12

This should be avoided as much as possible. We should avoid adding more includes to headers, especially common ones.

590

This looks unused.

mlir/lib/IR/StandardTypes.cpp

715

Why do these have to be re-simplified? Seems like this could just be done on construction, what is the benefit of doing this here?

[nicolasvasilache]

mlir/include/mlir/IR/StandardTypes.h

12

ugh .. of course, thanks!

mlir/lib/IR/StandardTypes.cpp

715

I can avoid this and just return the original type for now.
I do not remember the rationale for why simplifyAffineMap cannot be run by default but I remember opposition 1+ year ago.
It may be related to roundtripping and old ways of writing tests.

[nicolasvasilache]

Address review comments.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Thanks! Anything else?

[merge_guards_bot]

Unit tests: pass. 61250 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nicolasvasilache]

clang-format

[nicolasvasilache]

Landed the current version. I will address extra requests, if any, in a followup.
I have a separate diff with proper iterators and ranges, I'll update the AffineMap extraction part once https://reviews.llvm.org/D72310 lands.

[merge_guards_bot]

Unit tests: pass. 61250 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

mlir/include/mlir/Dialect/Linalg/EDSC/Intrinsics.h

 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Intrinsics.h"
 
 namespace mlir {
 namespace edsc {
 namespace intrinsics {
 
 using linalg_fill = OperationBuilder<linalg::FillOp>;
+using linalg_reshape = OperationBuilder<linalg::ReshapeOp>;
 using linalg_yield = OperationBuilder<linalg::YieldOp>;
 
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_EDSC_INTRINSICS_H_

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

[{
       auto rangeType = RangeType::get(builder->getContext());
       build(builder, result, rangeType, min, max, step);
     }]>];
 
   // Fully specified by traits.
   let verifier = ?;
 }
 
+def Linalg_ReshapeOp : Linalg_Op<"reshape", [NoSideEffect]>,
+    Arguments<(ins AnyStridedMemRef:$view, AffineMapArrayAttr:$reassociation)>,
+    Results<(outs AnyStridedMemRef)> {
+  let summary = "linalg.reshape produces a new view into the operand view";
+  let description = [{
+    The `linalg.reshape` op produces a new view 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 continous grouping of dimensions and is
+    represented with an affine map array attribute. In the future, non-continous
+    groupings may be allowed (i.e. permutations, reindexings etc).
+
+    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.
+
+    A reshape may either collapse or expand dimensions, depending on the
+    relationship between source and target memref ranks. The verification rule
+    is that the reassociation maps are applied to the memref with the larger
+    rank to obtain the memref with the smaller rank. In the case of a dimension
+    expansion, the reassociation maps can be interpreted as inverse maps.
+
+    Examples:
+
+    ```mlir
+       // Dimension collapse (i, j) -> i' and k -> k'
+       %1 = linalg.reshape %0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+         memref<?x?x?xf32, stride_spec> into memref<?x?xf32, stride_spec_2>
+    ```
+
+    ```mlir
+       // Dimension expansion i -> (i', j') and (k) -> (k')
+       %1 = linalg.reshape %0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+         memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
+    ```
+  }];
+
+  let builders = [OpBuilder<
+    "Builder *b, OperationState &result, Value view, "
+    "ArrayAttr reassociation, ArrayRef<NamedAttribute> attrs = {}">];
+
+  let extraClassDeclaration = [{
+    static StringRef getReassociationAttrName() { return "reassociation"; }
+    MemRefType getViewType() { return view().getType().cast<MemRefType>(); }
+  }];
+}
+
 def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view,
                    Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {
   let summary = "Produce a rank-reduced `subview` of a base `view`.";
   let description = [{
     The `linalg.slice` op allows defining a subregion of a smaller rank than the
     operand `view` within the underlying buffer.

mlir/include/mlir/IR/AffineExpr.h

public:
   bool operator==(int64_t v) const;
   bool operator!=(int64_t v) const { return !(*this == v); }
   explicit operator bool() const { return expr; }
 
   bool operator!() const { return expr == nullptr; }
 
   template <typename U> bool isa() const;
   template <typename U> U dyn_cast() const;
+  template <typename U>
+  U dyn_cast_or_null() const;
   template <typename U> U cast() const;
 
   MLIRContext *getContext() const;
 
   /// Return the classification for this type.
   AffineExprKind getKind() const;
 
   void print(raw_ostream &os) const;

 /// format [dims, symbols, locals, constant term].
 AffineExpr toAffineExpr(ArrayRef<int64_t> eq, unsigned numDims,
                         unsigned numSymbols, ArrayRef<AffineExpr> localExprs,
                         MLIRContext *context);
 
 raw_ostream &operator<<(raw_ostream &os, AffineExpr &expr);
 
 template <typename U> bool AffineExpr::isa() const {
-  if (std::is_same<U, AffineBinaryOpExpr>::value) {
+  if (std::is_same<U, AffineBinaryOpExpr>::value)
     return getKind() <= AffineExprKind::LAST_AFFINE_BINARY_OP;
-  }
-  if (std::is_same<U, AffineDimExpr>::value) {
+  if (std::is_same<U, AffineDimExpr>::value)
     return getKind() == AffineExprKind::DimId;
-  }
-  if (std::is_same<U, AffineSymbolExpr>::value) {
+  if (std::is_same<U, AffineSymbolExpr>::value)
     return getKind() == AffineExprKind::SymbolId;
-  }
-  if (std::is_same<U, AffineConstantExpr>::value) {
+  if (std::is_same<U, AffineConstantExpr>::value)
     return getKind() == AffineExprKind::Constant;
-  }
-}
 template <typename U> U AffineExpr::dyn_cast() const {
-  if (isa<U>()) {
+  if (isa<U>())
     return U(expr);
-  }
   return U(nullptr);
 }
+template <typename U>
+U AffineExpr::dyn_cast_or_null() const {
+  return (!*this || !isa<U>()) ? U(nullptr) : U(expr);
+}
 template <typename U> U AffineExpr::cast() const {
   assert(isa<U>());
   return U(expr);
 }
 
 /// Simplify an affine expression by flattening and some amount of
 /// simple analysis. This has complexity linear in the number of nodes in
 /// 'expr'. Returns the simplified expression, which is the same as the input

mlir/include/mlir/IR/StandardTypes.h

 //===- StandardTypes.h - MLIR Standard Type Classes -------------*- C++ -*-===//
 //
 // Part of the MLIR 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_IR_STANDARDTYPES_H
 #define MLIR_IR_STANDARDTYPES_H
 
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Types.h"
 
 namespace llvm {
 struct fltSemantics;
 } // namespace llvm
 
 namespace mlir {
 class AffineMap;
return type.getKind() == StandardTypes::Vector ||
            type.getKind() == StandardTypes::RankedTensor ||
            type.getKind() == StandardTypes::UnrankedTensor ||
            type.getKind() == StandardTypes::UnrankedMemRef ||
            type.getKind() == StandardTypes::MemRef;
   }
 
   /// Whether the given dimension size indicates a dynamic dimension.
   static constexpr bool isDynamic(int64_t dSize) { return dSize < 0; }
+  static constexpr bool isDynamicStrideOrOffset(int64_t dStrideOrOffset) {
+    return dStrideOrOffset == kDynamicStrideOrOffset;
+  }
 };
 
 /// Vector types represent multi-dimensional SIMD vectors, and have a fixed
 /// known constant shape with one or more dimension.
 class VectorType
     : public Type::TypeBase<VectorType, ShapedType, detail::VectorTypeStorage> {
 public:
   using Base::Base;

 /// If a simple strided form cannot be extracted from the composition of the
 /// layout map, returns llvm::None.
 ///
 /// The convention is that the strides for dimensions d0, .. dn appear in
 /// order to make indexing intuitive into the result.
 LogicalResult getStridesAndOffset(MemRefType t,
                                   SmallVectorImpl<int64_t> &strides,
                                   int64_t &offset);
+LogicalResult getStridesAndOffset(MemRefType t,
+                                  SmallVectorImpl<AffineExpr> &strides,
+                                  AffineExpr &offset);
 
 /// Given a list of strides (in which MemRefType::getDynamicStrideOrOffset()
 /// represents a dynamic value), return the single result AffineMap which
 /// represents the linearized strided layout map. Dimensions correspond to the
 /// offset followed by the strides in order. Symbols are inserted for each
 /// dynamic dimension in order. A stride cannot take value `0`.
 ///
 /// Examples:
 /// =========
 ///
 ///   1. For offset: 0 strides: ?, ?, 1 return
 ///         (i, j, k)[M, N]->(M * i + N * j + k)
 ///
 ///   2. For offset: 3 strides: 32, ?, 16 return
 ///         (i, j, k)[M]->(3 + 32 * i + M * j + 16 * k)
 ///
 ///   3. For offset: ? strides: ?, ?, ? return
 ///         (i, j, k)[off, M, N, P]->(off + M * i + N * j + P * k)
 AffineMap makeStridedLinearLayoutMap(ArrayRef<int64_t> strides, int64_t offset,
                                      MLIRContext *context);
 
+/// Return a version of `t` with identity layout if it can be determined
+/// statically that the layout is the canonical contiguous strided layout.
+/// Otherwise pass `t`'s layout into `simplifyAffineMap` and return a copy of
+/// `t` with simplifed layout.
+MemRefType canonicalizeStridedLayout(MemRefType t);
+
+/// Return true if the layout for `t` is compatible with strided semantics.
 bool isStrided(MemRefType t);
 
+/// Return true if the layout for `t` is compatible with a contiguous (empty)
+/// layout.
+bool isContiguous(MemRefType t);
+
 } // end namespace mlir
 
 #endif // MLIR_IR_STANDARDTYPES_H

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

 //===- LinalgOps.cpp - Implementation of the linalg operations ------------===//
 //
 // Part of the MLIR 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 file implements a the Linalg operations.
+// This file implements the Linalg operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/Support/Functional.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
return failure(parser.parseOperand(rangeInfo[0]) || parser.parseColon() ||
                  parser.parseOperand(rangeInfo[2]) ||
                  parser.parseOptionalAttrDict(result.attributes) ||
                  parser.parseColonType(type) ||
                  parser.resolveOperands(rangeInfo, indexTy, result.operands) ||
                  parser.addTypeToList(type, result.types));
 }
 
 //===----------------------------------------------------------------------===//
+// ReshapeOp
+//===----------------------------------------------------------------------===//
+
+/// Return true if the reassociation specification is valid, false otherwise.
+/// When false, the `invalidIndex` integer pointer is optionally filled with the
+/// index of the offending reassociation map.
+static bool isReassociationValid(ArrayRef<AffineMap> reassociation,
+                                 int *invalidIndex = nullptr) {
+  if (reassociation.empty())
+    return true;
+  unsigned nDims = reassociation[0].getNumDims();
+  unsigned nextExpectedDim = 0;
+  for (auto it : llvm::enumerate(reassociation)) {
+    auto m = it.value();
+    if (m.getNumDims() != nDims || m.getNumSymbols() != 0) {
+      if (invalidIndex)
+        *invalidIndex = it.index();
+      return false;
+    }
+    for (auto e : m.getResults()) {
+      auto d = e.dyn_cast<AffineDimExpr>();
+      if (!d || d.getPosition() != nextExpectedDim++) {
+        if (invalidIndex)
+          *invalidIndex = it.index();
+        return false;
+      }
+    }
+  }
+  if (nextExpectedDim != nDims) {
+    if (invalidIndex)
+      *invalidIndex = reassociation.size() - 1;
+    return false;
+  }
+  return true;
+}
+
+/// 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(ntv) 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::get(newSizes, type.getElementType(), {});
+
+  // 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::get(newSizes, type.getElementType(), {layout}));
+}
+
+/// Helper functions assert Attribute of the proper type in attr and returns the
+/// corresponding vector.
+/// TODO(rridle,ntv) this should be evolved into a generic
+/// `getRangeOfType<AffineMap>(ArrayAttr attrs)` that does not copy.
+static SmallVector<AffineMap, 4> getAffineMaps(ArrayAttr attrs) {
+  return functional::map(
+      [](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }, attrs);
+}
+
+void mlir::linalg::ReshapeOp::build(Builder *b, OperationState &result,
+                                    Value view, ArrayAttr reassociation,
+                                    ArrayRef<NamedAttribute> attrs) {
+  auto maps = getAffineMaps(reassociation);
+  auto memRefType = view.getType().cast<MemRefType>();
+  auto resultType = computeReshapeCollapsedType(memRefType, maps);
+  build(b, result, resultType, view, attrs);
+  result.addAttribute(ReshapeOp::getReassociationAttrName(), reassociation);
+}
+
+static void print(OpAsmPrinter &p, ReshapeOp op) {
+  p << op.getOperationName() << " " << op.view() << " " << op.reassociation();
+  p.printOptionalAttrDict(op.getAttrs(),
+                          {ReshapeOp::getReassociationAttrName()});
+  p << " : " << op.getViewType() << " into " << op.getResult().getType();
+}
+
+static ParseResult parseReshapeOp(OpAsmParser &parser, OperationState &result) {
+  OpAsmParser::OperandType view;
+  ArrayAttr reassociation;
+  MemRefType type, resultType;
+  return failure(parser.parseOperand(view) ||
+                 parser.parseAttribute(reassociation,
+                                       ReshapeOp::getReassociationAttrName(),
+                                       result.attributes) ||
+                 parser.parseOptionalAttrDict(result.attributes) ||
+                 parser.parseColonType(type) ||
+                 parser.parseKeywordType("into", resultType) ||
+                 parser.resolveOperand(view, type, result.operands) ||
+                 parser.addTypeToList(resultType, result.types));
+}
+
+static LogicalResult verify(ReshapeOp op) {
+  MemRefType expandedType = op.getViewType();
+  MemRefType collapsedType = op.getResult().getType().cast<MemRefType>();
+  unsigned expandedRank = expandedType.getRank();
+  unsigned collapsedRank = collapsedType.getRank();
+  bool isCollapse = expandedRank > collapsedRank;
+  if (!isCollapse) {
+    std::swap(expandedRank, collapsedRank);
+    std::swap(expandedType, collapsedType);
+  }
+  if (expandedRank == 0 || collapsedRank == 0)
+    return op.emitOpError("expected non-zero memref ranks");
+  if (expandedRank == collapsedRank)
+    return op.emitOpError("expected to collapse or expand dims");
+
+  if (collapsedRank != op.reassociation().size())
+    return op.emitOpError("expected rank of the collapsed view(")
+           << collapsedRank << ") to be the number of reassociation maps("
+           << op.reassociation().size() << ")";
+  auto maps = getAffineMaps(op.reassociation());
+  for (auto it : llvm::enumerate(maps))
+    if (it.value().getNumDims() != expandedRank)
+      return op.emitOpError("expected reassociation map #")
+             << it.index() << " of same rank as expanded memref("
+             << expandedRank << "), but got " << it.value().getNumDims();
+  int invalidIdx = 0;
+  if (!isReassociationValid(maps, &invalidIdx))
+    return op.emitOpError("expected reassociation map #")
+           << invalidIdx << " to be valid and contiguous";
+  MemRefType expectedType = computeReshapeCollapsedType(expandedType, maps);
+  if (collapsedType != expectedType)
+    return op.emitOpError("expected collapsed type to be ")
+           << expectedType << ", but got " << collapsedType;
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // SliceOp
 //===----------------------------------------------------------------------===//
 void mlir::linalg::SliceOp::build(Builder *b, OperationState &result,
                                   Value base, ValueRange indexings) {
   result.addOperands(base);
   result.addOperands(indexings);
 
   auto memRefType = base->getType().cast<MemRefType>();

mlir/lib/IR/StandardTypes.cpp

 //===- StandardTypes.cpp - MLIR Standard Type Classes ---------------------===//
 //
 // Part of the MLIR 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/IR/StandardTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//
if (bin.getKind() == AffineExprKind::Add) {
     auto res2 =
         extractStrides(bin.getRHS(), multiplicativeFactor, strides, offset);
     return success(succeeded(res1) && succeeded(res2));
   }
 
   llvm_unreachable("unexpected binary operation");
 }
 
-static LogicalResult getStridesAndOffset(MemRefType t,
-                                         SmallVectorImpl<AffineExpr> &strides,
-                                         AffineExpr &offset) {
+LogicalResult mlir::getStridesAndOffset(MemRefType t,
   auto affineMaps = t.getAffineMaps();
   // For now strides are only computed on a single affine map with a single
   // result (i.e. the closed subset of linearization maps that are compatible
   // with striding semantics).
   // TODO(ntv): support more forms on a per-need basis.
   if (affineMaps.size() > 1)
     return failure();
   if (affineMaps.size() == 1 && affineMaps[0].getNumResults() != 1)
LogicalResult mlir::getStridesAndOffset(MemRefType t,
   if (llvm::any_of(strides, [](AffineExpr e) {
         return e == getAffineConstantExpr(0, e.getContext());
       })) {
     offset = AffineExpr();
     strides.clear();
     return failure();
   }
 
   return success();
 }
 
 LogicalResult mlir::getStridesAndOffset(MemRefType t,
                                         SmallVectorImpl<int64_t> &strides,
                                         int64_t &offset) {
   AffineExpr offsetExpr;
   SmallVector<AffineExpr, 4> strideExprs;
   if (failed(::getStridesAndOffset(t, strideExprs, offsetExpr)))
     return failure();
   if (auto cst = offsetExpr.dyn_cast<AffineConstantExpr>())
     offset = cst.getValue();
   else
     offset = ShapedType::kDynamicStrideOrOffset;
else
       // Dynamic case, new symbol for each new stride.
       mult = getAffineSymbolExpr(nSymbols++, context);
     expr = expr + d * mult;
   }
 
   return AffineMap::get(strides.size(), nSymbols, expr);
 }
 
+/// Return a version of `t` with identity layout if it can be determined
+/// statically that the layout is the canonical contiguous strided layout.
+/// Otherwise pass `t`'s layout into `simplifyAffineMap` and return a copy of
+/// `t` with simplifed layout.
+MemRefType mlir::canonicalizeStridedLayout(MemRefType t) {
+  auto affineMaps = t.getAffineMaps();
+  // Already in canonical form.
+  if (affineMaps.empty())
+    return t;
+
+  // Can't reduce to canonical identity form, return in canonical form.
+  if (affineMaps.size() > 1 || affineMaps[0].getNumResults() > 1) {
+    auto maps = llvm::to_vector<4>(affineMaps);
+    for (auto &m : maps)
+      m = simplifyAffineMap(m);
+    return MemRefType::get(t.getShape(), t.getElementType(), maps);
+  }
+
+  // If the canonical strided layout for the sizes of `t` is equal to the
+  // simplified layout of `t` we can just return an empty layout. Otherwise,
+  // just simplify the existing layout.
+  AffineExpr expr =
+      makeCanonicalStridedLayoutExpr(t.getShape(), t.getContext());
+  auto m = affineMaps[0];
+  auto simplifiedLayoutExpr =
+      simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
+  if (expr != simplifiedLayoutExpr)
+    return MemRefType::get(t.getShape(), t.getElementType(),
+                           {AffineMap::get(m.getNumDims(), m.getNumSymbols(),
+                                           {simplifiedLayoutExpr})});
+
+  return MemRefType::get(t.getShape(), t.getElementType(), {});
+}
+
+/// Return true if the layout for `t` is compatible with strided semantics.
 bool mlir::isStrided(MemRefType t) {
   int64_t offset;
   SmallVector<int64_t, 4> stridesAndOffset;
   auto res = getStridesAndOffset(t, stridesAndOffset, offset);
   return succeeded(res);
 }

mlir/test/Dialect/Linalg/invalid.mlir

 
 // expected-error @+1 {{unknown Linalg type}}
 !invalid_type = type !linalg.unknown
 
 // -----
 
 // 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.reshape %arg0 [()->(0)] : memref<f32> into memref<f32>
+}
+
+// -----
+
+func @reshape(%arg0: memref<?xf32>) {
+  // expected-error @+1 {{expected to collapse or expand dims}}
+  %0 = linalg.reshape %arg0 [(i)->(i)] : memref<?xf32> into memref<?xf32>
+}
+
+// -----
+
+func @reshape(%arg0: memref<?x?x?xf32>) {
+  // expected-error @+1 {{expected rank of the collapsed view(2) to be the number of reassociation maps(1)}}
+  %0 = linalg.reshape %arg0 [(i, j, k) -> (i, j)] :
+    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 #0 of same rank as expanded memref(3), but got 1}}
+  %0 = linalg.reshape %arg0 [(i) -> (i), (i, j, k) -> (k)] :
+    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.reshape %arg0 [(i, j, k) -> (i, j), (i, j, k) -> (k, j)] :
+    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, (d0, d1)[s0] -> (d0 * s0 + d1)>'}}
+  %0 = linalg.reshape %arg0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+    memref<?x?x?xf32> into memref<?x?xf32, (d0, d1)[s0] -> (d0 * s0 + d1)>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 // TODO(pifon): Re-enable LLVM lowering test after IndexedGenericOp is lowered.
 //
 // 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: #[[strided1D:.*]] = (d0)[s0] -> (d0 + s0)
 // CHECK-DAG: #[[strided2D:.*]] = (d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)
+// CHECK-DAG: #[[strided2DOFF0:.*]] = (d0, d1)[s0] -> (d0 * s0 + d1)
 // CHECK-DAG: #[[strided3D:.*]] = (d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)
+// CHECK-DAG: #[[strided3DOFF0:.*]] = (d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2)
 // CHECK-DAG: #[[strided6D:.*]] = (d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4 + d4 * s5 + d5)
 
 // CHECK-DAG: #[[map0:.*]] = (d0, d1, d2) -> (d0, d2, d1)
 // CHECK-DAG: #[[map1:.*]] = (d0, d1, d2) -> (d2, d1, d0)
 
+// CHECK-DAG: #[[reshapeD01:.*]] = (d0, d1, d2) -> (d0, d1)
+// CHECK-DAG: #[[reshapeD2:.*]] = (d0, d1, d2) -> (d2)
+// CHECK-DAG: #[[reshapeD0:.*]] = (d0, d1, d2) -> (d0)
+// CHECK-DAG: #[[reshapeD12:.*]] = (d0, d1, d2) -> (d1, d2)
+// CHECK-DAG: #[[reshapeD012:.*]] = (d0, d1, d2) -> (d0, d1, d2)
+// CHECK-DAG: #[[reshape5D01:.*]] = (d0, d1, d2, d3, d4) -> (d0, d1)
+// CHECK-DAG: #[[reshape5D2:.*]] = (d0, d1, d2, d3, d4) -> (d2)
+// CHECK-DAG: #[[reshape5D34:.*]] = (d0, d1, d2, d3, d4) -> (d3, d4)
+
 func @range(%arg0: index, %arg1: index, %arg2: index) {
   %0 = linalg.range %arg0:%arg1:%arg2 : !linalg.range
   return
 }
 // CHECK-LABEL: func @range(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 //  CHECK-NEXT:  linalg.range %{{.*}}:%{{.*}}:%{{.*}} : !linalg.range
 
 func @views(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) {
func @generic_region(%arg0: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>, %arg1: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>) {
   } {foo = 1}: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>, memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>
   return
 }
 // CHECK-LABEL: func @generic_region
 //       CHECK:   linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"], library_call = "some_external_function_name_2"} %{{.*}}, %{{.*}} {
 //       CHECK:    ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32):    // no predecessors
 //       CHECK:      linalg.yield %{{.*}} : f32
 //       CHECK:    } {foo = 1 : i64}: memref<?x?xvector<3x4xi4>, #[[strided2D]]>, memref<?x?x?xf32, #[[strided3D]]>
-
 func @indexed_generic(%arg0: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>,
                       %arg1: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>) {
   linalg.indexed_generic #trait2 %arg0, %arg1 {
   ^bb(%i: index, %j: index, %k: index, %a: vector<3x4xi4>, %b: f32) :
       linalg.yield %b : f32
   } {foo = 1}: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>, memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>
   return
 }
 // CHECK-LABEL: func @indexed_generic
 //       CHECK:   linalg.indexed_generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#{{.*}}, #{{.*}}],  iterator_types = ["parallel", "parallel", "parallel"], library_call = "some_external_function_name_2"} %{{.*}}, %{{.*}} {
 //       CHECK:    ^{{.*}}(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: vector<3x4xi4>, %{{.*}}: f32):
 //       CHECK:      linalg.yield %{{.*}} : f32
 //       CHECK:    } {foo = 1 : i64}: memref<?x?xvector<3x4xi4>, #[[strided2D]]>, memref<?x?x?xf32, #[[strided3D]]>
+
+func @reshape_static(%arg0: memref<3x4x5xf32>) {
+  // Reshapes that collapse and expand back a contiguous tensor.
+  %0 = linalg.reshape %arg0 [(i, j, k) -> (i, j),
+                             (i, j, k) -> (k)] :
+    memref<3x4x5xf32> into memref<12x5xf32>
+  %r0 = linalg.reshape %0 [(i, j, k) -> (i, j),
+                           (i, j, k) -> (k)] :
+    memref<12x5xf32> into memref<3x4x5xf32>
+  %1 = linalg.reshape %arg0 [(i, j, k) -> (i),
+                             (i, j, k) -> (j, k)] :
+    memref<3x4x5xf32> into memref<3x20xf32>
+  %r1 = linalg.reshape %1 [(i, j, k) -> (i),
+                           (i, j, k) -> (j, k)] :
+    memref<3x20xf32> into memref<3x4x5xf32>
+  %2 = linalg.reshape %arg0 [(i, j, k) -> (i, j, k)] :
+    memref<3x4x5xf32> into memref<60xf32>
+  %r2 = linalg.reshape %2 [(i, j, k) -> (i, j, k)] :
+    memref<60xf32> into memref<3x4x5xf32>
+  // Reshapes that expand and collapse back a contiguous tensor with some 1's.
+  %3 = linalg.reshape %arg0 [(i, j, k, l, m) -> (i, j),
+                             (i, j, k, l, m) -> (k),
+                             (i, j, k, l, m) -> (l, m)] :
+    memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+  %r3 = linalg.reshape %3 [(i, j, k, l, m) -> (i, j),
+                           (i, j, k, l, m) -> (k),
+                           (i, j, k, l, m) -> (l, m)] :
+    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+  return
+}
+// CHECK-LABEL: func @reshape_static
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD0]], #[[reshapeD12]]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD0]], #[[reshapeD12]]]
+//  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD012]]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD012]]]
+//  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshape5D01]], #[[reshape5D2]], #[[reshape5D34]]]
+//  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshape5D01]], #[[reshape5D2]], #[[reshape5D34]]]
+//  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+
+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.reshape %arg0 [(i, j, k) -> (i, j),
+                             (i, j, k) -> (k)] :
+    memref<?x?x?xf32> into memref<?x?xf32>
+  %r0 = linalg.reshape %0 [(i, j, k) -> (i, j),
+                           (i, j, k) -> (k)] :
+    memref<?x?xf32> into memref<?x?x?xf32>
+  %1 = linalg.reshape %arg1 [(i, j, k) -> (i, j),
+                             (i, j, k) -> (k)] :
+    memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]> into
+    memref<?x?xf32, offset : 0, strides : [?, 1]>
+  %r1 = linalg.reshape %1 [(i, j, k) -> (i, j),
+                           (i, j, k) -> (k)] :
+    memref<?x?xf32, offset : 0, strides : [?, 1]> into
+    memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>
+  %2 = linalg.reshape %arg2 [(i, j, k) -> (i, j),
+                             (i, j, k) -> (k)] :
+    memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
+    memref<?x?xf32, offset : ?, strides : [?, 1]>
+  %r2 = linalg.reshape %2 [(i, j, k) -> (i, j),
+                           (i, j, k) -> (k)] :
+    memref<?x?xf32, offset : ?, strides : [?, 1]> into
+    memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>
+  return
+}
+// CHECK-LABEL: func @reshape
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?xf32> into memref<?x?x?xf32>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[strided3DOFF0]]> into memref<?x?xf32, #[[strided2DOFF0]]>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?xf32, #[[strided2DOFF0]]> into memref<?x?x?xf32, #[[strided3DOFF0]]>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?x?xf32, #[[strided3D]]> into memref<?x?xf32, #[[strided2D]]>
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
+//  CHECK-SAME:     memref<?x?xf32, #[[strided2D]]> into memref<?x?x?xf32, #[[strided3D]]>