1 //===- ElementwiseToLinalg.cpp - conversion of elementwise to linalg ------===//
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 #include "mlir/Dialect/Linalg/Passes.h"
10
11 #include "mlir/Dialect/Arith/Utils/Utils.h"
12 #include "mlir/Dialect/Linalg/IR/Linalg.h"
13 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
14 #include "mlir/Dialect/Linalg/Utils/Utils.h"
15 #include "mlir/Transforms/DialectConversion.h"
16
17 namespace mlir {
18 #define GEN_PASS_DEF_CONVERTELEMENTWISETOLINALGPASS
19 #include "mlir/Dialect/Linalg/Passes.h.inc"
20 } // namespace mlir
21
22 using namespace mlir;
23
isElementwiseMappableOpOnRankedTensors(Operation * op)24 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
25 if (!OpTrait::hasElementwiseMappableTraits(op))
26 return false;
27
28 // TODO: The conversion pattern can be made to work for `any_of` here, but
29 // it's more complex as it requires tracking which operands are scalars.
30 return llvm::all_of(op->getOperandTypes(), llvm::IsaPred<RankedTensorType>);
31 }
32
33 /// Given `op` assumed `isElementwiseMappableOpOnRankedTensors`, iterate over
34 /// the result types and return a list of values such that, for each result type
35 /// `t` and value `v` at the same index `idx`:
36 /// 1. `v.getType() == t`
37 /// 2. If an operand of `op` has type `t`, let `operand_first` be the first
38 /// such operand. Then`v == operand_first`.
39 /// 3. Otherwise, v is a newly created `tensor::EmptyOp` with:
40 /// a. Static and dynamic dims extracted from the first operand of `op`.
41 /// b. Elemental type equal to the elemental type of `t`.
42 ///
43 /// This is sufficient because ElementwiseMappable guarantees that "The static
44 /// types of all vector (resp. tensor) operands and results must have the same
45 /// shape".
46 static SmallVector<Value, 4>
getOrCreateOperandsMatchingResultTypes(OpBuilder & b,Operation * op)47 getOrCreateOperandsMatchingResultTypes(OpBuilder &b, Operation *op) {
48 assert(isElementwiseMappableOpOnRankedTensors(op));
49 Location loc = op->getLoc();
50 ValueRange operands = op->getOperands();
51 TypeRange rankedTensorTypes = op->getResultTypes();
52 SmallVector<Value, 4> res;
53 res.reserve(rankedTensorTypes.size());
54 for (Type t : rankedTensorTypes) {
55 // Try to find an operand with type matching the result tensor.
56 bool found = false;
57 for (Value v : operands) {
58 if (v.getType() == t) {
59 found = true;
60 res.push_back(v);
61 break;
62 }
63 }
64 if (found)
65 continue;
66
67 // Extract static / dynamic shape mix from the first operand.
68 res.push_back(b.create<tensor::EmptyOp>(
69 loc, tensor::getMixedSizes(b, loc, operands.front()),
70 cast<RankedTensorType>(t).getElementType()));
71 }
72 return res;
73 }
74
75 namespace {
76 struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
ConvertAnyElementwiseMappableOpOnRankedTensors__anonab5fbb810111::ConvertAnyElementwiseMappableOpOnRankedTensors77 ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context)
78 : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
matchAndRewrite__anonab5fbb810111::ConvertAnyElementwiseMappableOpOnRankedTensors79 LogicalResult matchAndRewrite(Operation *op,
80 PatternRewriter &rewriter) const final {
81 if (!isElementwiseMappableOpOnRankedTensors(op))
82 return rewriter.notifyMatchFailure(
83 op, "requires elementwise op on ranked tensors");
84
85 auto rank = cast<RankedTensorType>(op->getResult(0).getType()).getRank();
86 SmallVector<AffineMap, 3> indexingMaps(
87 op->getNumResults() + op->getNumOperands(),
88 rewriter.getMultiDimIdentityMap(rank));
89 SmallVector<utils::IteratorType, 6> iteratorTypes(
90 rank, utils::IteratorType::parallel);
91 auto outputs = getOrCreateOperandsMatchingResultTypes(rewriter, op);
92 rewriter.replaceOpWithNewOp<linalg::GenericOp>(
93 op, /*resultTensorTypes=*/op->getResultTypes(),
94 /*inputs=*/op->getOperands(),
95 /*outputs=*/outputs,
96 /*indexingMaps=*/indexingMaps,
97 /*iteratorTypes=*/iteratorTypes,
98 /*bodyBuilder=*/
99 [&](OpBuilder &builder, Location loc, ValueRange regionArgs) {
100 auto resultTypes = llvm::to_vector<6>(
101 llvm::map_range(op->getResultTypes(), [](Type type) {
102 return cast<TensorType>(type).getElementType();
103 }));
104 auto *scalarOp =
105 builder.create(loc, op->getName().getIdentifier(),
106 regionArgs.take_front(op->getNumOperands()),
107 resultTypes, op->getAttrs());
108 builder.create<linalg::YieldOp>(loc, scalarOp->getResults());
109 });
110 return success();
111 }
112 };
113 } // namespace
114
populateElementwiseToLinalgConversionPatterns(RewritePatternSet & patterns)115 void mlir::linalg::populateElementwiseToLinalgConversionPatterns(
116 RewritePatternSet &patterns) {
117 patterns.add<ConvertAnyElementwiseMappableOpOnRankedTensors>(
118 patterns.getContext());
119 }
120
121 namespace {
122 class ConvertElementwiseToLinalgPass
123 : public impl::ConvertElementwiseToLinalgPassBase<
124 ConvertElementwiseToLinalgPass> {
125 using impl::ConvertElementwiseToLinalgPassBase<
126 ConvertElementwiseToLinalgPass>::ConvertElementwiseToLinalgPassBase;
127
runOnOperation()128 void runOnOperation() final {
129 auto *func = getOperation();
130 auto *context = &getContext();
131 ConversionTarget target(*context);
132 RewritePatternSet patterns(context);
133
134 mlir::linalg::populateElementwiseToLinalgConversionPatterns(patterns);
135 target.markUnknownOpDynamicallyLegal([](Operation *op) {
136 return !isElementwiseMappableOpOnRankedTensors(op);
137 });
138
139 if (failed(applyPartialConversion(func, target, std::move(patterns))))
140 signalPassFailure();
141 }
142 };
143 } // namespace
144