1 //===- TosaDecomposeTransposeConv.cpp
2 //------------------------------------------===//
3 //
4 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5 // See https://llvm.org/LICENSE.txt for license information.
6 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // Insert reshape to binary op's input if needed to match rank
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Dialect/StandardOps/IR/Ops.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/ShapeUtils.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 namespace {
26 
27 template <typename T>
28 static void getValuesFromIntArrayAttribute(ArrayAttr attr,
29                                            SmallVector<T> &arrayValues) {
30   for (Attribute val : attr.getValue()) {
31     arrayValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
32   }
33 }
34 
35 template <typename TosaOp, typename... Args>
36 TosaOp CreateOpAndInfer(PatternRewriter &rewriter, Location loc, Type result_ty,
37                         Args &&...args) {
38   auto op = rewriter.create<TosaOp>(loc, result_ty, args...);
39 
40   InferShapedTypeOpInterface shapeInterface =
41       dyn_cast<InferShapedTypeOpInterface>(op.getOperation());
42   if (!shapeInterface)
43     return op;
44 
45   SmallVector<ShapedTypeComponents> returnedShapes;
46   if (shapeInterface
47           .inferReturnTypeComponents(op.getContext(), op.getLoc(),
48                                      op->getOperands(), op->getAttrDictionary(),
49                                      op->getRegions(), returnedShapes)
50           .failed())
51     return op;
52 
53   // We need to use the element type of the existing result type to generate
54   // the new result shaped type. This is because rescale can include a cast to
55   // different bit-width types and does not have a TypeAttr to define the
56   // target type.
57   auto result = op->getResult(0);
58   auto predictedShape = returnedShapes[0];
59   auto currentKnowledge =
60       mlir::tosa::ValueKnowledge::getKnowledgeFromType(result_ty);
61 
62   // Compute the knowledge based on the inferred type.
63   auto inferredKnowledge =
64       mlir::tosa::ValueKnowledge::getPessimisticValueState();
65   inferredKnowledge.dtype = result_ty.cast<ShapedType>().getElementType();
66   inferredKnowledge.hasRank = predictedShape.hasRank();
67   if (predictedShape.hasRank()) {
68     for (auto dim : predictedShape.getDims()) {
69       inferredKnowledge.sizes.push_back(dim);
70     }
71   }
72 
73   // Compute the new type based on the joined version.
74   auto newKnowledge =
75       mlir::tosa::ValueKnowledge::join(currentKnowledge, inferredKnowledge);
76   auto new_ty = newKnowledge.getType();
77   result.setType(new_ty);
78   return op;
79 }
80 
81 class TransposeConvDilatedConverter
82     : public OpRewritePattern<tosa::TransposeConv2DOp> {
83 public:
84   using OpRewritePattern<tosa::TransposeConv2DOp>::OpRewritePattern;
85   LogicalResult matchAndRewrite(tosa::TransposeConv2DOp op,
86                                 PatternRewriter &rewriter) const final {
87     Location loc = op->getLoc();
88     Value input = op->getOperand(0);
89     Value weight = op->getOperand(1);
90     Value bias = op->getOperand(2);
91 
92     ShapedType inputTy = input.getType().cast<ShapedType>();
93     ShapedType weightTy = weight.getType().cast<ShapedType>();
94     ShapedType biasTy = bias.getType().cast<ShapedType>();
95     ShapedType resultTy = op->getResult(0).getType().cast<ShapedType>();
96 
97     llvm::SmallVector<int64_t> pad;
98     llvm::SmallVector<int64_t> stride;
99     llvm::SmallVector<int64_t> dilation;
100 
101     getValuesFromIntArrayAttribute(op.out_pad().cast<ArrayAttr>(), pad);
102     getValuesFromIntArrayAttribute(op.stride().cast<ArrayAttr>(), stride);
103     getValuesFromIntArrayAttribute(op.dilation().cast<ArrayAttr>(), dilation);
104 
105     // If striding is all 1 we can modify padding and reverse the kernel along
106     // the x/y direction to make it a regular convolution. This is much simpler
107     // then handling striding....
108     if (llvm::any_of(stride, [](int64_t v) { return v != 1; }))
109       return failure();
110 
111     if (!inputTy.hasStaticShape() || !weightTy.hasStaticShape() ||
112         !biasTy.hasStaticShape() || !resultTy.hasStaticShape())
113       return failure();
114 
115     int64_t kernelHeight = (weightTy.getDimSize(1) - 1) * dilation[0] + 1;
116     int64_t kernelWidth = (weightTy.getDimSize(2) - 1) * dilation[1] + 1;
117     int64_t requiredInputHeight = resultTy.getDimSize(1) + kernelHeight - 1;
118     int64_t requiredInputWidth = resultTy.getDimSize(2) + kernelWidth - 1;
119 
120     llvm::SmallVector<int64_t> convPad(4, 0);
121     convPad[0] = kernelHeight - 1 - pad[0];
122     convPad[2] = kernelWidth - 1 - pad[1];
123     convPad[1] = requiredInputHeight - convPad[0] - inputTy.getDimSize(1);
124     convPad[3] = requiredInputWidth - convPad[2] - inputTy.getDimSize(2);
125 
126     auto reverse1 = rewriter.create<tosa::ReverseOp>(
127         loc, weightTy, weight, rewriter.getI64IntegerAttr(1));
128     auto reverse2 = rewriter.create<tosa::ReverseOp>(
129         loc, weightTy, reverse1, rewriter.getI64IntegerAttr(2));
130 
131     Value conv2d;
132     if (op.quantization_info().hasValue()) {
133       conv2d = rewriter.create<tosa::Conv2DOp>(
134           loc, resultTy, input, reverse2, bias,
135           rewriter.getI64ArrayAttr(convPad), rewriter.getI64ArrayAttr(stride),
136           rewriter.getI64ArrayAttr(dilation),
137           op.quantization_info().getValue());
138     } else {
139       conv2d = rewriter.create<tosa::Conv2DOp>(
140           loc, resultTy, input, reverse2, bias,
141           rewriter.getI64ArrayAttr(convPad), rewriter.getI64ArrayAttr(stride),
142           rewriter.getI64ArrayAttr(dilation));
143     }
144 
145     rewriter.replaceOp(op, conv2d);
146     return success();
147   }
148 };
149 
150 class TransposeConvStridedConverter
151     : public OpRewritePattern<tosa::TransposeConv2DOp> {
152 public:
153   using OpRewritePattern<tosa::TransposeConv2DOp>::OpRewritePattern;
154   LogicalResult matchAndRewrite(tosa::TransposeConv2DOp op,
155                                 PatternRewriter &rewriter) const final {
156     Location loc = op->getLoc();
157     Value input = op->getOperand(0);
158     Value weight = op->getOperand(1);
159     Value bias = op->getOperand(2);
160 
161     ShapedType inputTy = input.getType().cast<ShapedType>();
162     ShapedType weightTy = weight.getType().cast<ShapedType>();
163     ShapedType biasTy = bias.getType().cast<ShapedType>();
164     ShapedType resultTy = op->getResult(0).getType().cast<ShapedType>();
165 
166     Type inputETy = inputTy.getElementType();
167     Type weightETy = weightTy.getElementType();
168     Type biasETy = biasTy.getElementType();
169     Type resultETy = resultTy.getElementType();
170 
171     llvm::SmallVector<int64_t> pad;
172     llvm::SmallVector<int64_t> stride;
173     llvm::SmallVector<int64_t> dilation;
174 
175     getValuesFromIntArrayAttribute(op.out_pad().cast<ArrayAttr>(), pad);
176     getValuesFromIntArrayAttribute(op.stride().cast<ArrayAttr>(), stride);
177     getValuesFromIntArrayAttribute(op.dilation().cast<ArrayAttr>(), dilation);
178 
179     // If striding is all 1 we can modify padding and reverse the kernel along
180     // the x/y direction to make it a regular convolution. This is much simpler
181     // then handling striding....
182     if (llvm::any_of(dilation, [](int64_t v) { return v != 1; }))
183       return failure();
184 
185     // If strides are all 1 we dont need to use this one.
186     if (llvm::all_of(stride, [](int64_t v) { return v == 1; }))
187       return failure();
188 
189     if (!inputTy.hasStaticShape() || !weightTy.hasStaticShape() ||
190         !biasTy.hasStaticShape() || !resultTy.hasStaticShape())
191       return failure();
192 
193     int64_t batch = inputTy.getDimSize(0);
194 
195     int64_t outputChannels = weightTy.getDimSize(0);
196     int64_t weightHeight = weightTy.getDimSize(1);
197     int64_t weightWidth = weightTy.getDimSize(2);
198     int64_t inputChannels = weightTy.getDimSize(3);
199 
200     // Pad the weight so that it is modulo of the striding.
201     llvm::SmallVector<int32_t, 8> weightPadding = {0, 0, 0, 0, 0, 0, 0, 0};
202     weightPadding[3] =
203         weightHeight % stride[0] ? stride[0] - weightHeight % stride[0] : 0;
204     weightPadding[5] =
205         weightWidth % stride[1] ? stride[1] - weightWidth % stride[1] : 0;
206     DenseElementsAttr weightPaddingAttr = DenseIntElementsAttr::get(
207         RankedTensorType::get({4, 2}, rewriter.getI32Type()), weightPadding);
208     Value weightPaddingVal = CreateOpAndInfer<tosa::ConstOp>(
209         rewriter, loc, weightPaddingAttr.getType(), weightPaddingAttr);
210 
211     if (op.quantization_info().hasValue()) {
212       auto quantInfo = op.quantization_info().getValue();
213       weight = CreateOpAndInfer<tosa::PadOp>(
214           rewriter, loc, UnrankedTensorType::get(weightETy), weight,
215           weightPaddingVal, nullptr,
216           PadOpQuantizationAttr::get(quantInfo.weight_zp(),
217                                      rewriter.getContext()));
218 
219     } else {
220       weight = CreateOpAndInfer<tosa::PadOp>(rewriter, loc,
221                                              UnrankedTensorType::get(weightETy),
222                                              weight, weightPaddingVal);
223     }
224 
225     weightTy = weight.getType().cast<ShapedType>();
226     weightHeight = weightTy.getDimSize(1);
227     weightWidth = weightTy.getDimSize(2);
228 
229     // Split out the width / height by the stride dimensions.
230     llvm::SmallVector<int64_t, 6> weightReshapeDims0 = {
231         outputChannels, weightHeight / stride[0],
232         stride[0],      weightWidth / stride[1],
233         stride[1],      inputChannels};
234     weight = CreateOpAndInfer<tosa::ReshapeOp>(
235         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
236         rewriter.getI64ArrayAttr(weightReshapeDims0));
237 
238     // Transpose the factored-out stride to the output channels.
239     Value transposeWeightVal = rewriter.create<tosa::ConstOp>(
240         loc, RankedTensorType::get({6}, rewriter.getI32Type()),
241         rewriter.getI32TensorAttr({2, 4, 0, 1, 3, 5}));
242 
243     weight = CreateOpAndInfer<tosa::TransposeOp>(
244         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
245         transposeWeightVal);
246 
247     // Collapse the strides and output channels into a single dimension.
248     llvm::SmallVector<int64_t, 6> weightReshapeDims1 = {
249         outputChannels * stride[0] * stride[1], weightHeight / stride[0],
250         weightWidth / stride[1], inputChannels};
251     weight = CreateOpAndInfer<tosa::ReshapeOp>(
252         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
253         rewriter.getI64ArrayAttr(weightReshapeDims1));
254     ShapedType restridedWeightTy = weight.getType().cast<ShapedType>();
255 
256     weight = CreateOpAndInfer<tosa::ReverseOp>(
257         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
258         rewriter.getI64IntegerAttr(1));
259     weight = CreateOpAndInfer<tosa::ReverseOp>(
260         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
261         rewriter.getI64IntegerAttr(2));
262 
263     // We need to pad the input far enough that we can pull all values.
264     llvm::SmallVector<int32_t, 8> inputPadding = {0, 0, 0, 0, 0, 0, 0, 0};
265     inputPadding[2] += restridedWeightTy.getDimSize(1) - 1;
266     inputPadding[3] += restridedWeightTy.getDimSize(1) - 1;
267     inputPadding[4] += restridedWeightTy.getDimSize(2) - 1;
268     inputPadding[5] += restridedWeightTy.getDimSize(2) - 1;
269 
270     DenseElementsAttr inputPaddingAttr = DenseIntElementsAttr::get(
271         RankedTensorType::get({4, 2}, rewriter.getI32Type()), inputPadding);
272 
273     Value inputPaddingVal = CreateOpAndInfer<tosa::ConstOp>(
274         rewriter, loc, inputPaddingAttr.getType(), inputPaddingAttr);
275 
276     if (op.quantization_info().hasValue()) {
277       auto quantInfo = op.quantization_info().getValue();
278       input = CreateOpAndInfer<tosa::PadOp>(
279           rewriter, loc, UnrankedTensorType::get(inputETy), input,
280           inputPaddingVal, nullptr,
281           PadOpQuantizationAttr::get(quantInfo.input_zp(),
282                                      rewriter.getContext()));
283     } else {
284       input = CreateOpAndInfer<tosa::PadOp>(rewriter, loc,
285                                             UnrankedTensorType::get(inputETy),
286                                             input, inputPaddingVal);
287     }
288 
289     // We use a zero bias as we need to broadcast the bias.
290     auto zeroBias = rewriter.create<tosa::ConstOp>(
291         loc,
292         RankedTensorType::get({outputChannels * stride[0] * stride[1]},
293                               biasETy),
294         DenseElementsAttr::get(
295             RankedTensorType::get({outputChannels * stride[0] * stride[1]},
296                                   biasETy),
297             rewriter.getZeroAttr(biasETy)));
298 
299     // Perform the convolution using the zero bias.
300     Value conv2d;
301     if (op.quantization_info().hasValue()) {
302       conv2d = CreateOpAndInfer<tosa::Conv2DOp>(
303                    rewriter, loc, UnrankedTensorType::get(resultETy), input,
304                    weight, zeroBias,
305                    /*pad=*/rewriter.getI64ArrayAttr({0, 0, 0, 0}),
306                    /*stride=*/rewriter.getI64ArrayAttr({1, 1}),
307                    /*dilation=*/rewriter.getI64ArrayAttr({1, 1}),
308                    op.quantization_info().getValue())
309                    .getResult();
310     } else {
311       conv2d = CreateOpAndInfer<tosa::Conv2DOp>(
312                    rewriter, loc, UnrankedTensorType::get(resultETy), input,
313                    weight, zeroBias,
314                    /*pad=*/rewriter.getI64ArrayAttr({0, 0, 0, 0}),
315                    /*stride=*/rewriter.getI64ArrayAttr({1, 1}),
316                    /*dilation=*/rewriter.getI64ArrayAttr({1, 1}))
317                    .getResult();
318     }
319 
320     // Factor the resulting width / height.
321     ShapedType convTy = conv2d.getType().cast<ShapedType>();
322     Type convETy = convTy.getElementType();
323 
324     int64_t convHeight = convTy.getDimSize(1);
325     int64_t convWidth = convTy.getDimSize(2);
326 
327     // Factor striding out of the convolution result.
328     llvm::SmallVector<int64_t, 6> convReshapeDims0 = {
329         batch, convHeight, convWidth, stride[0], stride[1], outputChannels};
330     conv2d = CreateOpAndInfer<tosa::ReshapeOp>(
331         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
332         rewriter.getI64ArrayAttr(convReshapeDims0));
333 
334     // Transpose the factored-out stride to the output channels.
335     Value transposeConvVal = rewriter.create<tosa::ConstOp>(
336         loc, RankedTensorType::get({6}, rewriter.getI32Type()),
337         rewriter.getI32TensorAttr({0, 1, 3, 2, 4, 5}));
338 
339     conv2d = CreateOpAndInfer<tosa::TransposeOp>(
340         rewriter, loc, UnrankedTensorType::get(convETy), conv2d,
341         transposeConvVal);
342 
343     // Fuse striding behavior back into width / height.
344     llvm::SmallVector<int64_t, 6> convReshapeDims1 = {
345         batch, convHeight * stride[0], convWidth * stride[1], outputChannels};
346     conv2d = CreateOpAndInfer<tosa::ReshapeOp>(
347         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
348         rewriter.getI64ArrayAttr(convReshapeDims1));
349 
350     // Slice out the final result.
351     llvm::SmallVector<int64_t, 4> sliceBegin = {0, 0, 0, 0};
352     llvm::SmallVector<int64_t, 4> sliceSize(resultTy.getShape().begin(),
353                                             resultTy.getShape().begin());
354     sliceBegin[1] = pad[0];
355     sliceBegin[2] = pad[1];
356 
357     auto slice = CreateOpAndInfer<tosa::SliceOp>(
358                      rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
359                      rewriter.getI64ArrayAttr(sliceBegin),
360                      rewriter.getI64ArrayAttr(resultTy.getShape()))
361                      .getResult();
362 
363     auto addBias =
364         CreateOpAndInfer<tosa::AddOp>(rewriter, loc, op.getType(), slice, bias);
365 
366     rewriter.replaceOp(op, addBias.getResult());
367 
368     return success();
369   }
370 };
371 
372 /// Pass that enables broadcast by making all input arrays have the same
373 /// number of dimensions. Insert RESHAPE operations to lower rank operand
374 struct TosaDecomposeTransposeConv
375     : public TosaDecomposeTransposeConvBase<TosaDecomposeTransposeConv> {
376 public:
377   void runOnFunction() override {
378     auto func = getFunction();
379     RewritePatternSet patterns(func.getContext());
380     patterns
381         .insert<TransposeConvDilatedConverter, TransposeConvStridedConverter>(
382             func.getContext());
383     (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
384   }
385 };
386 } // end anonymous namespace
387 
388 std::unique_ptr<Pass> mlir::tosa::createTosaDecomposeTransposeConvPass() {
389   return std::make_unique<TosaDecomposeTransposeConv>();
390 }
391