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
17namespace mlir {
18#define GEN_PASS_DEF_CONVERTELEMENTWISETOLINALGPASS
19#include "mlir/Dialect/Linalg/Passes.h.inc"
20} // namespace mlir
21
22using namespace mlir;
23
24static 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".
46static SmallVector<Value, 4>
47getOrCreateOperandsMatchingResultTypes(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(N: 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(Elt: 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
75namespace {
76struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
77 ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context)
78 : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
79 LogicalResult matchAndRewrite(Operation *op,
80 PatternRewriter &rewriter) const final {
81 if (!isElementwiseMappableOpOnRankedTensors(op))
82 return rewriter.notifyMatchFailure(
83 arg&: op, msg: "requires elementwise op on ranked tensors");
84
85 auto rank = cast<RankedTensorType>(op->getResult(idx: 0).getType()).getRank();
86 SmallVector<AffineMap, 3> indexingMaps(
87 op->getNumResults() + op->getNumOperands(),
88 rewriter.getMultiDimIdentityMap(rank: rank));
89 SmallVector<utils::IteratorType, 6> iteratorTypes(
90 rank, utils::IteratorType::parallel);
91 auto outputs = getOrCreateOperandsMatchingResultTypes(b&: 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
115void mlir::linalg::populateElementwiseToLinalgConversionPatterns(
116 RewritePatternSet &patterns) {
117 patterns.add<ConvertAnyElementwiseMappableOpOnRankedTensors>(
118 arg: patterns.getContext());
119}
120
121namespace {
122class ConvertElementwiseToLinalgPass
123 : public impl::ConvertElementwiseToLinalgPassBase<
124 ConvertElementwiseToLinalgPass> {
125 using impl::ConvertElementwiseToLinalgPassBase<
126 ConvertElementwiseToLinalgPass>::ConvertElementwiseToLinalgPassBase;
127
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(fn: [](Operation *op) {
136 return !isElementwiseMappableOpOnRankedTensors(op);
137 });
138
139 if (failed(applyPartialConversion(func, target, std::move(patterns))))
140 signalPassFailure();
141 }
142};
143} // namespace
144

source code of mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp