1 //===- TosaDecomposeTransposeConv.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 // Decompose TOSA TransposeConv operation to a series of TOSA Ops specifically
10 // (1) Convert a Dilated TransposeConv2D to Conv2D including reversing/reshaping
11 // etc.. of the weights (2) Convert a Strided TransposeConv2D to Conv2D
12 // including transposing/reversing/reshaping etc..
13 //     of the weights and input/output tenors and reversing/reshaping etc .. of
14 //     the weights
15 //
16 //===----------------------------------------------------------------------===//
17 
18 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
19 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
20 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
21 #include "mlir/Pass/Pass.h"
22 
23 using namespace mlir;
24 using namespace mlir::tosa;
25 
26 namespace {
27 
28 template <typename T>
29 static void getValuesFromIntArrayAttribute(ArrayAttr attr,
30                                            SmallVector<T> &arrayValues) {
31   for (Attribute val : attr.getValue()) {
32     arrayValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
33   }
34 }
35 
36 template <typename TosaOp, typename... Args>
37 TosaOp createOpAndInfer(PatternRewriter &rewriter, Location loc, Type resultTy,
38                         Args &&...args) {
39   auto op = rewriter.create<TosaOp>(loc, resultTy, args...);
40 
41   InferShapedTypeOpInterface shapeInterface =
42       dyn_cast<InferShapedTypeOpInterface>(op.getOperation());
43   if (!shapeInterface)
44     return op;
45 
46   SmallVector<ShapedTypeComponents> returnedShapes;
47   if (shapeInterface
48           .inferReturnTypeComponents(op.getContext(), op.getLoc(),
49                                      op->getOperands(), op->getAttrDictionary(),
50                                      op->getRegions(), returnedShapes)
51           .failed())
52     return op;
53 
54   // We need to use the element type of the existing result type to generate
55   // the new result shaped type. This is because rescale can include a cast to
56   // different bit-width types and does not have a TypeAttr to define the
57   // target type.
58   auto result = op->getResult(0);
59   auto predictedShape = returnedShapes[0];
60   auto currentKnowledge =
61       mlir::tosa::ValueKnowledge::getKnowledgeFromType(resultTy);
62 
63   // Compute the knowledge based on the inferred type.
64   auto inferredKnowledge =
65       mlir::tosa::ValueKnowledge::getPessimisticValueState();
66   inferredKnowledge.dtype = resultTy.cast<ShapedType>().getElementType();
67   inferredKnowledge.hasRank = predictedShape.hasRank();
68   if (predictedShape.hasRank()) {
69     for (auto dim : predictedShape.getDims()) {
70       inferredKnowledge.sizes.push_back(dim);
71     }
72   }
73 
74   // Compute the new type based on the joined version.
75   auto newKnowledge =
76       mlir::tosa::ValueKnowledge::join(currentKnowledge, inferredKnowledge);
77   auto newTy = newKnowledge.getType();
78   result.setType(newTy);
79   return op;
80 }
81 
82 class TransposeConvDilatedConverter
83     : public OpRewritePattern<tosa::TransposeConv2DOp> {
84 public:
85   using OpRewritePattern<tosa::TransposeConv2DOp>::OpRewritePattern;
86   LogicalResult matchAndRewrite(tosa::TransposeConv2DOp op,
87                                 PatternRewriter &rewriter) const final {
88     Location loc = op->getLoc();
89     Value input = op->getOperand(0);
90     Value weight = op->getOperand(1);
91     Value bias = op->getOperand(2);
92 
93     ShapedType inputTy = input.getType().cast<ShapedType>();
94     ShapedType weightTy = weight.getType().cast<ShapedType>();
95     ShapedType biasTy = bias.getType().cast<ShapedType>();
96     ShapedType resultTy = op->getResult(0).getType().cast<ShapedType>();
97 
98     llvm::SmallVector<int64_t> pad;
99     llvm::SmallVector<int64_t> stride;
100     llvm::SmallVector<int64_t> dilation;
101 
102     getValuesFromIntArrayAttribute(op.out_pad().cast<ArrayAttr>(), pad);
103     getValuesFromIntArrayAttribute(op.stride().cast<ArrayAttr>(), stride);
104     getValuesFromIntArrayAttribute(op.dilation().cast<ArrayAttr>(), dilation);
105 
106     // If striding is all 1 we can modify padding and reverse the kernel along
107     // the x/y direction to make it a regular convolution. This is much simpler
108     // then handling striding....
109     if (llvm::any_of(stride, [](int64_t v) { return v != 1; }))
110       return failure();
111 
112     if (!inputTy.hasStaticShape() || !weightTy.hasStaticShape() ||
113         !biasTy.hasStaticShape() || !resultTy.hasStaticShape())
114       return failure();
115 
116     int64_t kernelHeight = (weightTy.getDimSize(1) - 1) * dilation[0] + 1;
117     int64_t kernelWidth = (weightTy.getDimSize(2) - 1) * dilation[1] + 1;
118     int64_t requiredInputHeight = resultTy.getDimSize(1) + kernelHeight - 1;
119     int64_t requiredInputWidth = resultTy.getDimSize(2) + kernelWidth - 1;
120 
121     llvm::SmallVector<int64_t> convPad(4, 0);
122     convPad[0] = kernelHeight - 1 - pad[0];
123     convPad[2] = kernelWidth - 1 - pad[1];
124     convPad[1] = requiredInputHeight - convPad[0] - inputTy.getDimSize(1);
125     convPad[3] = requiredInputWidth - convPad[2] - inputTy.getDimSize(2);
126 
127     auto reverse1 = rewriter.create<tosa::ReverseOp>(
128         loc, weightTy, weight, rewriter.getI64IntegerAttr(1));
129     auto reverse2 = rewriter.create<tosa::ReverseOp>(
130         loc, weightTy, reverse1, rewriter.getI64IntegerAttr(2));
131 
132     Value conv2d;
133     if (op.quantization_info()) {
134       conv2d = rewriter.create<tosa::Conv2DOp>(
135           loc, resultTy, input, reverse2, bias,
136           rewriter.getI64ArrayAttr(convPad), rewriter.getI64ArrayAttr(stride),
137           rewriter.getI64ArrayAttr(dilation), *op.quantization_info());
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           rewriter.getAttr<PadOpQuantizationAttr>(quantInfo.getWeightZp()));
217 
218     } else {
219       weight = createOpAndInfer<tosa::PadOp>(rewriter, loc,
220                                              UnrankedTensorType::get(weightETy),
221                                              weight, weightPaddingVal);
222     }
223 
224     weightTy = weight.getType().cast<ShapedType>();
225     weightHeight = weightTy.getDimSize(1);
226     weightWidth = weightTy.getDimSize(2);
227 
228     // Split out the width / height by the stride dimensions.
229     llvm::SmallVector<int64_t, 6> weightReshapeDims0 = {
230         outputChannels, weightHeight / stride[0],
231         stride[0],      weightWidth / stride[1],
232         stride[1],      inputChannels};
233     weight = createOpAndInfer<tosa::ReshapeOp>(
234         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
235         rewriter.getI64ArrayAttr(weightReshapeDims0));
236 
237     // Transpose the factored-out stride to the output channels.
238     Value transposeWeightVal = rewriter.create<tosa::ConstOp>(
239         loc, RankedTensorType::get({6}, rewriter.getI32Type()),
240         rewriter.getI32TensorAttr({2, 4, 0, 1, 3, 5}));
241 
242     weight = createOpAndInfer<tosa::TransposeOp>(
243         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
244         transposeWeightVal);
245 
246     // Collapse the strides and output channels into a single dimension.
247     llvm::SmallVector<int64_t, 6> weightReshapeDims1 = {
248         outputChannels * stride[0] * stride[1], weightHeight / stride[0],
249         weightWidth / stride[1], inputChannels};
250     weight = createOpAndInfer<tosa::ReshapeOp>(
251         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
252         rewriter.getI64ArrayAttr(weightReshapeDims1));
253     ShapedType restridedWeightTy = weight.getType().cast<ShapedType>();
254 
255     weight = createOpAndInfer<tosa::ReverseOp>(
256         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
257         rewriter.getI64IntegerAttr(1));
258     weight = createOpAndInfer<tosa::ReverseOp>(
259         rewriter, loc, UnrankedTensorType::get(weightETy), weight,
260         rewriter.getI64IntegerAttr(2));
261 
262     // We need to pad the input far enough that we can pull all values.
263     llvm::SmallVector<int32_t, 8> inputPadding = {0, 0, 0, 0, 0, 0, 0, 0};
264     inputPadding[2] += restridedWeightTy.getDimSize(1) - 1;
265     inputPadding[3] += restridedWeightTy.getDimSize(1) - 1;
266     inputPadding[4] += restridedWeightTy.getDimSize(2) - 1;
267     inputPadding[5] += restridedWeightTy.getDimSize(2) - 1;
268 
269     DenseElementsAttr inputPaddingAttr = DenseIntElementsAttr::get(
270         RankedTensorType::get({4, 2}, rewriter.getI32Type()), inputPadding);
271 
272     Value inputPaddingVal = createOpAndInfer<tosa::ConstOp>(
273         rewriter, loc, inputPaddingAttr.getType(), inputPaddingAttr);
274 
275     if (op.quantization_info().hasValue()) {
276       auto quantInfo = op.quantization_info().getValue();
277       input = createOpAndInfer<tosa::PadOp>(
278           rewriter, loc, UnrankedTensorType::get(inputETy), input,
279           inputPaddingVal, nullptr,
280           rewriter.getAttr<PadOpQuantizationAttr>(quantInfo.getInputZp()));
281     } else {
282       input = createOpAndInfer<tosa::PadOp>(rewriter, loc,
283                                             UnrankedTensorType::get(inputETy),
284                                             input, inputPaddingVal);
285     }
286 
287     // We use a zero bias as we need to broadcast the bias.
288     auto zeroBias = rewriter.create<tosa::ConstOp>(
289         loc,
290         RankedTensorType::get({outputChannels * stride[0] * stride[1]},
291                               biasETy),
292         DenseElementsAttr::get(
293             RankedTensorType::get({outputChannels * stride[0] * stride[1]},
294                                   biasETy),
295             rewriter.getZeroAttr(biasETy)));
296 
297     // Perform the convolution using the zero bias.
298     Value conv2d;
299     if (op.quantization_info()) {
300       conv2d = createOpAndInfer<tosa::Conv2DOp>(
301                    rewriter, loc, UnrankedTensorType::get(resultETy), input,
302                    weight, zeroBias,
303                    /*pad=*/rewriter.getI64ArrayAttr({0, 0, 0, 0}),
304                    /*stride=*/rewriter.getI64ArrayAttr({1, 1}),
305                    /*dilation=*/rewriter.getI64ArrayAttr({1, 1}),
306                    *op.quantization_info())
307                    .getResult();
308     } else {
309       conv2d = createOpAndInfer<tosa::Conv2DOp>(
310                    rewriter, loc, UnrankedTensorType::get(resultETy), input,
311                    weight, zeroBias,
312                    /*pad=*/rewriter.getI64ArrayAttr({0, 0, 0, 0}),
313                    /*stride=*/rewriter.getI64ArrayAttr({1, 1}),
314                    /*dilation=*/rewriter.getI64ArrayAttr({1, 1}))
315                    .getResult();
316     }
317 
318     // Factor the resulting width / height.
319     ShapedType convTy = conv2d.getType().cast<ShapedType>();
320     Type convETy = convTy.getElementType();
321 
322     int64_t convHeight = convTy.getDimSize(1);
323     int64_t convWidth = convTy.getDimSize(2);
324 
325     // Factor striding out of the convolution result.
326     llvm::SmallVector<int64_t, 6> convReshapeDims0 = {
327         batch, convHeight, convWidth, stride[0], stride[1], outputChannels};
328     conv2d = createOpAndInfer<tosa::ReshapeOp>(
329         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
330         rewriter.getI64ArrayAttr(convReshapeDims0));
331 
332     // Transpose the factored-out stride to the output channels.
333     Value transposeConvVal = rewriter.create<tosa::ConstOp>(
334         loc, RankedTensorType::get({6}, rewriter.getI32Type()),
335         rewriter.getI32TensorAttr({0, 1, 3, 2, 4, 5}));
336 
337     conv2d = createOpAndInfer<tosa::TransposeOp>(
338         rewriter, loc, UnrankedTensorType::get(convETy), conv2d,
339         transposeConvVal);
340 
341     // Fuse striding behavior back into width / height.
342     llvm::SmallVector<int64_t, 6> convReshapeDims1 = {
343         batch, convHeight * stride[0], convWidth * stride[1], outputChannels};
344     conv2d = createOpAndInfer<tosa::ReshapeOp>(
345         rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
346         rewriter.getI64ArrayAttr(convReshapeDims1));
347 
348     // Slice out the final result.
349     llvm::SmallVector<int64_t, 4> sliceBegin = {0, 0, 0, 0};
350     llvm::SmallVector<int64_t, 4> sliceSize(resultTy.getShape().begin(),
351                                             resultTy.getShape().begin());
352     sliceBegin[1] = pad[0];
353     sliceBegin[2] = pad[1];
354 
355     auto slice = createOpAndInfer<tosa::SliceOp>(
356                      rewriter, loc, UnrankedTensorType::get(resultETy), conv2d,
357                      rewriter.getI64ArrayAttr(sliceBegin),
358                      rewriter.getI64ArrayAttr(resultTy.getShape()))
359                      .getResult();
360 
361     auto addBias =
362         createOpAndInfer<tosa::AddOp>(rewriter, loc, op.getType(), slice, bias);
363 
364     rewriter.replaceOp(op, addBias.getResult());
365 
366     return success();
367   }
368 };
369 
370 } // namespace
371 
372 void mlir::tosa::populateTosaDecomposeTransposeConv(
373     MLIRContext *ctx, RewritePatternSet &patterns) {
374   patterns.add<TransposeConvDilatedConverter>(ctx);
375   patterns.add<TransposeConvStridedConverter>(ctx);
376 }
377