[mlir][Linalg] Introduce folding patterns to remove certain MemRefCastOp

[nicolasvasilache]

Summary

[mlir][Linalg] Introduce folding patterns to remove certain MemRefCastOp

Summary:
Canonicalization and folding patterns in StandardOps may interfere with the needs
of Linalg. This revision introduces specific foldings for dynamic memrefs that can
be proven to be static.

Very concretely:

Determines whether it is possible to fold it away in the parent Linalg op:

```mlir
  %1 = memref_cast %0 : memref<8x16xf32> to memref<?x?xf32>
  %2 = linalg.slice %1 ... : memref<?x?xf32> ...
  // or
  %1 = memref_cast %0 : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
         to memref<?x?xf32>
  linalg.generic(%1 ...) : memref<?x?xf32> ...
```

into

```mlir
  %2 = linalg.slice %0 ... : memref<8x16xf32> ...
  // or
  linalg.generic(%0 ... : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
```

[nicolasvasilache]

Remove spurious pattern.
Fix comment.

[merge_guards_bot]

Unit tests: pass. 62018 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

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

[merge_guards_bot]

Unit tests: pass. 62018 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

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

[nicolasvasilache]

Fix the test, give up on a single pass in this iteration.

[merge_guards_bot]

Unit tests: pass. 62018 tests passed, 0 failed and 783 were skipped.

clang-tidy: pass.

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

[asaadaldien]

LGTM!

[ftynse]

mlir/test/Dialect/Linalg/matmul-to-vector.mlir

3 ↗

(On Diff #239756)

Is this comment relevant to the test?

25 ↗

(On Diff #239756)

Could you describe where do this magic sizes come from?

mlir/test/lib/DeclarativeTransforms/TestLinalgMatmulToVectorPatterns.td

3 ↗

(On Diff #239756)

MLIR Project

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

27 ↗

(On Diff #239756)

a MemrefCastOp, you don't change the LinalgOp

39 ↗

(On Diff #239756)

I don't understand this comment.

If you have a memref_cast %0 fed by %0 = %0 = memref_cast %? : ... to memref<?x?xf32>, why do we convert it to memref_cast %0 : ... to memref<8x16xf32>`. There's something wrong with the order in your description.

I understand it even less after reading the code.

48 ↗

(On Diff #239756)

Nit: cast

56 ↗

(On Diff #239756)

Can verified memref_cast apply to non-memrefs?

81 ↗

(On Diff #239756)

Leftover comment

83 ↗

(On Diff #239756)

Nit: I'd put it before iterating over the use list

95 ↗

(On Diff #239756)

Please add a comment, it looks like this block is useless (it isn't the trick is in the "linear" part).

112 ↗

(On Diff #239756)

I understand this is a test code, but having a separate test that exercises this pattern specifically would definitely help. One one hand, avoid spurious failures in a bigger test. On the other hand, better see what kind of transformation you are doing here.

[jsetoain]

mlir/test/lib/Transforms/CMakeLists.txt

29 ↗

(On Diff #239756)

I'm not sure if the policy is to keep these in lexycographical order?

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

3 ↗

(On Diff #239756)

Part of the LLVM Project or part of the MLIR Project?

[nicolasvasilache]

Thanks @ftynse I'll split out the new pattern in an independent test and rebase.

Note that comments in this diff are also relevant to the problem at hand: https://reviews.llvm.org/D73295.

[tetuante]

mlir/test/Dialect/Linalg/matmul-to-vector.mlir

25 ↗

(On Diff #239756)

I think they come from

def : Pat<(MatmulOp:$op $_, $_, $_),
          (TileLinalgOp<[8, 12, 16], "L1__with_perm__", [1, 0, 2]>),
         [(Constraint<HasLinalgTransformMarker<"L2__with_perm__">>)]>;

in TestLinalgMatmulToVectorPatterns.td

They are fixed tile sizes right now for testing, but should somehow be computed in the futre

[mehdi_amini]

mlir/test/lib/Transforms/TestLinalgMatmulToVector.cpp

3 ↗

(On Diff #239756)

Thanks for noticing. I just mass-updated all the MLIR files in the repo to mention "Part of the LLVM project".

[nicolasvasilache]

Format.

[nicolasvasilache]

Apologies for the delay, I recast this diff as an independent linalg canonicalization + test only as @ftynse suggested.
I'll be rebasing the rest on top of this in a new diff.

[rriddle]

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

73

When will this ever be hit? The desiredResultType doesn't appear to have any maps.

78

nit: getUsers

102

nit: Can we use early exit here?

1173

Drop the llvm:: and mlir:: on all of these.

[merge_guards_bot]

Unit tests: fail. 62196 tests passed, 1 failed and 815 were skipped.

failed: MLIR.mlir-cpu-runner/unranked_memref.mlir

clang-tidy: pass.

clang-format: pass.

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

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

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

78

incorrectly carried over, case removed.

[nicolasvasilache]

Address review comments.

[merge_guards_bot]

Unit tests: pass. 62197 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

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

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

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

87–91

Let's exit even earlier

if (sourceType.getAffineMaps.empty())
  return true;

AffineMap sourceMap = sourceType.getAffineMaps().front();

100

return sourceMap == stridedMap ?

[jsetoain]

LGTM.

[nicolasvasilache]

Thanks for your reviews!

[nicolasvasilache]

This landed here: https://github.com/llvm/llvm-project/commit/ea1e3369f7a8aa9729f8e2fc208b8f6a79392874

But since I made a mistake with phab I need to close manually.

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

let builders = [
     OpBuilder<"Builder *b, OperationState &result, Type resultType, Value view,"
     "ArrayRef<ArrayRef<AffineExpr>> reassociation, "
     "ArrayRef<NamedAttribute> attrs = {}">];
 
   let extraClassDeclaration = [{
     static StringRef getReassociationAttrName() { return "reassociation"; }
     MemRefType getViewType() { return view().getType().cast<MemRefType>(); }
   }];
+
+  let hasFolder = 1;
 }
 
 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 = [{
let extraClassDeclaration = [{
     SmallVector<Value, 8> getRanges() {
       SmallVector<Value, 8> res;
       for (auto operand : indexings())
         if (!operand.getType().isa<IndexType>())
           res.push_back(operand);
       return res;
     }
   }];
+
+  let hasFolder = 1;
 }
 
 def Linalg_TransposeOp : Linalg_Op<"transpose", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view, AffineMapAttr:$permutation)>,
     Results<(outs AnyStridedMemRef)> {
   let summary = "`transpose` produces a new strided memref (metadata-only)";
   let description = [{
     The `linalg.transpose` op produces a strided memref whose sizes and strides
if (permutation().getNumDims() != getShapedType().getRank())
       return emitOpError("expected a permutation map of same rank as the view");
     return success();
   }];
 
   let extraClassDeclaration = [{
     static StringRef getPermutationAttrName() { return "permutation"; }
     ShapedType getShapedType() { return view().getType().cast<ShapedType>(); }
   }];
+
+  let hasFolder = 1;
 }
 
 def Linalg_YieldOp : Linalg_Op<"yield", [NativeOpTrait<"IsTerminator">]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";
   let description = [{
     `linalg.yield` is a special terminator operation for blocks inside regions
     in `linalg` generic ops. It returns values to the immediately enclosing

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

ArrayAttr iterator_types() {
       unsigned nPar = input().getType().cast<ShapedType>().getRank();
       MLIRContext *ctx = getContext();
       SmallVector<Attribute, 8> iters(
         nPar, StringAttr::get(getParallelIteratorTypeName(), ctx));
       return ArrayAttr::get(iters, ctx);
     }
   }];
   let verifier = [{ return ::verify(*this); }];
+
+  let hasFolder = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", [NInputs<0>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRef:$input,
                    AnyTypeOf<[AnyFloat, AnyInteger, AnyVector]>:$value);
   let extraClassDeclaration = libraryCallName # [{
     ArrayAttr indexing_maps();
 
     ArrayAttr iterator_types() {
       unsigned nPar = input().getType().cast<ShapedType>().getRank();
       MLIRContext *ctx = getContext();
       SmallVector<Attribute, 8> iters(
         nPar, StringAttr::get(getParallelIteratorTypeName(), ctx));
       return ArrayAttr::get(iters, ctx);
     }
   }];
   let verifier = [{ return ::verify(*this); }];
+
+  let hasFolder = 1;
 }
 
 def DotOp : LinalgStructured_Op<"dot", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<0>);
   let extraClassDeclaration = libraryCallName # [{
     ArrayAttr indexing_maps();
 
     ArrayAttr iterator_types() {
       MLIRContext *ctx = getContext();
       return ArrayAttr::get(
         StringAttr::get(getReductionIteratorTypeName(), ctx), ctx);
     }
   }];
+
+  let hasFolder = 1;
 }
 
 def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<1>);
   let extraClassDeclaration = libraryCallName # [{
     ArrayAttr indexing_maps();
 
     ArrayAttr iterator_types() {
       MLIRContext *ctx = getContext();
       Attribute iters[2]{
         StringAttr::get(getParallelIteratorTypeName(), ctx),
         StringAttr::get(getReductionIteratorTypeName(), ctx)};
       return ArrayAttr::get(iters, ctx);
     }
   }];
+
+  let hasFolder = 1;
 }
 
 def MatmulOp : LinalgStructured_Op<"matmul", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<2>);
   let extraClassDeclaration = libraryCallName # [{
     ArrayAttr indexing_maps();
 
     ArrayAttr iterator_types() {
       MLIRContext *ctx = getContext();
       Attribute iters[3]{
         StringAttr::get(getParallelIteratorTypeName(), ctx),
         StringAttr::get(getParallelIteratorTypeName(), ctx),
         StringAttr::get(getReductionIteratorTypeName(), ctx)};
       return ArrayAttr::get(iters, ctx);
     }
   }];
+
+  let hasFolder = 1;
 }
 
 def ConvOp : LinalgStructured_Op<"conv", [NInputs<2>, NOutputs<1>]> {
   let description = [{
     Generic n-D convolution as described in the TF documentation:
     https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/convolution
 
     ```
let extraClassDeclaration = libraryCallName # [{
 
     int64_t getDilation(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!dilations().hasValue()) return 1;
       return dilations()->getValue()[i]
         .cast<IntegerAttr>().getValue().getSExtValue();
     }
   }];
+
   let verifier = [{ return ::verify(*this); }];
+
+  let hasFolder = 1;
 }
 
 def LinalgOperand: Type<
   Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>;
 
 class LinalgOperandOfRank<int rank>: Type<
   And<[
     LinalgOperand.predicate,
let description = [{
     into output buffer operands and updates the region arguments accordingly.
 
     Transformations that create control-flow around linalg.indexed_generic
     operations are not expected to work with tensors because SSA values do not
     escape naturally. Still, transformations and rewrites that take advantage of
     tensor SSA values are expected to be useful and will be added in the near
     future.
   }];
+
   let verifier = [{ return ::verify(*this); }];
+
+  let hasFolder = 1;
 }
 
 def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
   let description = [{
     Indexed Generic Linalg op form where the key properties of the computation
     are specified as attributes. In pretty form, a linalg.indexed_generic op is
     written as:
 
let description = [{
     into output buffer operands and updates the region argument accordingly.
 
     Transformations that create control-flow around linalg.indexed_generic
     operations are not expected to work with tensors because SSA values do not
     escape naturally. Still, transformations and rewrites that take advantage of
     tensor SSA values are expected to be useful and will be added in the near
     future.
   }];
+
   let verifier = [{ return ::verify(*this); }];
+
+  let hasFolder = 1;
 }
 
 #endif // LINALG_STRUCTURED_OPS

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 the Linalg operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
+#include "mlir/Dialect/StandardOps/Ops.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 "llvm/ADT/StringSet.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
+/// Determines whether it is possible to fold it away in the parent Linalg op:
+///
+/// ```mlir
+///   %1 = memref_cast %0 : memref<8x16xf32> to memref<?x?xf32>
+///   %2 = linalg.slice %1 ... : memref<?x?xf32> ...
+///   // or
+///   %1 = memref_cast %0 : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
+///          to memref<?x?xf32>
+///   linalg.generic(%1 ...) : memref<?x?xf32> ...
+/// ```
+///
+/// into
+///
+/// ```mlir
+///   %2 = linalg.slice %0 ... : memref<8x16xf32> ...
+///   // or
+///   linalg.generic(%0 ... : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
+/// ```
+///
+static bool canFold(MemRefCastOp castOp) {
+  MemRefType sourceType = castOp.source().getType().dyn_cast<MemRefType>();
+  MemRefType resultType = castOp.getType().dyn_cast<MemRefType>();
+
+  // If we don't have MemRefType as source and destination, bail out.
+  if (!sourceType || !resultType)
+    return false;
+
+  // If resultType has a map, it needs to be the same as the source type to
+  // canonicalize.
+  if (!resultType.getAffineMaps().empty() &&
+      sourceType.getAffineMaps() != resultType.getAffineMaps())
+    return false;
+
+  // Ensure that:
+  //   1. source is static
+  //   2. source and target have the same rank (will be extended when needed)
+  //   3. if result is partially static, ensure sizes match.
+  if (!sourceType.hasStaticShape() ||
+      sourceType.getRank() != resultType.getRank())
+    return false;
+
+  for (auto it : llvm::zip(sourceType.getShape(), resultType.getShape())) {
+    auto sourceSize = std::get<0>(it);
+    auto resultSize = std::get<1>(it);
+    if (ShapedType::isDynamic(resultSize))
+      continue;
+    if (sourceSize != resultSize)
+      return false;
+  }
+
+  // If source has a map, it can only canonicalize if it is the canonical
+  // strided layout map.
+  AffineMap sourceMap = (sourceType.getAffineMaps().empty())
+                            ? AffineMap()
+                            : sourceType.getAffineMaps().front();
+  if (!sourceMap)
+    return true;
+
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto res = getStridesAndOffset(sourceType, strides, offset);
+  (void)res;
+  assert(succeeded(res));
+  auto stridedMap =
+      makeStridedLinearLayoutMap(strides, offset, sourceMap.getContext());
+  if (sourceMap != stridedMap)
+    return false;
+
+  return true;
+}
+
+/// This is a common class used for patterns of the form
+/// ```
+///    someop(memrefcast) -> someop
+/// ```
+/// It folds the source of any memref_cast into the root operation directly.
+static LogicalResult foldMemRefCast(Operation *op) {
+  bool folded = false;
+  for (OpOperand &operand : op->getOpOperands()) {
+    auto castOp = dyn_cast_or_null<MemRefCastOp>(operand.get().getDefiningOp());
+    if (castOp && canFold(castOp)) {
+      operand.set(castOp.getOperand());
+      folded = true;
+    }
+  }
+  return success(folded);
+}
+
 ///////////////////// Operations defined with Tablegen /////////////////////////
 // For such operations that do not correspond to library calls (i.e. defined in
 // LinalgOps.td), we define an overloaded `print` function and a
 // parse`className` function.
 
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
ArrayAttr mlir::linalg::FillOp::indexing_maps() {
   return getIndexingMaps(getOperation());
 }
 ArrayAttr mlir::linalg::MatmulOp::indexing_maps() {
   return getIndexingMaps(getOperation());
 }
 ArrayAttr mlir::linalg::MatvecOp::indexing_maps() {
   return getIndexingMaps(getOperation());
 }
+
+// TODO(ntv, rriddle): 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.
+LogicalResult ConvOp::fold(ArrayRef<Attribute>,
+                           SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult CopyOp::fold(ArrayRef<Attribute>,
+                           SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult DotOp::fold(ArrayRef<Attribute>,
+                          SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult FillOp::fold(ArrayRef<Attribute>,
+                           SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult GenericOp::fold(ArrayRef<Attribute>,
+                              SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult IndexedGenericOp::fold(ArrayRef<Attribute>,
+                                     SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult MatvecOp::fold(ArrayRef<Attribute>,
+                             SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+LogicalResult MatmulOp::fold(ArrayRef<Attribute>,
+                             SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
+OpFoldResult ReshapeOp::fold(ArrayRef<Attribute>) {
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return {};
+}
+OpFoldResult SliceOp::fold(ArrayRef<Attribute>) {
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return {};
+}
+OpFoldResult TransposeOp::fold(ArrayRef<Attribute>) {
+  if (succeeded(foldMemRefCast(*this)))
+    return getResult();
+  return {};
+}

mlir/test/Dialect/Linalg/canonicalize.mlir

+// RUN: mlir-opt %s -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @memref_cast(
+func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c8 = constant 8 : index
+  %c16 = constant 16 : index
+  %1 = alloc (%b) : memref<?xi8>
+  %2 = view %1[][] : memref<?xi8> to memref<16x16xf32>
+  %3 = memref_cast %2 : memref<16x16xf32> to memref<?x?xf32>
+  %r0 = linalg.range %c0:%c8:%c1 : !linalg.range
+
+  // CHECK:  linalg.slice {{.*}} : memref<16x16xf32>, !linalg.range, !linalg.range, memref<?x?xf32>
+  %4 = linalg.slice %3[%r0, %r0] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32>
+
+  // CHECK:  linalg.matmul{{.*}}: memref<16x16xf32>, memref<16x16xf32>, memref<16x16xf32>
+  linalg.matmul(%3, %3, %3) : memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  return %4: memref<?x?xf32>
+}