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

[nicolasvasilache]

Summary

This diff adds support for linalg.reshape which takes a view and collapses
its dimensions according to a reassociation list of maps to return a new view.

This is a pure metadata operation which does not allocate or copy.
In the general case, such a reshape may require allocating and moving data, these
cases are currently undefined behavior.

This is captured by the op documentation:

For now, it is assumed that all reassociation occur on contiguous
dimensions or that the reshape op will be folded into its consumers.
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.

Relevant invalid and roundtrip test are added.

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 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

[rriddle]

mlir/include/mlir/IR/Builders.h

380

I don't see why this should be here. This seems unrelated to Builders.

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

351

Duplicated assert.

354

Cache the end iterator.

367

This should have a comment.

383

nit: use AffineMap instead of auto

406

Do not use -> or * on Value unless you really need to. They are going away and it just adds extra work for me. I see many other places in this revision that also use them.

[nicolasvasilache]

fix unused symbol in dbg mode.

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 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

[nicolasvasilache]

This needs more work, I have a better way.

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 strided memref";
+  let description = [{
+    The `linalg.reshape` op produces a strided memref whose sizes are a
+    reassociation of the original `view`. Depending on whether or not the
+    reassociated dimensions are contiguous, the resulting memref may require
+    explicit alloc and copies.
+
+    A reassociation is defined as a contiguous grouping of dimensions and is
+    represented with a affine map array attribute. In the future, non-contiguous
+    groupings may be allowed (i.e. permutations, reindexings etc).
+
+    For now, it is assumed that all reassociation occur on contiguous
+    dimensions or that the reshape op will be folded into its consumers.
+    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.
+
+    Example:
+
+    ```mlir
+       %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>
+    ```
+  }];
+
+  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/Builders.h

 //===- Builders.h - Helpers for constructing MLIR 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_BUILDERS_H
 #define MLIR_IR_BUILDERS_H
 
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 
 class AffineExpr;
 class BlockAndValueMapping;
 class ModuleOp;
 class UnknownLoc;
 class FileLineColLoc;
template <typename OpT> OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
 private:
   Block *block = nullptr;
   Block::iterator insertPoint;
 };
 
+/// Helper functions assert Attribute of the proper type in attr and returns the
+/// corresponding vector.
+SmallVector<AffineMap, 4> getAffineMaps(ArrayAttr attrs);
+
 } // namespace mlir
 
 #endif

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/Types.h"
 
 namespace llvm {
 struct fltSemantics;
 } // namespace llvm
 
 namespace mlir {
 class AffineMap;
 class FloatType;
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;

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.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
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<int64_t> strides) {
+  assert(sizes.size() == strides.size() && "mismatched ranks");
+  assert(sizes.size() == strides.size() && "mismatched ranks");
+  // off by 1 indexing to avoid out of bounds
+  //                       V
+  for (auto idx = dim; idx + 1 < dim + extent; ++idx) {
+    // All values involved must be static, symbolic dims would require more
+    // information in the type.
+    if (ShapedType::isDynamic(sizes[dim + 1]) ||
+        ShapedType::isDynamicStrideOrOffset(strides[dim]) ||
+        ShapedType::isDynamicStrideOrOffset(strides[dim + 1]))
+      return false;
+    if (strides[idx] != strides[idx + 1] * sizes[idx + 1])
+      return false;
+  }
+  return true;
+}
+
+static MemRefType computeReshapeResultType(MemRefType type,
+                                           ArrayRef<AffineMap> reassociation) {
+  auto sizes = type.getShape();
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto status = getStridesAndOffset(type, strides, offset);
+  assert(succeeded(status) && "expected strided memref");
+
+  SmallVector<int64_t, 4> newSizes, newStrides;
+  newSizes.reserve(reassociation.size());
+  newStrides.reserve(reassociation.size());
+  unsigned currentDim = 0;
+
+  // Use the fact that reassociation is valid to simplify the logic: only use
+  // each map's rank.
+  assert(isReassociationValid(reassociation) && "invalid reassociation");
+  for (auto m : reassociation) {
+    unsigned dim = m.getNumResults();
+    int64_t size = 1, stride = strides[currentDim];
+    if (!isReshapableDimBand(currentDim, dim, sizes, strides)) {
+      size = ShapedType::kDynamicSize;
+      stride = ShapedType::kDynamicStrideOrOffset;
+    } else {
+      for (unsigned d = 0; d < dim; ++d)
+        size *= sizes[currentDim + d];
+    }
+    newSizes.push_back(size);
+    newStrides.push_back(stride);
+    currentDim += dim;
+  }
+  auto layout =
+      makeStridedLinearLayoutMap(newStrides, offset, type.getContext());
+  return MemRefType::get(newSizes, type.getElementType(), {layout});
+}
+
+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 = computeReshapeResultType(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) {
+  unsigned sourceRank = op.getViewType().getRank();
+  unsigned resultRank = op.getResult().getType().cast<MemRefType>().getRank();
+  if (sourceRank == 0 || resultRank == 0)
+    return op.emitOpError("expected non-zero memref ranks");
+  if (resultRank != op.reassociation().size())
+    return op.emitOpError("expected rank of the result view(")
+           << resultRank << ") 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() != sourceRank)
+      return op.emitOpError("expected reassociation map #")
+             << it.index() << " of same rank as source memref(" << sourceRank
+             << "), 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 expectedResultType =
+      computeReshapeResultType(op.getViewType(), maps);
+  if (op.getResult().getType() != expectedResultType)
+    return op.emitOpError("expected result to be ")
+           << expectedResultType << ", but got " << op.getResult().getType();
+  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/Builders.cpp

 //===- Builders.cpp - Helpers for constructing MLIR 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/Builders.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/Functional.h"
LogicalResult OpBuilder::tryFold(Operation *op,
   }
 
   // If we were successful, insert any generated constants.
   for (Operation *cst : generatedConstants)
     insert(cst);
 
   return success();
 }
+
+SmallVector<AffineMap, 4> mlir::getAffineMaps(ArrayAttr attrs) {
+  return functional::map(
+      [](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }, attrs);
+}

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<?x?x?xf32>) {
+  // expected-error @+1 {{expected rank of the result 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 source 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 result to be 'memref<?x?xf32, (d0, d1)[s0] -> (d0 * s0 + d1)>', but got 'memref<?x?xf32>'}}
+  %0 = linalg.reshape %arg0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+    memref<?x?x?xf32> into memref<?x?xf32>
+}

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: #[[strided3D:.*]] = (d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + 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: #[[strided2D_off0:.*]] = (d0, d1)[s0] -> (d0 * s0 + d1)
+// CHECK-DAG: #[[reshapeD0:.*]] = (d0, d1, d2) -> (d0, d1)
+// CHECK-DAG: #[[reshapeD1:.*]] = (d0, d1, d2) -> (d2)
+
 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 @indexed_generic(%arg0: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>,
   } {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(%arg0: memref<?x?x?xf32>) -> memref<?x?xf32, offset: 0, strides: [?, 1]> {
+  %0 = linalg.reshape %arg0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
+    memref<?x?x?xf32> into memref<?x?xf32, offset: 0, strides: [?, 1]>
+  return %0: memref<?x?xf32, offset: 0, strides: [?, 1]>
+}
+// CHECK-LABEL: func @reshape
+//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD0]], #[[reshapeD1]]]
+//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32, #[[strided2D_off0]]>