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/Tensor/IR/Tensor.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 {
32   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
39 ConvertTosaNegateOp::matchAndRewrite(Operation *op,
40                                      PatternRewriter &rewriter) const {
41 
42   auto tosaNegateOp = cast<tosa::NegateOp>(op);
43 
44   auto inputType =
45       tosaNegateOp.getInput1().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 narrowRange = outputElementType.getStorageTypeMin() == 1;
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(narrowRange)));
84 
85   ElementsAttr inputElems;
86   if (!matchPattern(tosaNegateOp.getInput1(), 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 {
102   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
109 ConvertTosaConv2DOp::matchAndRewrite(Operation *op,
110                                      PatternRewriter &rewriter) const {
111 
112   auto tosaConv2DOp = cast<tosa::Conv2DOp>(op);
113 
114   auto inputType =
115       tosaConv2DOp.getInput().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.getWeight().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.getInput(),
150       tosaConv2DOp.getWeight(), tosaConv2DOp.getBias(), tosaConv2DOp.getPad(),
151       tosaConv2DOp.getStride(), tosaConv2DOp.getDilation());
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, OperationPass<func::FuncOp>> {
182   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TosaTestQuantUtilAPI)
183 
184   StringRef getArgument() const final { return PASS_NAME; }
185   StringRef getDescription() const final {
186     return "TOSA Test: Exercise the APIs in QuantUtils.cpp.";
187   }
188   void runOnOperation() override;
189 };
190 
191 void TosaTestQuantUtilAPI::runOnOperation() {
192   auto *ctx = &getContext();
193   RewritePatternSet patterns(ctx);
194   auto func = getOperation();
195 
196   patterns.add<ConvertTosaNegateOp>(ctx);
197   patterns.add<ConvertTosaConv2DOp>(ctx);
198   (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
199 }
200 
201 } // namespace
202 
203 namespace mlir {
204 void registerTosaTestQuantUtilAPIPass() {
205   PassRegistration<TosaTestQuantUtilAPI>();
206 }
207 } // namespace mlir
208