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

source code of mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeDepthwise.cpp