MLIR  18.0.0git
ComposeSubView.cpp
Go to the documentation of this file.
1 //===- ComposeSubView.cpp - Combining composed subview ops ----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file contains patterns for combining composed subview ops (i.e. subview
10 // of a subview becomes a single subview).
11 //
12 //===----------------------------------------------------------------------===//
13 
18 #include "mlir/IR/OpDefinition.h"
19 #include "mlir/IR/PatternMatch.h"
22 
23 using namespace mlir;
24 
25 namespace {
26 
27 // Replaces a subview of a subview with a single subview. Only supports subview
28 // ops with static sizes and static strides of 1 (both static and dynamic
29 // offsets are supported).
30 struct ComposeSubViewOpPattern : public OpRewritePattern<memref::SubViewOp> {
32 
33  LogicalResult matchAndRewrite(memref::SubViewOp op,
34  PatternRewriter &rewriter) const override {
35  // 'op' is the 'SubViewOp' we're rewriting. 'sourceOp' is the op that
36  // produces the input of the op we're rewriting (for 'SubViewOp' the input
37  // is called the "source" value). We can only combine them if both 'op' and
38  // 'sourceOp' are 'SubViewOp'.
39  auto sourceOp = op.getSource().getDefiningOp<memref::SubViewOp>();
40  if (!sourceOp)
41  return failure();
42 
43  // A 'SubViewOp' can be "rank-reducing" by eliminating dimensions of the
44  // output memref that are statically known to be equal to 1. We do not
45  // allow 'sourceOp' to be a rank-reducing subview because then our two
46  // 'SubViewOp's would have different numbers of offset/size/stride
47  // parameters (just difficult to deal with, not impossible if we end up
48  // needing it).
49  if (sourceOp.getSourceType().getRank() != sourceOp.getType().getRank()) {
50  return failure();
51  }
52 
53  // Offsets, sizes and strides OpFoldResult for the combined 'SubViewOp'.
54  SmallVector<OpFoldResult> offsets, sizes, strides;
55 
56  // Because we only support input strides of 1, the output stride is also
57  // always 1.
58  if (llvm::all_of(strides, [](OpFoldResult &valueOrAttr) {
59  Attribute attr = llvm::dyn_cast_if_present<Attribute>(valueOrAttr);
60  return attr && cast<IntegerAttr>(attr).getInt() == 1;
61  })) {
62  strides = SmallVector<OpFoldResult>(sourceOp.getMixedStrides().size(),
63  rewriter.getI64IntegerAttr(1));
64  } else {
65  return failure();
66  }
67 
68  // The rules for calculating the new offsets and sizes are:
69  // * Multiple subview offsets for a given dimension compose additively.
70  // ("Offset by m" followed by "Offset by n" == "Offset by m + n")
71  // * Multiple sizes for a given dimension compose by taking the size of the
72  // final subview and ignoring the rest. ("Take m values" followed by "Take
73  // n values" == "Take n values") This size must also be the smallest one
74  // by definition (a subview needs to be the same size as or smaller than
75  // its source along each dimension; presumably subviews that are larger
76  // than their sources are disallowed by validation).
77  for (auto it : llvm::zip(op.getMixedOffsets(), sourceOp.getMixedOffsets(),
78  op.getMixedSizes())) {
79  auto opOffset = std::get<0>(it);
80  auto sourceOffset = std::get<1>(it);
81  auto opSize = std::get<2>(it);
82 
83  // We only support static sizes.
84  if (opSize.is<Value>()) {
85  return failure();
86  }
87 
88  sizes.push_back(opSize);
89  Attribute opOffsetAttr = llvm::dyn_cast_if_present<Attribute>(opOffset),
90  sourceOffsetAttr =
91  llvm::dyn_cast_if_present<Attribute>(sourceOffset);
92 
93  if (opOffsetAttr && sourceOffsetAttr) {
94  // If both offsets are static we can simply calculate the combined
95  // offset statically.
96  offsets.push_back(rewriter.getI64IntegerAttr(
97  cast<IntegerAttr>(opOffsetAttr).getInt() +
98  cast<IntegerAttr>(sourceOffsetAttr).getInt()));
99  } else {
100  // When either offset is dynamic, we must emit an additional affine
101  // transformation to add the two offsets together dynamically.
102  AffineExpr expr = rewriter.getAffineConstantExpr(0);
103  SmallVector<Value> affineApplyOperands;
104  for (auto valueOrAttr : {opOffset, sourceOffset}) {
105  if (auto attr = llvm::dyn_cast_if_present<Attribute>(valueOrAttr)) {
106  expr = expr + cast<IntegerAttr>(attr).getInt();
107  } else {
108  expr =
109  expr + rewriter.getAffineSymbolExpr(affineApplyOperands.size());
110  affineApplyOperands.push_back(valueOrAttr.get<Value>());
111  }
112  }
113 
114  AffineMap map = AffineMap::get(0, affineApplyOperands.size(), expr);
115  Value result = rewriter.create<affine::AffineApplyOp>(
116  op.getLoc(), map, affineApplyOperands);
117  offsets.push_back(result);
118  }
119  }
120 
121  // This replaces 'op' but leaves 'sourceOp' alone; if it no longer has any
122  // uses it can be removed by a (separate) dead code elimination pass.
123  rewriter.replaceOpWithNewOp<memref::SubViewOp>(op, sourceOp.getSource(),
124  offsets, sizes, strides);
125  return success();
126  }
127 };
128 
129 } // namespace
130 
132  MLIRContext *context) {
133  patterns.add<ComposeSubViewOpPattern>(context);
134 }
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:44
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
Attributes are known-constant values of operations.
Definition: Attributes.h:25
AffineExpr getAffineSymbolExpr(unsigned position)
Definition: Builders.cpp:357
AffineExpr getAffineConstantExpr(int64_t constant)
Definition: Builders.cpp:361
IntegerAttr getI64IntegerAttr(int64_t value)
Definition: Builders.cpp:128
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:446
This class represents a single result from folding an operation.
Definition: OpDefinition.h:266
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:727
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:539
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:93
void populateComposeSubViewPatterns(RewritePatternSet &patterns, MLIRContext *context)
This header declares functions that assist transformations in the MemRef dialect.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:357
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:361