1 //===- TosaMakeBroadcastable.cpp ------------------------------------------===//
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 // Insert reshape to binary op's input if needed to match rank
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/StandardOps/IR/Ops.h"
14 #include "mlir/Dialect/Tensor/IR/Tensor.h"
15 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
16 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
17 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
18 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
19 #include "mlir/Pass/Pass.h"
20 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
21 
22 using namespace mlir;
23 using namespace mlir::tosa;
24 
25 /// There are two potential ways implementing broadcast:
26 /// a. https://www.tensorflow.org/xla/broadcasting#formal_definition
27 /// b. https://numpy.org/doc/stable/user/basics.broadcasting.html
28 /// TBD: picking option (a) now.
29 
30 /// In this pass, we insert RESHAPE operators to increase the rank of the
31 /// lower rank operand as a first step in the broadcasting process. The TOSA
32 /// operators that support broadcast require that the rank of the operands
33 /// are equal.
34 
35 // Examples:
36 // If lower=[a], target=[a, b, c], [a] reshaped into [a, 1, 1].
37 // TODO: If lower=[b], target=[a, b, c], [b] should but NOT YET reshaped into
38 // [1, b, 1].
39 // If lower=[c], target=[a, b, c], [c] reshaped into [1, 1, c].
40 // If lower=[a, c], target=[a, b, c], [a, c] reshaped into [a, 1, c].
41 // If lower=[a, b], target=[a, b, c], [a, b] reshaped into [a, b, 1].
42 // If lower=[b, c], target=[a, b, c], [b, c] reshaped into [1, b, c].
43 // If lower=[a], target=[a, a], [a] reshaped into [1, a] instead of [a, 1].
44 // If lower=[a], target=[a, b, a], [a] reshaped into [1, 1, a].
45 // If lower=[], target=[a, b, c], [] reshaped into [1, 1, 1].
46 
47 static void computeReshapeOutput(ArrayRef<int64_t> higherRankShape,
48                                  ArrayRef<int64_t> lowerRankShape,
49                                  SmallVectorImpl<int64_t> &reshapeOutputShape) {
50   // Initialize new shapes with [1] * higherRank.
51   int64_t higherRank = higherRankShape.size();
52   int64_t lowerRank = lowerRankShape.size();
53 
54   reshapeOutputShape.assign(higherRank, 1);
55 
56   int64_t higherLeftIndex = 0;
57   int64_t higherRightIndex = higherRank;
58   int64_t lowerLeftIndex = 0;
59   int64_t lowerRightIndex = lowerRank;
60   int64_t higherRankDim, lowerRankDim;
61 
62   if (lowerRightIndex != 0 && higherRightIndex != 0) {
63     // Matches lower rank shape from right dimension first, until not
64     // matching high rank shape or reaching dimension 0.
65     while (true) {
66       higherRankDim = higherRankShape[higherRightIndex - 1];
67       lowerRankDim = lowerRankShape[lowerRightIndex - 1];
68       if (higherRankDim != lowerRankDim)
69         break;
70 
71       reshapeOutputShape[higherRightIndex - 1] = higherRankDim;
72 
73       if (higherRightIndex > 0)
74         higherRightIndex--;
75 
76       if (lowerRightIndex > 0)
77         lowerRightIndex--;
78 
79       if (higherRightIndex == 0 || lowerRightIndex == 0)
80         break;
81     }
82     if (lowerRightIndex != 0 && higherRightIndex != 0) {
83       // Matches lower rank shape from left dimension, until not matching
84       // high rank shape or reaching right index.
85       while (true) {
86         higherRankDim = higherRankShape[higherLeftIndex];
87         lowerRankDim = lowerRankShape[lowerLeftIndex];
88         if (higherRankDim != lowerRankDim)
89           break;
90 
91         reshapeOutputShape[higherLeftIndex] = higherRankDim;
92 
93         if (higherLeftIndex < higherRightIndex)
94           higherLeftIndex++;
95 
96         if (lowerLeftIndex < lowerRightIndex)
97           lowerLeftIndex++;
98 
99         if (higherLeftIndex == higherRightIndex ||
100             lowerLeftIndex == lowerRightIndex)
101           break;
102       }
103     }
104   }
105 }
106 
107 /// Common code to create the reshape op where necessary to make the rank of the
108 /// operations equal. Returns the updated input1 and input2 for the original
109 /// input. The caller is expected to use these to rewrite the original operator
110 /// with the RESHAPE now in the graph.
111 static LogicalResult reshapeLowerToHigher(PatternRewriter &rewriter,
112                                           Location loc,
113                                           RankedTensorType outputType,
114                                           Value input1, Value input2,
115                                           Value &outInput1, Value &outInput2) {
116   auto input1Ty = input1.getType().dyn_cast<RankedTensorType>();
117   auto input2Ty = input2.getType().dyn_cast<RankedTensorType>();
118 
119   if (!input1Ty || !input2Ty)
120     return failure();
121 
122   int64_t input1Rank = input1Ty.getRank();
123   int64_t input2Rank = input2Ty.getRank();
124 
125   Value higherTensorValue, lowerTensorValue;
126   // Cannot rewrite as its already correct.
127   if (input1Rank == input2Rank)
128     return failure();
129 
130   if (input1Rank > input2Rank) {
131     higherTensorValue = input1;
132     lowerTensorValue = input2;
133   } else {
134     higherTensorValue = input2;
135     lowerTensorValue = input1;
136   }
137 
138   ArrayRef<int64_t> higherRankShape =
139       higherTensorValue.getType().cast<RankedTensorType>().getShape();
140   (void)higherRankShape;
141   ArrayRef<int64_t> lowerRankShape =
142       lowerTensorValue.getType().cast<RankedTensorType>().getShape();
143 
144   SmallVector<int64_t, 4> reshapeOutputShape;
145 
146   computeReshapeOutput(outputType.getShape(), lowerRankShape,
147                        reshapeOutputShape);
148 
149   auto reshapeInputType = lowerTensorValue.getType().cast<RankedTensorType>();
150   auto reshapeOutputType = RankedTensorType::get(
151       ArrayRef<int64_t>(reshapeOutputShape), reshapeInputType.getElementType());
152 
153   // Verify the rank agrees with the output type if the output type is ranked.
154   if (outputType) {
155     if (outputType.getShape().size() != reshapeOutputShape.size() ||
156         outputType.getShape().size() != higherRankShape.size())
157       return failure();
158   }
159 
160   auto reshapeLower = rewriter.create<tosa::ReshapeOp>(
161       loc, reshapeOutputType, lowerTensorValue,
162       rewriter.getI64ArrayAttr(reshapeOutputShape));
163 
164   if (input1Rank > input2Rank) {
165     outInput1 = higherTensorValue;
166     outInput2 = reshapeLower.getResult();
167   } else {
168     outInput1 = reshapeLower.getResult();
169     outInput2 = higherTensorValue;
170   }
171 
172   return success();
173 }
174 
175 namespace {
176 template <typename OpTy>
177 struct ConvertTosaOp : public OpRewritePattern<OpTy> {
178   using OpRewritePattern<OpTy>::OpRewritePattern;
179 
180   LogicalResult matchAndRewrite(OpTy tosaBinaryOp,
181                                 PatternRewriter &rewriter) const override {
182 
183     Value input1 = tosaBinaryOp.input1();
184     Value input2 = tosaBinaryOp.input2();
185     Value output = tosaBinaryOp.getResult();
186 
187     auto outputType = output.getType().dyn_cast<RankedTensorType>();
188 
189     Value outInput1, outInput2;
190     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
191                              input1, input2, outInput1, outInput2)
192             .failed())
193       return failure();
194 
195     rewriter.replaceOpWithNewOp<OpTy>(tosaBinaryOp, outputType, outInput1,
196                                       outInput2);
197 
198     return success();
199   }
200 };
201 
202 // The MulOp has an extra parameter 'shift' not present in other elementwise
203 // binary ops, that necessitates special handling of its builder.
204 template <>
205 struct ConvertTosaOp<tosa::MulOp> : public OpRewritePattern<tosa::MulOp> {
206   using OpRewritePattern<tosa::MulOp>::OpRewritePattern;
207 
208   LogicalResult matchAndRewrite(tosa::MulOp tosaBinaryOp,
209                                 PatternRewriter &rewriter) const override {
210 
211     Value input1 = tosaBinaryOp.input1();
212     Value input2 = tosaBinaryOp.input2();
213     int32_t shift = tosaBinaryOp.shift();
214     Value output = tosaBinaryOp.getResult();
215     auto outputType = output.getType().dyn_cast<RankedTensorType>();
216 
217     Value outInput1, outInput2;
218     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
219                              input1, input2, outInput1, outInput2)
220             .failed())
221       return failure();
222 
223     rewriter.replaceOpWithNewOp<tosa::MulOp>(tosaBinaryOp, outputType,
224                                              outInput1, outInput2, shift);
225 
226     return success();
227   }
228 };
229 
230 // The ArithmeticRightShiftOp has an extra parameter 'round' not present in
231 // other elementwise binary ops, that necessitates special handling of its
232 // builder.
233 template <>
234 struct ConvertTosaOp<tosa::ArithmeticRightShiftOp>
235     : public OpRewritePattern<tosa::ArithmeticRightShiftOp> {
236   using OpRewritePattern<tosa::ArithmeticRightShiftOp>::OpRewritePattern;
237 
238   LogicalResult matchAndRewrite(tosa::ArithmeticRightShiftOp tosaBinaryOp,
239                                 PatternRewriter &rewriter) const override {
240 
241     Value input1 = tosaBinaryOp.input1();
242     Value input2 = tosaBinaryOp.input2();
243     int32_t round = tosaBinaryOp.round();
244     Value output = tosaBinaryOp.getResult();
245     auto outputType = output.getType().dyn_cast<RankedTensorType>();
246 
247     Value outInput1, outInput2;
248     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
249                              input1, input2, outInput1, outInput2)
250             .failed())
251       return failure();
252 
253     rewriter.replaceOpWithNewOp<tosa::ArithmeticRightShiftOp>(
254         tosaBinaryOp, outputType, outInput1, outInput2, round);
255 
256     return success();
257   }
258 };
259 } // end anonymous namespace
260 
261 namespace {
262 /// Pass that enables broadcast by making all input arrays have the same
263 /// number of dimensions. Insert RESHAPE operations to lower rank operand
264 struct TosaMakeBroadcastable
265     : public TosaMakeBroadcastableBase<TosaMakeBroadcastable> {
266 public:
267   void runOnFunction() override {
268     auto func = getFunction();
269     RewritePatternSet patterns(func.getContext());
270     MLIRContext *ctx = func.getContext();
271     // Add the generated patterns to the list.
272     patterns.add<ConvertTosaOp<tosa::BitwiseAndOp>>(ctx);
273     patterns.add<ConvertTosaOp<tosa::BitwiseOrOp>>(ctx);
274     patterns.add<ConvertTosaOp<tosa::BitwiseXorOp>>(ctx);
275     patterns.add<ConvertTosaOp<tosa::AddOp>>(ctx);
276     patterns.add<ConvertTosaOp<tosa::SubOp>>(ctx);
277     patterns.add<ConvertTosaOp<tosa::MulOp>>(ctx);
278     patterns.add<ConvertTosaOp<tosa::DivOp>>(ctx);
279     patterns.add<ConvertTosaOp<tosa::MaximumOp>>(ctx);
280     patterns.add<ConvertTosaOp<tosa::MinimumOp>>(ctx);
281     patterns.add<ConvertTosaOp<tosa::EqualOp>>(ctx);
282     patterns.add<ConvertTosaOp<tosa::GreaterOp>>(ctx);
283     patterns.add<ConvertTosaOp<tosa::GreaterEqualOp>>(ctx);
284     patterns.add<ConvertTosaOp<tosa::LogicalLeftShiftOp>>(ctx);
285     patterns.add<ConvertTosaOp<tosa::ArithmeticRightShiftOp>>(ctx);
286     patterns.add<ConvertTosaOp<tosa::LogicalRightShiftOp>>(ctx);
287     patterns.add<ConvertTosaOp<tosa::LogicalAndOp>>(ctx);
288     patterns.add<ConvertTosaOp<tosa::LogicalOrOp>>(ctx);
289     patterns.add<ConvertTosaOp<tosa::LogicalXorOp>>(ctx);
290     patterns.add<ConvertTosaOp<tosa::PowOp>>(ctx);
291     (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
292   }
293 };
294 } // end anonymous namespace
295 
296 std::unique_ptr<Pass> mlir::tosa::createTosaMakeBroadcastablePass() {
297   return std::make_unique<TosaMakeBroadcastable>();
298 }
299