1 | //===- TosaDecomposeDepthwise.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 Depthwise operation to a series of TOSA Ops specifically |
10 | // (1) Convert a 1x1 Depthwise to Reshape -> Mul -> Reshape -> Add |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "mlir/Dialect/Tosa/IR/TosaOps.h" |
15 | #include "mlir/Dialect/Tosa/Transforms/Passes.h" |
16 | #include "mlir/Dialect/Tosa/Utils/ConversionUtils.h" |
17 | #include "mlir/IR/BuiltinTypes.h" |
18 | #include "mlir/Pass/Pass.h" |
19 | |
20 | using namespace mlir; |
21 | using namespace mlir::tosa; |
22 | |
23 | namespace { |
24 | |
25 | struct DepthwiseConv2DIsMul : public OpRewritePattern<tosa::DepthwiseConv2DOp> { |
26 | explicit DepthwiseConv2DIsMul(MLIRContext *context) |
27 | : OpRewritePattern(context) {} |
28 | |
29 | LogicalResult matchAndRewrite(tosa::DepthwiseConv2DOp op, |
30 | PatternRewriter &rewriter) const override { |
31 | Value input = op.getInput(); |
32 | Value weight = op.getWeight(); |
33 | ShapedType inputType = cast<ShapedType>(input.getType()); |
34 | ShapedType weightType = cast<ShapedType>(weight.getType()); |
35 | ShapedType resultType = cast<ShapedType>(op.getOutput().getType()); |
36 | |
37 | if (!(inputType.hasStaticShape() && weightType.hasStaticShape() && |
38 | resultType.hasStaticShape())) { |
39 | return failure(); |
40 | } |
41 | |
42 | if (!llvm::all_of(op.getStride(), [](int64_t v) { return v == 1; })) |
43 | return failure(); |
44 | |
45 | // Only works for a 1x1 kernel. |
46 | ArrayRef<int64_t> weightShape = weightType.getShape(); |
47 | if (weightShape[0] != 1 || weightShape[1] != 1) { |
48 | return failure(); |
49 | } |
50 | |
51 | Type inputETy = inputType.getElementType(); |
52 | Type weightETy = weightType.getElementType(); |
53 | if (!inputETy.isIntOrFloat() || !weightETy.isIntOrFloat()) |
54 | return rewriter.notifyMatchFailure(op, "unsupported type" ); |
55 | |
56 | // Get and verify zero points. |
57 | FailureOr<int64_t> maybeIZp = op.getInputZeroPoint(); |
58 | if (failed(Result: maybeIZp)) |
59 | return rewriter.notifyMatchFailure( |
60 | op, "input zero point cannot be statically determined" ); |
61 | |
62 | FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint(); |
63 | if (failed(Result: maybeWZp)) |
64 | return rewriter.notifyMatchFailure( |
65 | op, "weight zero point cannot be statically determined" ); |
66 | |
67 | int64_t iZp = *maybeIZp; |
68 | int64_t wZp = *maybeWZp; |
69 | if (op.verifyInputZeroPoint(iZp).failed()) |
70 | return rewriter.notifyMatchFailure( |
71 | op, "input zero point must be zero for non-int8 integer types" ); |
72 | if (op.verifyWeightZeroPoint(wZp).failed()) |
73 | return rewriter.notifyMatchFailure( |
74 | op, "weight zero point must be zero for non-int8 integer types" ); |
75 | |
76 | // Reshape input to [N, H, W, C] -> [N, H, W, C, 1]. |
77 | ArrayRef<int64_t> inputShape = inputType.getShape(); |
78 | llvm::SmallVector<int64_t, 2> revisedInputShape{ |
79 | inputShape[0], inputShape[1], inputShape[2], inputShape[3], 1}; |
80 | inputType = RankedTensorType::get( |
81 | revisedInputShape, |
82 | dyn_cast<RankedTensorType>(input.getType()).getElementType()); |
83 | auto revisedInputShapeValue = |
84 | getTosaConstShape(rewriter, op.getLoc(), revisedInputShape); |
85 | input = rewriter |
86 | .create<tosa::ReshapeOp>(op.getLoc(), inputType, input, |
87 | revisedInputShapeValue) |
88 | .getResult(); |
89 | |
90 | Type resultETy = resultType.getElementType(); |
91 | |
92 | if (inputETy != resultETy) { |
93 | inputType = inputType.clone(resultETy); |
94 | input = rewriter.create<tosa::CastOp>(op.getLoc(), inputType, input); |
95 | } |
96 | |
97 | if (weightETy != resultETy) { |
98 | weightType = weightType.clone(resultETy); |
99 | weight = rewriter.create<tosa::CastOp>(op.getLoc(), weightType, weight); |
100 | } |
101 | |
102 | if (iZp != 0 || wZp != 0) { |
103 | |
104 | auto applyZp = [&](Value val, int64_t zp) -> Value { |
105 | if (zp == 0) |
106 | return val; |
107 | auto ety = cast<ShapedType>(val.getType()).getElementType(); |
108 | std::vector<int64_t> shape(cast<ShapedType>(val.getType()).getRank(), |
109 | 1); |
110 | auto zpTy = RankedTensorType::get(shape, ety); |
111 | auto zpAttr = |
112 | DenseElementsAttr::get(zpTy, rewriter.getIntegerAttr(ety, zp)); |
113 | auto zpVal = rewriter.create<tosa::ConstOp>(op.getLoc(), zpTy, zpAttr); |
114 | return rewriter.create<tosa::SubOp>(op.getLoc(), val.getType(), val, |
115 | zpVal); |
116 | }; |
117 | |
118 | input = applyZp(input, iZp); |
119 | weight = applyZp(weight, wZp); |
120 | } |
121 | |
122 | ArrayRef<int64_t> padAttr = op.getPad(); |
123 | llvm::SmallVector<int64_t> pad(10, 0); |
124 | for (const auto &it : llvm::enumerate(padAttr)) |
125 | pad[it.index() + 2] = it.value(); |
126 | |
127 | if (llvm::any_of(Range&: pad, P: [](int64_t p) { return p != 0; })) { |
128 | Attribute zeroAttr = rewriter.getZeroAttr(inputETy); |
129 | |
130 | llvm::SmallVector<int64_t> newShape(inputType.getShape()); |
131 | for (int i = 0, s = pad.size(); i < s; ++i) { |
132 | if (newShape[i / 2] != ShapedType::kDynamic) { |
133 | newShape[i / 2] += pad[i]; |
134 | } |
135 | } |
136 | |
137 | Value padSizeVal = getTosaConstShape(rewriter, op->getLoc(), pad); |
138 | |
139 | auto padTy = RankedTensorType::get({1}, inputETy); |
140 | auto padAttr = DenseElementsAttr::get(padTy, zeroAttr); |
141 | Value padVal = |
142 | rewriter.create<tosa::ConstOp>(op->getLoc(), padTy, padAttr); |
143 | inputType = RankedTensorType::get(newShape, inputETy); |
144 | input = rewriter.create<tosa::PadOp>(op->getLoc(), inputType, input, |
145 | padSizeVal, padVal); |
146 | } |
147 | |
148 | // Perform an elementwise mul over the reshaped input and weight. |
149 | llvm::SmallVector<int64_t, 2> mulShape{ |
150 | inputType.getDimSize(0), inputType.getDimSize(1), |
151 | inputType.getDimSize(2), inputType.getDimSize(3), weightShape[3]}; |
152 | auto mulShapeType = RankedTensorType::get( |
153 | mulShape, |
154 | dyn_cast<RankedTensorType>(weight.getType()).getElementType()); |
155 | |
156 | if (EqualizeRanks(rewriter, op.getLoc(), input, weight).failed()) { |
157 | return failure(); |
158 | } |
159 | |
160 | auto shiftElementType = IntegerType::get(rewriter.getContext(), 8); |
161 | auto shiftType = RankedTensorType::get({1}, shiftElementType); |
162 | auto shiftZeroAttr = DenseElementsAttr::get( |
163 | shiftType, rewriter.getIntegerAttr(shiftElementType, 0)); |
164 | Value constZero = |
165 | rewriter.create<tosa::ConstOp>(op.getLoc(), shiftType, shiftZeroAttr); |
166 | Value mulValue = rewriter |
167 | .create<tosa::MulOp>(op.getLoc(), mulShapeType, input, |
168 | weight, constZero) |
169 | .getResult(); |
170 | |
171 | // Reshape output to [N, H, W, C * M]. |
172 | auto outputShape = cast<ShapedType>(op.getOutput().getType()).getShape(); |
173 | auto outputShapeType = RankedTensorType::get( |
174 | outputShape, |
175 | dyn_cast<RankedTensorType>(input.getType()).getElementType()); |
176 | auto outputShapeValue = |
177 | getTosaConstShape(rewriter, op->getLoc(), outputShape); |
178 | Value outputValue = rewriter.create<tosa::ReshapeOp>( |
179 | op.getLoc(), outputShapeType, mulValue, outputShapeValue); |
180 | |
181 | Value bias = op.getBias(); |
182 | if (EqualizeRanks(rewriter, op.getLoc(), outputValue, bias).failed()) { |
183 | return failure(); |
184 | } |
185 | |
186 | // Add in the bias. |
187 | rewriter |
188 | .replaceOpWithNewOp<tosa::AddOp>(op, outputShapeType, outputValue, bias) |
189 | .getResult(); |
190 | return success(); |
191 | } |
192 | }; |
193 | |
194 | } // namespace |
195 | |
196 | void mlir::tosa::populateTosaDecomposeDepthwise(MLIRContext *ctx, |
197 | RewritePatternSet &patterns) { |
198 | patterns.add<DepthwiseConv2DIsMul>(arg&: ctx); |
199 | } |
200 | |