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 "PassDetail.h"
12 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
13 #include "mlir/Dialect/StandardOps/IR/Ops.h"
14 #include "mlir/Transforms/DialectConversion.h"
15
16 using namespace mlir;
17
isElementwiseMappableOpOnRankedTensors(Operation * op)18 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
19 if (!op->hasTrait<OpTrait::ElementwiseMappable>())
20 return false;
21
22 // TODO: The conversion pattern can be made to work for `any_of` here, but
23 // it's more complex as it requires tracking which operands are scalars.
24 return llvm::all_of(op->getOperandTypes(),
25 [](Type type) { return type.isa<RankedTensorType>(); });
26 }
27
28 namespace {
29 struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
ConvertAnyElementwiseMappableOpOnRankedTensors__anon075008450211::ConvertAnyElementwiseMappableOpOnRankedTensors30 ConvertAnyElementwiseMappableOpOnRankedTensors()
31 : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
matchAndRewrite__anon075008450211::ConvertAnyElementwiseMappableOpOnRankedTensors32 LogicalResult matchAndRewrite(Operation *op,
33 PatternRewriter &rewriter) const final {
34 if (!isElementwiseMappableOpOnRankedTensors(op))
35 return rewriter.notifyMatchFailure(
36 op, "requires elementwise op on ranked tensors");
37
38 auto rank = op->getResult(0).getType().cast<RankedTensorType>().getRank();
39 SmallVector<AffineMap, 3> indexingMaps(
40 op->getNumResults() + op->getNumOperands(),
41 rewriter.getMultiDimIdentityMap(rank));
42 SmallVector<StringRef, 6> iteratorTypes(rank,
43 getParallelIteratorTypeName());
44 rewriter.replaceOpWithNewOp<linalg::GenericOp>(
45 op, /*resultTensorTypes=*/op->getResultTypes(),
46 /*inputs=*/op->getOperands(),
47 /*outputBuffers=*/ValueRange(),
48 /*initTensors=*/ValueRange(),
49 /*indexingMaps=*/indexingMaps,
50 /*iteratorTypes=*/iteratorTypes,
51 /*bodyBuilder=*/
52 [&](OpBuilder &builder, Location loc, ValueRange regionArgs) {
53 OperationState state(loc, op->getName());
54 state.addAttributes(op->getAttrs());
55 state.addOperands(regionArgs);
56 auto resultTypes = llvm::to_vector<6>(
57 llvm::map_range(op->getResultTypes(), [](Type type) {
58 return type.cast<TensorType>().getElementType();
59 }));
60 state.addTypes(resultTypes);
61 auto *scalarOp = builder.createOperation(state);
62 builder.create<linalg::YieldOp>(loc, scalarOp->getResults());
63 });
64 return success();
65 }
66 };
67 } // namespace
68
populateElementwiseToLinalgConversionPatterns(OwningRewritePatternList & patterns,MLIRContext *)69 void mlir::populateElementwiseToLinalgConversionPatterns(
70 OwningRewritePatternList &patterns, MLIRContext *) {
71 patterns.insert<ConvertAnyElementwiseMappableOpOnRankedTensors>();
72 }
73
74 namespace {
75 class ConvertElementwiseToLinalgPass
76 : public ConvertElementwiseToLinalgBase<ConvertElementwiseToLinalgPass> {
77
runOnFunction()78 void runOnFunction() final {
79 auto func = getOperation();
80 auto *context = &getContext();
81 ConversionTarget target(*context);
82 OwningRewritePatternList patterns;
83
84 populateElementwiseToLinalgConversionPatterns(patterns, context);
85 target.markUnknownOpDynamicallyLegal([](Operation *op) {
86 return !isElementwiseMappableOpOnRankedTensors(op);
87 });
88
89 if (failed(applyPartialConversion(func, target, std::move(patterns))))
90 signalPassFailure();
91 }
92 };
93 } // namespace
94
95 std::unique_ptr<OperationPass<FuncOp>>
createConvertElementwiseToLinalgPass()96 mlir::createConvertElementwiseToLinalgPass() {
97 return std::make_unique<ConvertElementwiseToLinalgPass>();
98 }
99