1 //===- TosaTestPasses.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 // Test passes to exercise TOSA helper functions.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "mlir/Dialect/StandardOps/IR/Ops.h"
14 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
15 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
16 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
17 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
18 #include "mlir/IR/BuiltinTypes.h"
19 #include "mlir/IR/Matchers.h"
20 #include "mlir/Pass/Pass.h"
21 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
22
23 #define PASS_NAME "tosa-test-quant-utils"
24
25 using namespace mlir;
26 using namespace mlir::tosa;
27
28 // This transformation converts quantized uint8 to quantized int8. The
29 // construction of the new type invokes buildQTypeFromMinMax. Extracted from
30 // TOSA legalization infrastructure.
31 struct ConvertTosaNegateOp : public RewritePattern {
ConvertTosaNegateOpConvertTosaNegateOp32 explicit ConvertTosaNegateOp(MLIRContext *context)
33 : RewritePattern(tosa::NegateOp::getOperationName(), 1, context) {}
34 LogicalResult matchAndRewrite(Operation *op,
35 PatternRewriter &rewriter) const override;
36 };
37
38 LogicalResult
matchAndRewrite(Operation * op,PatternRewriter & rewriter) const39 ConvertTosaNegateOp::matchAndRewrite(Operation *op,
40 PatternRewriter &rewriter) const {
41
42 auto tosaNegateOp = cast<tosa::NegateOp>(op);
43
44 auto inputType =
45 tosaNegateOp.input1().getType().dyn_cast<mlir::RankedTensorType>();
46 // skip if input is not ranked tensor type
47 if (!inputType)
48 return failure();
49
50 // skip if it's not ranked tensor type.
51 auto outputType =
52 tosaNegateOp.getResult().getType().dyn_cast<mlir::RankedTensorType>();
53 if (!outputType)
54 return failure();
55
56 // skip if output is not per-tensor quantized type.
57 auto outputElementType =
58 outputType.getElementType().dyn_cast<mlir::quant::UniformQuantizedType>();
59 if (!outputElementType)
60 return failure();
61
62 // skip if output is not uint8.
63 if (outputElementType.isSigned() ||
64 outputElementType.getStorageTypeIntegralWidth() != 8)
65 return failure();
66
67 double typeRangeMin = double(outputElementType.getStorageTypeMin() -
68 outputElementType.getZeroPoint()) *
69 outputElementType.getScale();
70 double typeRangeMax = double(outputElementType.getStorageTypeMax() -
71 outputElementType.getZeroPoint()) *
72 outputElementType.getScale();
73 bool narrow_range = outputElementType.getStorageTypeMin() == 1 ? true : false;
74
75 auto dstQConstType = RankedTensorType::get(
76 outputType.getShape(),
77 buildQTypeFromMinMax(rewriter, outputElementType.getExpressedType(),
78 rewriter.getF64FloatAttr(typeRangeMin),
79 rewriter.getF64FloatAttr(typeRangeMax),
80 rewriter.getI32IntegerAttr(
81 outputElementType.getStorageTypeIntegralWidth()),
82 0, true /* signed */,
83 rewriter.getBoolAttr(narrow_range)));
84
85 ElementsAttr inputElems;
86 if (!matchPattern(tosaNegateOp.input1(), m_Constant(&inputElems)))
87 return failure();
88
89 auto newConstOp =
90 rewriter.create<tosa::ConstOp>(op->getLoc(), dstQConstType, inputElems);
91 auto newNegateOp = rewriter.create<tosa::NegateOp>(
92 op->getLoc(), dstQConstType, newConstOp.getResult());
93
94 rewriter.replaceOp(op, {newNegateOp.getResult()});
95 return success();
96 }
97
98 // This transformation modifies the quantized output of a test conv2d input and
99 // appends a TOSA rescale after it. The rescale op requires the invocation of
100 // computeMultiplierAndShift. From TOSA legalization infrastructure.
101 struct ConvertTosaConv2DOp : public RewritePattern {
ConvertTosaConv2DOpConvertTosaConv2DOp102 explicit ConvertTosaConv2DOp(MLIRContext *context)
103 : RewritePattern(tosa::Conv2DOp::getOperationName(), 1, context) {}
104 LogicalResult matchAndRewrite(Operation *op,
105 PatternRewriter &rewriter) const override;
106 };
107
108 LogicalResult
matchAndRewrite(Operation * op,PatternRewriter & rewriter) const109 ConvertTosaConv2DOp::matchAndRewrite(Operation *op,
110 PatternRewriter &rewriter) const {
111
112 auto tosaConv2DOp = cast<tosa::Conv2DOp>(op);
113
114 auto inputType =
115 tosaConv2DOp.input().getType().dyn_cast<mlir::RankedTensorType>();
116
117 // skip if input is not ranked tensor type
118 if (!inputType)
119 return failure();
120
121 auto weightType =
122 tosaConv2DOp.weight().getType().dyn_cast<mlir::RankedTensorType>();
123
124 // skip if wt is not ranked tensor type
125 if (!weightType)
126 return failure();
127
128 // skip if it's not ranked tensor type.
129 auto outputType =
130 tosaConv2DOp.getResult().getType().dyn_cast<mlir::RankedTensorType>();
131 if (!outputType)
132 return failure();
133
134 auto inputQType =
135 inputType.getElementType().dyn_cast<mlir::quant::UniformQuantizedType>();
136 auto weightQType =
137 weightType.getElementType().dyn_cast<mlir::quant::UniformQuantizedType>();
138 auto outputQType =
139 outputType.getElementType().dyn_cast<mlir::quant::UniformQuantizedType>();
140
141 // Works on quantized type only.
142 if (!(inputQType && weightQType && outputQType))
143 return failure();
144
145 auto newTosaConv2DOpType =
146 RankedTensorType::get(outputType.getShape(), rewriter.getIntegerType(32));
147
148 auto newTosaConv2DOp = rewriter.create<tosa::Conv2DOp>(
149 op->getLoc(), newTosaConv2DOpType, tosaConv2DOp.input(),
150 tosaConv2DOp.weight(), tosaConv2DOp.bias(), tosaConv2DOp.pad(),
151 tosaConv2DOp.stride(), tosaConv2DOp.dilation());
152
153 // Create rescale to quantized type
154 double inputScale = inputQType.getScale();
155 double weightScale = weightQType.getScale();
156 double outputScale = outputQType.getScale();
157 int64_t outputZp = outputQType.getZeroPoint();
158
159 double opTensorScale = (inputScale * weightScale) / outputScale;
160
161 int32_t multiplier;
162 int32_t shift;
163
164 // Obtain the quantized scale = multiplier and shift.
165 computeMultiplierAndShift(opTensorScale, multiplier, shift, 32);
166
167 auto newTosaRescaleOp = rewriter.create<tosa::RescaleOp>(
168 op->getLoc(), outputType, newTosaConv2DOp.getResult(),
169 rewriter.getI32IntegerAttr(0), rewriter.getI32IntegerAttr(outputZp),
170 rewriter.getI32ArrayAttr({multiplier}), rewriter.getI32ArrayAttr({shift}),
171 rewriter.getBoolAttr(true), rewriter.getBoolAttr(true),
172 rewriter.getBoolAttr(false));
173
174 rewriter.replaceOp(op, {newTosaRescaleOp.getResult()});
175 return success();
176 }
177
178 namespace {
179
180 struct TosaTestQuantUtilAPI
181 : public PassWrapper<TosaTestQuantUtilAPI, FunctionPass> {
182 void runOnFunction() override;
183 };
184
runOnFunction()185 void TosaTestQuantUtilAPI::runOnFunction() {
186 OwningRewritePatternList patterns;
187 auto *ctx = &getContext();
188 auto func = getFunction();
189
190 patterns.insert<ConvertTosaNegateOp>(ctx);
191 patterns.insert<ConvertTosaConv2DOp>(ctx);
192 applyPatternsAndFoldGreedily(func, std::move(patterns));
193 }
194
195 } // anonymous namespace
196
197 namespace mlir {
registerTosaTestQuantUtilAPIPass()198 void registerTosaTestQuantUtilAPIPass() {
199 PassRegistration<TosaTestQuantUtilAPI>(
200 PASS_NAME, "TOSA Test: Exercise the APIs in QuantUtils.cpp.");
201 }
202 } // namespace mlir
203